1 files changed, 36 insertions, 12 deletions

diff --git a/third_party/rust/neqo-common/src/tos.rs b/third_party/rust/neqo-common/src/tos.rs
index 3610f72750..533c5447e2 100644
--- a/third_party/rust/neqo-common/src/tos.rs
+++ b/third_party/rust/neqo-common/src/tos.rs
@@ -36,7 +36,7 @@ impl From<IpTosEcn> for u8 {
 
 impl From<u8> for IpTosEcn {
     fn from(v: u8) -> Self {
-        match v & 0b11 {
+        match v & 0b0000_0011 {
             0b00 => IpTosEcn::NotEct,
             0b01 => IpTosEcn::Ect1,
             0b10 => IpTosEcn::Ect0,
@@ -47,8 +47,8 @@ impl From<u8> for IpTosEcn {
 }
 
 impl From<IpTos> for IpTosEcn {
-    fn from(value: IpTos) -> Self {
-        IpTosEcn::from(value.0 & 0x3)
+    fn from(v: IpTos) -> Self {
+        IpTosEcn::from(u8::from(v))
     }
 }
 
@@ -166,14 +166,13 @@ impl From<u8> for IpTosDscp {
 }
 
 impl From<IpTos> for IpTosDscp {
-    fn from(value: IpTos) -> Self {
-        IpTosDscp::from(value.0 & 0xfc)
+    fn from(v: IpTos) -> Self {
+        IpTosDscp::from(u8::from(v))
     }
 }
 
 /// The type-of-service field in an IP packet.
-#[allow(clippy::module_name_repetitions)]
-#[derive(Copy, Clone, PartialEq, Eq)]
+#[derive(Copy, Clone, PartialEq, Eq, Default)]
 pub struct IpTos(u8);
 
 impl From<IpTosEcn> for IpTos {
@@ -215,15 +214,19 @@ impl From<u8> for IpTos {
 impl Debug for IpTos {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_tuple("IpTos")
-            .field(&IpTosDscp::from(self.0 & 0xfc))
-            .field(&IpTosEcn::from(self.0 & 0x3))
+            .field(&IpTosDscp::from(*self))
+            .field(&IpTosEcn::from(*self))
             .finish()
     }
 }
 
-impl Default for IpTos {
-    fn default() -> Self {
-        (IpTosDscp::default(), IpTosEcn::default()).into()
+impl IpTos {
+    pub fn set_ecn(&mut self, ecn: IpTosEcn) {
+        self.0 = u8::from(IpTosDscp::from(*self)) | u8::from(ecn);
+    }
+
+    pub fn set_dscp(&mut self, dscp: IpTosDscp) {
+        self.0 = u8::from(IpTosEcn::from(*self)) | u8::from(dscp);
     }
 }
 
@@ -322,4 +325,25 @@ mod tests {
         assert_eq!(tos, u8::from(iptos));
         assert_eq!(IpTos::from(tos), iptos);
     }
+
+    #[test]
+    fn iptos_to_iptosdscp() {
+        let tos = IpTos::from((IpTosDscp::Af41, IpTosEcn::NotEct));
+        let dscp = IpTosDscp::from(tos);
+        assert_eq!(dscp, IpTosDscp::Af41);
+    }
+
+    #[test]
+    fn tos_modify_ecn() {
+        let mut iptos: IpTos = (IpTosDscp::Af41, IpTosEcn::NotEct).into();
+        iptos.set_ecn(IpTosEcn::Ce);
+        assert_eq!(u8::from(iptos), 0b1000_1011);
+    }
+
+    #[test]
+    fn tos_modify_dscp() {
+        let mut iptos: IpTos = (IpTosDscp::Af41, IpTosEcn::Ect1).into();
+        iptos.set_dscp(IpTosDscp::Le);
+        assert_eq!(u8::from(iptos), 0b0000_0101);
+    }
 }