diff --git a/src/decode/domain_name.rs b/src/decode/domain_name.rs index e73ef59..a57fe22 100644 --- a/src/decode/domain_name.rs +++ b/src/decode/domain_name.rs @@ -82,6 +82,82 @@ impl<'a, 'b: 'a> Decoder<'a, 'b> { } } } + + pub(super) fn domain_name_without_compression(&mut self) -> DecodeResult { + let mut domain_name = DomainName::default(); + + loop { + match self.domain_name_length()? { + DomainNameLength::Compressed(offset) => { + return Err(DecodeError::DomainNameCompressed(offset)) + } + DomainNameLength::Label(0) => return Ok(domain_name), + DomainNameLength::Label(length) => { + self.domain_name_label(&mut domain_name, length)? + } + } + } + } } impl_decode!(DomainName, domain_name); +impl_decode!( + DomainName, + domain_name_without_compression, + decode_without_compression +); + +#[test] +fn domain_name_without_compression_1() { + let domain_name: DomainName = "example.org.".parse().unwrap(); + let mut decoder = Decoder::main(bytes::Bytes::copy_from_slice( + &b"\x07example\x03org\0\x07example\x03org\0"[..], + )); + assert_eq!( + decoder.domain_name_without_compression().unwrap(), + domain_name + ); + assert_eq!( + decoder.domain_name_without_compression().unwrap(), + domain_name + ); + assert!(decoder.is_finished().unwrap()); +} + +#[test] +fn domain_name_without_compression_2() { + let domain_name: DomainName = "example.org.".parse().unwrap(); + let mut decoder = Decoder::main(bytes::Bytes::copy_from_slice( + &b"\x07example\x03org\0\x07example\x03org\0"[..], + )); + assert_eq!( + decoder.domain_name_without_compression().unwrap(), + domain_name + ); + assert_eq!( + decoder.domain_name_without_compression().unwrap(), + domain_name + ); + assert!(decoder.is_finished().unwrap()); +} + +#[test] +fn domain_name_without_compression_3() { + let domain_name: DomainName = "example.org.".parse().unwrap(); + let mut decoder = Decoder::main(bytes::Bytes::copy_from_slice( + &b"\x07example\x03org\0\x07example\x03org\0\xc0\r"[..], + )); + assert_eq!( + decoder.domain_name_without_compression().unwrap(), + domain_name + ); + assert_eq!( + decoder.domain_name_without_compression().unwrap(), + domain_name + ); + assert_eq!( + decoder.domain_name_without_compression(), + Err(DecodeError::DomainNameCompressed(13)) + ); + assert!(decoder.is_finished().unwrap()); +} diff --git a/src/decode/error.rs b/src/decode/error.rs index dcfe06b..e27403f 100644 --- a/src/decode/error.rs +++ b/src/decode/error.rs @@ -105,4 +105,6 @@ pub enum DecodeError { TagError(#[from] TagError), #[error("ECH length mismatch. Expected {0} got {1}")] ECHLengthMismatch(usize, usize), + #[error("The domain name is compressed: {0}")] + DomainNameCompressed(u16), } diff --git a/src/decode/macros.rs b/src/decode/macros.rs index a3f472e..645890c 100644 --- a/src/decode/macros.rs +++ b/src/decode/macros.rs @@ -1,7 +1,10 @@ macro_rules! impl_decode { ($i:path, $m:ident) => { + impl_decode!($i, $m, decode); + }; + ($i:path, $m:ident, $n:ident) => { impl $i { - pub fn decode(bytes: bytes::Bytes) -> crate::DecodeResult<$i> { + pub fn $n(bytes: bytes::Bytes) -> crate::DecodeResult<$i> { let mut decoder = crate::decode::Decoder::main(bytes); decoder.$m() } diff --git a/src/encode/domain_name.rs b/src/encode/domain_name.rs index 214f413..d4ff39c 100644 --- a/src/encode/domain_name.rs +++ b/src/encode/domain_name.rs @@ -71,6 +71,17 @@ impl Encoder { self.merge_domain_name_index(domain_name_index, 0)?; Ok(()) } + + pub(super) fn domain_name_without_compression( + &mut self, + domain_name: &DomainName, + ) -> EncodeResult<()> { + for (label, _) in domain_name.iter() { + self.label(&label)?; + } + self.string_with_len("")?; + Ok(()) + } } impl DomainName { @@ -98,3 +109,37 @@ impl<'a> Iterator for DomainNameIter<'a> { } } impl_encode!(DomainName, domain_name); +impl_encode!( + DomainName, + domain_name_without_compression, + encode_without_compression +); + +#[test] +fn domain_name_without_compression_1() { + let domain_name: DomainName = "example.org.".parse().unwrap(); + let mut encoder = Encoder::default(); + encoder + .domain_name_without_compression(&domain_name) + .unwrap(); + encoder + .domain_name_without_compression(&domain_name) + .unwrap(); + assert_eq!( + encoder.bytes, + &b"\x07example\x03org\0\x07example\x03org\0"[..] + ); +} + +#[test] +fn domain_name_without_compression_2() { + let domain_name = DomainName::default(); + let mut encoder = Encoder::default(); + encoder + .domain_name_without_compression(&domain_name) + .unwrap(); + encoder + .domain_name_without_compression(&domain_name) + .unwrap(); + assert_eq!(encoder.bytes, &b"\0\0"[..]); +} diff --git a/src/encode/macros.rs b/src/encode/macros.rs index 78d414d..3c29fb0 100644 --- a/src/encode/macros.rs +++ b/src/encode/macros.rs @@ -12,8 +12,11 @@ macro_rules! impl_encode_without_result { macro_rules! impl_encode { ($i:path, $m:ident) => { + impl_encode!($i, $m, encode); + }; + ($i:path, $m:ident, $n:ident) => { impl $i { - pub fn encode(&self) -> crate::EncodeResult { + pub fn $n(&self) -> crate::EncodeResult { let mut encoder = crate::encode::Encoder::default(); encoder.$m(self)?; Ok(encoder.bytes)