net: Fix token magic not being added to ConnectAccept

heinrich5991 · Jul 12, 2023 · e8d432e · e8d432e
1 parent c61dda2
commit e8d432e
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 14 deletions.
diff --git a/net/src/connection.rs b/net/src/connection.rs
@@ -903,7 +903,7 @@ mod test {
         let packet = cb.0.pop_front().unwrap();
         assert!(cb.0.is_empty());
         hexdump(&packet);
-        assert!(&packet == b"\x10\x00\x00\x02\x12\x34\x56\x78");
+        assert!(&packet == b"\x10\x00\x00\x02TKEN\x12\x34\x56\x78");
 
         // Accept
         assert!(client.feed(cb, &mut Panic, &packet, &mut buffer[..]).0.collect_vec()

diff --git a/net/src/protocol.rs b/net/src/protocol.rs
@@ -45,7 +45,7 @@ pub const TOKEN_NONE:     Token = Token([0xff, 0xff, 0xff, 0xff]);
 pub const TOKEN_RESERVED: Token = Token([0x00, 0x00, 0x00, 0x00]);
 
 pub const CTRLMSG_CLOSE_REASON_LENGTH: usize = 127;
-pub const CTRLMSG_CONNECT_TOKEN_MAGIC: &[u8; 4] = b"TKEN";
+pub const CTRLMSG_TOKEN_MAGIC: &[u8; 4] = b"TKEN";
 pub const CHUNK_FLAGS_BITS: u32 = 2;
 pub const CHUNK_SIZE_BITS: u32 = 10;
 pub const PACKET_FLAGS_BITS: u32 = 4;
@@ -308,8 +308,8 @@ fn has_token_heuristic(control: bool, num_chunks: u8, payload: &[u8]) -> bool {
     let payload_end_heuristic = if control {
         let (&control, payload) = unwrap_or_return!(payload.split_first(), false);
         match control {
-            CTRLMSG_CONNECT => {
-                if payload.len() < 4 || &payload[0..4] != CTRLMSG_CONNECT_TOKEN_MAGIC {
+            CTRLMSG_CONNECT | CTRLMSG_CONNECTACCEPT => {
+                if payload.len() < 4 || &payload[0..4] != CTRLMSG_TOKEN_MAGIC {
                     return false;
                 }
                 1 + 4
@@ -475,18 +475,16 @@ impl<'a> Packet<'a> {
 
             let (&control, payload) = unwrap_or_return!(payload.split_first(),
                                                         Err(ControlMissing));
-            if control != CTRLMSG_CLOSE && control != CTRLMSG_CONNECT && payload.len() != 0 {
-                warn.warn(Warning::ControlExcessData);
-            }
-            if control == CTRLMSG_CONNECT {
-                if token.is_some() {
-                    if !payload.starts_with(CTRLMSG_CONNECT_TOKEN_MAGIC) {
+            // check for excess data
+            match control {
+                CTRLMSG_CONNECT | CTRLMSG_CONNECTACCEPT => if token.is_some() {
+                    if !payload.starts_with(CTRLMSG_TOKEN_MAGIC) {
                         warn.warn(Warning::ControlConnectMissingTokenMagic);
                         if !payload.is_empty() {
                             warn.warn(Warning::ControlExcessData);
                         }
                     } else {
-                        if payload.len() > CTRLMSG_CONNECT_TOKEN_MAGIC.len() {
+                        if payload.len() > CTRLMSG_TOKEN_MAGIC.len() {
                             warn.warn(Warning::ControlExcessData);
                         }
                     }
@@ -495,6 +493,10 @@ impl<'a> Packet<'a> {
                         warn.warn(Warning::ControlExcessData);
                     }
                 }
+                CTRLMSG_CLOSE => {}, // handled later
+                _ => if payload.len() != 0 {
+                    warn.warn(Warning::ControlExcessData);
+                }
             }
             let control = match control {
                 CTRLMSG_KEEPALIVE => ControlPacket::KeepAlive,
@@ -657,8 +659,8 @@ impl<'a> ControlPacket<'a> {
             ControlPacket::Close(..) => CTRLMSG_CLOSE,
         };
         buffer.write(&[magic])?;
-        if matches!(*self, ControlPacket::Connect) && token.is_some() {
-            buffer.write(CTRLMSG_CONNECT_TOKEN_MAGIC)?;
+        if matches!(*self, ControlPacket::Connect | ControlPacket::ConnectAccept) && token.is_some() {
+            buffer.write(CTRLMSG_TOKEN_MAGIC)?;
         }
         match *self {
             ControlPacket::Close(m) => {

diff --git a/wireshark-dissector/src/tw.rs b/wireshark-dissector/src/tw.rs
@@ -309,7 +309,7 @@ unsafe fn dissect_impl(
                 ctrl_raw,
             );
             match ctrl {
-                Connect => {
+                Connect | ConnectAccept => {
                     if token.is_some() {
                         field_none!(tree, HF_PACKET_TOKEN_MAGIC, 1, 4,
                             "Magic bytes for DDNet token protocol: \"TKEN\"",