attributes: add some tests for some write_into() implementations

Author: ystreet

stun-types/src/attribute/alternate.rs

@@ -217,12 +217,18 @@ mod tests {
                 })
             ));
             // provide incorrectly typed data
-            let mut data: Vec<_> = raw.into();
+            let mut data: Vec<_> = raw.clone().into();
             BigEndian::write_u16(&mut data[0..2], 0);
             assert!(matches!(
                 AlternateServer::try_from(&RawAttribute::from_bytes(data.as_ref()).unwrap()),
                 Err(StunParseError::WrongAttributeImplementation)
             ));
+
+            let mut dest = vec![0; raw.padded_len()];
+            mapped.write_into(&mut dest).unwrap();
+            let raw = RawAttribute::from_bytes(&dest).unwrap();
+            let mapped2 = AlternateServer::try_from(&raw).unwrap();
+            assert_eq!(mapped2.server(), *addr);
         }
     }
 
@@ -246,12 +252,18 @@ mod tests {
             AlternateDomain::try_from(&RawAttribute::from_bytes(data.as_ref()).unwrap()),
             Err(StunParseError::WrongAttributeImplementation)
         ));
-        let mut data: Vec<_> = raw.into();
+        let mut data: Vec<_> = raw.clone().into();
         // invalid utf-8 data
         data[8] = 0x88;
         assert!(matches!(
             AlternateDomain::try_from(&RawAttribute::from_bytes(data.as_ref()).unwrap()),
             Err(StunParseError::InvalidAttributeData)
         ));
+
+        let mut dest = vec![0; raw.padded_len()];
+        attr.write_into(&mut dest).unwrap();
+        let raw = RawAttribute::from_bytes(&dest).unwrap();
+        let mapped2 = AlternateDomain::try_from(&raw).unwrap();
+        assert_eq!(mapped2.domain(), dns);
     }
 }

stun-types/src/attribute/error.rs

@@ -49,14 +49,8 @@ impl AttributeWrite for ErrorCode {
 
     fn write_into_unchecked(&self, dest: &mut [u8]) {
         let len = self.padded_len();
-        let offset = self.write_header_unchecked(dest);
-        dest[offset] = 0u8;
-        dest[offset + 1] = 0u8;
-        dest[offset + 2] = (self.code / 100) as u8;
-        dest[offset + 3] = (self.code % 100) as u8;
-        let offset = offset + 4;
-        dest[offset..offset + self.reason.len()].copy_from_slice(self.reason.as_bytes());
-        let offset = offset + self.reason.len();
+        let mut offset = self.write_header_unchecked(dest);
+        offset += self.write_into_data(&mut dest[offset..]);
         if len - offset > 0 {
             dest[offset..len].fill(0);
         }
@@ -247,6 +241,16 @@ impl ErrorCode {
             _ => "Unknown",
         }
     }
+
+    fn write_into_data(&self, dest: &mut [u8]) -> usize {
+        dest[0] = 0u8;
+        dest[1] = 0u8;
+        dest[2] = (self.code / 100) as u8;
+        dest[3] = (self.code % 100) as u8;
+        let bytes = self.reason.as_bytes();
+        dest[4..4 + bytes.len()].copy_from_slice(bytes);
+        4 + bytes.len()
+    }
 }
 
 impl std::fmt::Display for ErrorCode {
@@ -273,23 +277,15 @@ impl Attribute for UnknownAttributes {
 }
 impl AttributeWrite for UnknownAttributes {
     fn to_raw(&self) -> RawAttribute {
-        let mut data = Vec::with_capacity(self.length() as usize);
-        for attr in &self.attributes {
-            let mut encoded = [0; 2];
-            BigEndian::write_u16(&mut encoded, (*attr).into());
-            data.extend(encoded);
-        }
+        let mut data = vec![0; self.length() as usize];
+        self.write_into_data(&mut data);
         RawAttribute::new_owned(UnknownAttributes::TYPE, data.into_boxed_slice())
     }
 
     fn write_into_unchecked(&self, dest: &mut [u8]) {
         let len = self.padded_len();
         let mut offset = self.write_header_unchecked(dest);
-        for attr in &self.attributes {
-            let mut encoded = [0; 2];
-            BigEndian::write_u16(&mut encoded, (*attr).into());
-            offset += 2;
-        }
+        offset += self.write_into_data(&mut dest[offset..]);
         if len - offset > 0 {
             dest[offset..len].fill(0);
         }
@@ -360,6 +356,15 @@ impl UnknownAttributes {
     pub fn has_attribute(&self, attr: AttributeType) -> bool {
         self.attributes.iter().any(|&a| a == attr)
     }
+
+    fn write_into_data(&self, dest: &mut [u8]) -> usize {
+        let mut offset = 0;
+        for attr in &self.attributes {
+            BigEndian::write_u16(&mut dest[offset..offset + 2], (*attr).into());
+            offset += 2;
+        }
+        offset
+    }
 }
 
 impl std::fmt::Display for UnknownAttributes {
@@ -392,6 +397,13 @@ mod tests {
             let err2 = ErrorCode::try_from(&raw).unwrap();
             assert_eq!(err2.code(), code);
             assert_eq!(err2.reason(), reason);
+
+            let mut dest = vec![0; raw.padded_len()];
+            err.write_into(&mut dest).unwrap();
+            let raw = RawAttribute::from_bytes(&dest).unwrap();
+            let err2 = ErrorCode::try_from(&raw).unwrap();
+            assert_eq!(err2.code(), code);
+            assert_eq!(err2.reason(), reason);
         }
     }
 
@@ -527,11 +539,19 @@ mod tests {
             })
         ));
         // provide incorrectly typed data
-        let mut data: Vec<_> = raw.into();
+        let mut data: Vec<_> = raw.clone().into();
         BigEndian::write_u16(&mut data[0..2], 0);
         assert!(matches!(
             UnknownAttributes::try_from(&RawAttribute::from_bytes(data.as_ref()).unwrap()),
             Err(StunParseError::WrongAttributeImplementation)
         ));
+
+        let mut dest = vec![0; raw.padded_len()];
+        unknown.write_into(&mut dest).unwrap();
+        tracing::error!("{dest:?}");
+        let raw = RawAttribute::from_bytes(&dest).unwrap();
+        let unknown2 = UnknownAttributes::try_from(&raw).unwrap();
+        assert!(unknown2.has_attribute(Realm::TYPE));
+        assert!(unknown2.has_attribute(AlternateServer::TYPE));
     }
 }

stun-types/src/attribute/fingerprint.rs

@@ -116,6 +116,8 @@ impl std::fmt::Display for Fingerprint {
 
 #[cfg(test)]
 mod tests {
+    use crate::prelude::AttributeExt;
+
     use super::*;
     use byteorder::{BigEndian, ByteOrder};
     use tracing::trace;
@@ -145,11 +147,17 @@ mod tests {
             })
         ));
         // provide incorrectly typed data
-        let mut data: Vec<_> = raw.into();
+        let mut data: Vec<_> = raw.clone().into();
         BigEndian::write_u16(&mut data[0..2], 0);
         assert!(matches!(
             Fingerprint::try_from(&RawAttribute::from_bytes(data.as_ref()).unwrap()),
             Err(StunParseError::WrongAttributeImplementation)
         ));
+
+        let mut dest = vec![0; raw.padded_len()];
+        attr.write_into(&mut dest).unwrap();
+        let raw = RawAttribute::from_bytes(&dest).unwrap();
+        let attr2 = Fingerprint::try_from(&raw).unwrap();
+        assert_eq!(attr2.fingerprint(), &val);
     }
 }