From c292a3c6a60b9b67c82dfaf5b2dab212d8f4a530 Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Sat, 31 Aug 2024 10:11:16 -0700
Subject: [PATCH] Fix and test validation of IDCID length

---
 quinn-proto/src/endpoint.rs  |  4 ++--
 quinn-proto/src/tests/mod.rs | 23 ++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
index 1e0d77f95..36546448e 100644
--- a/quinn-proto/src/endpoint.rs
+++ b/quinn-proto/src/endpoint.rs
@@ -706,8 +706,8 @@ impl Endpoint {
         // length. If we ever issue non-Retry address validation tokens via `NEW_TOKEN`, then we'll
         // also need to validate CID length for those after decoding the token.
         if header.dst_cid.len() < 8
-            && (!header.token_pos.is_empty()
-                && header.dst_cid.len() != self.local_cid_generator.cid_len())
+            && (header.token_pos.is_empty()
+                || header.dst_cid.len() != self.local_cid_generator.cid_len())
         {
             debug!(
                 "rejecting connection due to invalid DCID length {}",
diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs
index f67ab75ab..2c940caa6 100644
--- a/quinn-proto/src/tests/mod.rs
+++ b/quinn-proto/src/tests/mod.rs
@@ -7,7 +7,7 @@ use std::{
 };
 
 use assert_matches::assert_matches;
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use hex_literal::hex;
 use rand::RngCore;
 use ring::hmac;
@@ -3143,3 +3143,24 @@ fn voluntary_ack_with_large_datagrams() {
         "client should have sent some ACK-only packets"
     );
 }
+
+#[test]
+fn reject_short_idcid() {
+    let _guard = subscribe();
+    let client_addr = "[::2]:7890".parse().unwrap();
+    let mut server = Endpoint::new(
+        Default::default(),
+        Some(Arc::new(server_config())),
+        true,
+        None,
+    );
+    let now = Instant::now();
+    let mut buf = Vec::with_capacity(server.config().get_max_udp_payload_size() as usize);
+    // Initial header that has an empty DCID but is otherwise well-formed
+    let mut initial = BytesMut::from(hex!("c4 00000001 00 00 00 3f").as_ref());
+    initial.resize(MIN_INITIAL_SIZE.into(), 0);
+    let event = server.handle(now, client_addr, None, None, initial, &mut buf);
+    let Some(DatagramEvent::Response(Transmit { .. })) = event else {
+        panic!("expected an initial close");
+    };
+}