add Tlv trait

Author: robjtede

actix-proxy-protocol/Cargo.toml

@@ -21,6 +21,7 @@ crc32fast = "1"
 futures-core = { version = "0.3.17", default-features = false, features = ["std"] }
 futures-util = { version = "0.3.17", default-features = false, features = ["std"] }
 itoa = "1"
+smallvec = "1"
 tokio = { version = "1.13.1", features = ["sync", "io-util"] }
 tracing = { version = "0.1.30", default-features = false, features = ["log"] }
 

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

@@ -11,7 +11,7 @@ use std::{
     },
 };
 
-use actix_proxy_protocol::{v1, v2, AddressFamily, Command, Tlv, TransportProtocol};
+use actix_proxy_protocol::{tlv, v1, v2, AddressFamily, Command, TransportProtocol};
 use actix_rt::net::TcpStream;
 use actix_server::Server;
 use actix_service::{fn_service, ServiceFactoryExt as _};
@@ -31,6 +31,7 @@ TLV = type-length-value
 TO DO:
 handle UNKNOWN transport
 v2 UNSPEC mode
+AF_UNIX socket
 */
 
 fn extend_with_ip_bytes(buf: &mut Vec<u8>, ip: IpAddr) {
@@ -71,18 +72,10 @@ async fn wrap_with_proxy_protocol_v2(mut stream: TcpStream) -> io::Result<()> {
         UPSTREAM.to_string()
     );
 
-    let mut proxy_header = v2::Header::new(
-        Command::Proxy,
-        TransportProtocol::Stream,
-        AddressFamily::Inet,
-        SocketAddr::from(([127, 0, 0, 1], 8082)),
-        *UPSTREAM,
-        vec![
-            Tlv::new(0x05, [0x34, 0x32, 0x36, 0x39]), // UNIQUE_ID
-            Tlv::new(0x04, "NOOP m9"),                // NOOP
-        ],
-    );
+    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_crc23c_checksum_tlv();
 
     proxy_header.write_to_tokio(&mut upstream).await?;

actix-proxy-protocol/src/lib.rs

@@ -16,6 +16,7 @@ use std::{
 use arrayvec::{ArrayString, ArrayVec};
 use tokio::io::{AsyncWrite, AsyncWriteExt as _};
 
+pub mod tlv;
 pub mod v1;
 pub mod v2;
 
@@ -165,80 +166,3 @@ enum ProxyProtocolHeader {
     V1(v1::Header),
     V2(v2::Header),
 }
-
-#[derive(Debug, Clone, Default, PartialEq, Eq)]
-pub struct Pp2Crc32c {
-    checksum: u32,
-}
-
-impl Pp2Crc32c {
-    fn to_tlv(&self) -> Tlv {
-        Tlv {
-            r#type: 0x03,
-            value: self.checksum.to_be_bytes().to_vec(),
-        }
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct Tlv {
-    r#type: u8,
-    value: Vec<u8>,
-}
-
-impl Tlv {
-    pub fn new(r#type: u8, value: impl Into<Vec<u8>>) -> Self {
-        let value = value.into();
-
-        assert!(
-            !value.is_empty(),
-            "TLV values must have length greater than 1",
-        );
-
-        Self { r#type, value }
-    }
-
-    fn len(&self) -> u16 {
-        // 1b type + 2b len
-        // + [len]b value
-        1 + 2 + (self.value.len() as u16)
-    }
-
-    pub fn write_to(&self, wrt: &mut impl io::Write) -> io::Result<()> {
-        wrt.write_all(&[self.r#type])?;
-        wrt.write_all(&(self.value.len() as u16).to_be_bytes())?;
-        wrt.write_all(&self.value)?;
-        Ok(())
-    }
-
-    pub fn as_crc32c(&self) -> Option<Pp2Crc32c> {
-        if self.r#type != 0x03 {
-            return None;
-        }
-
-        let checksum_bytes = <[u8; 4]>::try_from(self.value.as_slice()).ok()?;
-
-        Some(Pp2Crc32c {
-            checksum: u32::from_be_bytes(checksum_bytes),
-        })
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn tlv_as_crc32c() {
-        let noop = Tlv::new(0x04, vec![0x00]);
-        assert_eq!(noop.as_crc32c(), None);
-
-        let noop = Tlv::new(0x03, vec![0x08, 0x70, 0x17, 0x7b]);
-        assert_eq!(
-            noop.as_crc32c(),
-            Some(Pp2Crc32c {
-                checksum: 141563771
-            })
-        );
-    }
-}

actix-proxy-protocol/src/tlv.rs

@@ -0,0 +1,58 @@
+use std::convert::TryFrom;
+
+pub trait Tlv: Sized {
+    const TYPE: u8;
+
+    fn try_from_parts(typ: u8, value: &[u8]) -> Option<Self>;
+
+    fn value_bytes(&self) -> Vec<u8>;
+}
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct Crc32c {
+    pub(crate) checksum: u32,
+}
+
+impl Tlv for Crc32c {
+    const TYPE: u8 = 0x03;
+
+    fn try_from_parts(typ: u8, value: &[u8]) -> Option<Self> {
+        if typ != Self::TYPE {
+            return None;
+        }
+
+        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()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // #[test]
+    // #[should_panic]
+    // fn tlv_zero_len() {
+    //     Tlv::new(0x00, vec![]);
+    // }
+
+    #[test]
+    fn tlv_as_crc32c() {
+        // noop
+        assert_eq!(Crc32c::try_from_parts(0x04, &[0x00]), None);
+
+        assert_eq!(
+            Crc32c::try_from_parts(0x03, &[0x08, 0x70, 0x17, 0x7b]),
+            Some(Crc32c {
+                checksum: 141563771
+            })
+        );
+    }
+}