Clean up PSH flag handling.

smoltcp-rs · Sep 22, 2017 · af51f2a · af51f2a
1 parent c43446f
commit af51f2a
Showing 2 changed files with 10 additions and 12 deletions.
diff --git a/src/socket/tcp.rs b/src/socket/tcp.rs
@@ -939,16 +939,8 @@ impl<'a> TcpSocket<'a> {
             }
         }
 
-        // We don't care about the PSH flag.
-        let control =
-            if repr.control == TcpControl::Psh {
-                TcpControl::None
-            } else {
-                repr.control
-            };
-
         // Validate and update the state.
-        match (self.state, control) {
+        match (self.state, repr.control.quash_psh()) {
             // RSTs are not accepted in the LISTEN state.
             (State::Listen, TcpControl::Rst) =>
                 return Err(Error::Dropped),
@@ -1493,9 +1485,7 @@ mod test {
             (recv(&mut $socket, $time, |result| {
                 // Most of the time we don't care about the PSH flag.
                 let result = result.map(|mut repr| {
-                    if repr.control == TcpControl::Psh {
-                        repr.control = TcpControl::None;
-                    }
+                    repr.control = repr.control.quash_psh();
                     repr
                 });
                 assert_eq!(result, $result)

diff --git a/src/wire/tcp.rs b/src/wire/tcp.rs
@@ -617,6 +617,14 @@ impl Control {
             _ => 0
         }
     }
+
+    /// Turn the PSH flag into no flag, and keep the rest as-is.
+    pub fn quash_psh(self) -> Control {
+        match self {
+            Control::Psh => Control::None,
+            _ => self
+        }
+    }
 }
 
 /// A high-level representation of a Transmission Control Protocol packet.