add some typed TLVs

Author: robjtede

actix-proxy-protocol/examples/proxy-server.rs

@@ -74,8 +74,8 @@ async fn wrap_with_proxy_protocol_v2(mut stream: TcpStream) -> io::Result<()> {
 
     let mut proxy_header = v2::Header::new_tcp_ipv4_proxy(([127, 0, 0, 1], 8082), *UPSTREAM);
 
-    proxy_header.add_tlv(0x05, [0x34, 0x32, 0x36, 0x39]); // UNIQUE_ID
-    proxy_header.add_tlv(0x04, "NOOP m9"); // NOOP
+    proxy_header.add_typed_tlv(tlv::UniqueId::new("4269")); // UNIQUE_ID
+    proxy_header.add_typed_tlv(tlv::Noop::new("NOOP m8")); // NOOP
     proxy_header.add_crc23c_checksum_tlv();
 
     proxy_header.write_to_tokio(&mut upstream).await?;

actix-proxy-protocol/src/tlv.rs

@@ -1,11 +1,19 @@
-use std::convert::TryFrom;
+use std::{borrow::Cow, convert::TryFrom};
 
 pub trait Tlv: Sized {
     const TYPE: u8;
 
-    fn try_from_parts(typ: u8, value: &[u8]) -> Option<Self>;
+    fn try_from_value(value: &[u8]) -> Option<Self>;
 
-    fn value_bytes(&self) -> Vec<u8>;
+    fn value_bytes(&self) -> Cow<'_, [u8]>;
+
+    fn try_from_parts(typ: u8, value: &[u8]) -> Option<Self> {
+        if typ != Self::TYPE {
+            return None;
+        }
+
+        Self::try_from_value(value)
+    }
 }
 
 #[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -16,20 +24,82 @@ pub struct Crc32c {
 impl Tlv for Crc32c {
     const TYPE: u8 = 0x03;
 
-    fn try_from_parts(typ: u8, value: &[u8]) -> Option<Self> {
-        if typ != Self::TYPE {
-            return None;
-        }
-
+    fn try_from_value(value: &[u8]) -> Option<Self> {
         let checksum_bytes = <[u8; 4]>::try_from(value).ok()?;
 
         Some(Self {
             checksum: u32::from_be_bytes(checksum_bytes),
         })
     }
 
-    fn value_bytes(&self) -> Vec<u8> {
-        self.checksum.to_be_bytes().to_vec()
+    fn value_bytes(&self) -> Cow<'_, [u8]> {
+        Cow::Owned(self.checksum.to_be_bytes().to_vec())
+    }
+}
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct Noop {
+    value: Vec<u8>,
+}
+
+impl Noop {
+    ///
+    ///
+    /// # Panics
+    /// Panics if value is empty (i.e., has length of 0).
+    pub fn new(value: impl Into<Vec<u8>>) -> Self {
+        let value = value.into();
+
+        assert!(!value.is_empty(), "Noop TLV `value` cannot be empty");
+
+        Self { value }
+    }
+}
+
+impl Tlv for Noop {
+    const TYPE: u8 = 0x04;
+
+    fn try_from_value(value: &[u8]) -> Option<Self> {
+        Some(Self {
+            value: value.to_owned(),
+        })
+    }
+
+    fn value_bytes(&self) -> Cow<'_, [u8]> {
+        Cow::Borrowed(&self.value)
+    }
+}
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct UniqueId {
+    value: Vec<u8>,
+}
+
+impl UniqueId {
+    ///
+    ///
+    /// # Panics
+    /// Panics if value is empty (i.e., has length of 0).
+    pub fn new(value: impl Into<Vec<u8>>) -> Self {
+        let value = value.into();
+
+        assert!(!value.is_empty(), "UniqueId TLV `value` cannot be empty");
+
+        Self { value }
+    }
+}
+
+impl Tlv for UniqueId {
+    const TYPE: u8 = 0x05;
+
+    fn try_from_value(value: &[u8]) -> Option<Self> {
+        Some(Self {
+            value: value.to_owned(),
+        })
+    }
+
+    fn value_bytes(&self) -> Cow<'_, [u8]> {
+        Cow::Borrowed(&self.value)
     }
 }
 

actix-proxy-protocol/src/v2.rs

@@ -57,6 +57,10 @@ impl Header {
         self.tlvs.push((typ, SmallVec::from_slice(value.as_ref())));
     }
 
+    pub fn add_typed_tlv<T: Tlv>(&mut self, tlv: T) {
+        self.add_tlv(T::TYPE, tlv.value_bytes());
+    }
+
     fn v2_len(&self) -> u16 {
         let addr_len = if self.src.is_ipv4() {
             4 + 2 // 4b IPv4 + 2b port number
@@ -144,8 +148,7 @@ impl Header {
         //   the bits unchanged.
 
         // add zeroed checksum field to TLVs
-        let crc = (Crc32c::TYPE, Crc32c::default().value_bytes().to_smallvec());
-        self.tlvs.push(crc);
+        self.add_typed_tlv(Crc32c::default());
 
         // write PROXY header to buffer
         let mut buf = Vec::new();