diff --git a/actix-proxy-protocol/examples/proxy-server.rs b/actix-proxy-protocol/examples/proxy-server.rs index 05c80b68..6bbb9243 100644 --- a/actix-proxy-protocol/examples/proxy-server.rs +++ b/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?; diff --git a/actix-proxy-protocol/src/tlv.rs b/actix-proxy-protocol/src/tlv.rs index bf7a04ae..e2e001ab 100644 --- a/actix-proxy-protocol/src/tlv.rs +++ b/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; + fn try_from_value(value: &[u8]) -> Option; - fn value_bytes(&self) -> Vec; + fn value_bytes(&self) -> Cow<'_, [u8]>; + + fn try_from_parts(typ: u8, value: &[u8]) -> Option { + if typ != Self::TYPE { + return None; + } + + Self::try_from_value(value) + } } #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -16,11 +24,7 @@ pub struct Crc32c { impl Tlv for Crc32c { const TYPE: u8 = 0x03; - fn try_from_parts(typ: u8, value: &[u8]) -> Option { - if typ != Self::TYPE { - return None; - } - + fn try_from_value(value: &[u8]) -> Option { let checksum_bytes = <[u8; 4]>::try_from(value).ok()?; Some(Self { @@ -28,8 +32,74 @@ impl Tlv for Crc32c { }) } - fn value_bytes(&self) -> Vec { - 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, +} + +impl Noop { + /// + /// + /// # Panics + /// Panics if value is empty (i.e., has length of 0). + pub fn new(value: impl Into>) -> 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 { + 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, +} + +impl UniqueId { + /// + /// + /// # Panics + /// Panics if value is empty (i.e., has length of 0). + pub fn new(value: impl Into>) -> 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 { + Some(Self { + value: value.to_owned(), + }) + } + + fn value_bytes(&self) -> Cow<'_, [u8]> { + Cow::Borrowed(&self.value) } } diff --git a/actix-proxy-protocol/src/v2.rs b/actix-proxy-protocol/src/v2.rs index efabdd12..14335ade 100644 --- a/actix-proxy-protocol/src/v2.rs +++ b/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(&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();