From 36950c6bb193ea4294a568b2fed813cc6a4f4ee7 Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Mon, 25 Mar 2024 12:08:43 -0700
Subject: [PATCH] Factor initial packet plain header fields out into a struct

---
 quinn-proto/src/endpoint.rs | 29 +++++++++++++------------
 quinn-proto/src/packet.rs   | 43 +++++++++++++++++++++----------------
 2 files changed, 40 insertions(+), 32 deletions(-)

diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
index 614e96011..bd09b9e7f 100644
--- a/quinn-proto/src/endpoint.rs
+++ b/quinn-proto/src/endpoint.rs
@@ -210,27 +210,28 @@ impl Endpoint {
             }
         };
 
-        if let Some(version) = first_decode.initial_version() {
+        if let Some(header) = first_decode.initial_header() {
             if datagram_len < MIN_INITIAL_SIZE as usize {
                 debug!("ignoring short initial for connection {}", dst_cid);
                 return None;
             }
 
-            let crypto = match server_config
-                .crypto
-                .initial_keys(version, dst_cid, Side::Server)
-            {
-                Ok(keys) => keys,
-                Err(UnsupportedVersion) => {
-                    // This probably indicates that the user set supported_versions incorrectly in
-                    // `EndpointConfig`.
-                    debug!(
+            let crypto =
+                match server_config
+                    .crypto
+                    .initial_keys(header.version, dst_cid, Side::Server)
+                {
+                    Ok(keys) => keys,
+                    Err(UnsupportedVersion) => {
+                        // This probably indicates that the user set supported_versions incorrectly in
+                        // `EndpointConfig`.
+                        debug!(
                         "ignoring initial packet version {:#x} unsupported by cryptographic layer",
-                        version
+                        header.version
                     );
-                    return None;
-                }
-            };
+                        return None;
+                    }
+                };
             return match first_decode.finish(Some(&*crypto.header.remote)) {
                 Ok(packet) => {
                     self.handle_first_packet(now, addresses, ecn, packet, remaining, &crypto, buf)
diff --git a/quinn-proto/src/packet.rs b/quinn-proto/src/packet.rs
index b70eb98f7..d8f2611f7 100644
--- a/quinn-proto/src/packet.rs
+++ b/quinn-proto/src/packet.rs
@@ -60,11 +60,8 @@ impl PartialDecode {
         self.buf.get_ref()
     }
 
-    pub(crate) fn initial_version(&self) -> Option<u32> {
-        match self.plain_header {
-            PlainHeader::Initial { version, .. } => Some(version),
-            _ => None,
-        }
+    pub(crate) fn initial_header(&self) -> Option<&PlainInitialHeader> {
+        self.plain_header.as_initial()
     }
 
     pub(crate) fn has_long_header(&self) -> bool {
@@ -119,13 +116,13 @@ impl PartialDecode {
             mut buf,
         } = self;
 
-        if let Initial {
+        if let Initial(PlainInitialHeader {
             dst_cid,
             src_cid,
             token_pos,
             version,
             ..
-        } = plain_header
+        }) = plain_header
         {
             let number = Self::decrypt_header(&mut buf, header_crypto.unwrap())?;
             let header_len = buf.position() as usize;
@@ -481,13 +478,7 @@ impl PartialEncode {
 
 #[derive(Clone, Debug)]
 pub(crate) enum PlainHeader {
-    Initial {
-        dst_cid: ConnectionId,
-        src_cid: ConnectionId,
-        token_pos: Range<usize>,
-        len: u64,
-        version: u32,
-    },
+    Initial(PlainInitialHeader),
     Long {
         ty: LongType,
         dst_cid: ConnectionId,
@@ -512,10 +503,17 @@ pub(crate) enum PlainHeader {
 }
 
 impl PlainHeader {
+    pub(crate) fn as_initial(&self) -> Option<&PlainInitialHeader> {
+        match self {
+            Self::Initial(x) => Some(x),
+            _ => None,
+        }
+    }
+
     fn dst_cid(&self) -> &ConnectionId {
         use self::PlainHeader::*;
         match self {
-            Initial { dst_cid, .. } => dst_cid,
+            Initial(header) => &header.dst_cid,
             Long { dst_cid, .. } => dst_cid,
             Retry { dst_cid, .. } => dst_cid,
             Short { dst_cid, .. } => dst_cid,
@@ -526,7 +524,7 @@ impl PlainHeader {
     fn payload_len(&self) -> Option<u64> {
         use self::PlainHeader::*;
         match self {
-            Initial { len, .. } | Long { len, .. } => Some(*len),
+            Initial(PlainInitialHeader { len, .. }) | Long { len, .. } => Some(*len),
             _ => None,
         }
     }
@@ -587,13 +585,13 @@ impl PlainHeader {
                     buf.advance(token_len);
 
                     let len = buf.get_var()?;
-                    Ok(Self::Initial {
+                    Ok(Self::Initial(PlainInitialHeader {
                         dst_cid,
                         src_cid,
                         token_pos: token_start..token_start + token_len,
                         len,
                         version,
-                    })
+                    }))
                 }
                 LongHeaderType::Retry => Ok(Self::Retry {
                     dst_cid,
@@ -612,6 +610,15 @@ impl PlainHeader {
     }
 }
 
+#[derive(Clone, Debug)]
+pub(crate) struct PlainInitialHeader {
+    pub(crate) dst_cid: ConnectionId,
+    pub(crate) src_cid: ConnectionId,
+    pub(crate) token_pos: Range<usize>,
+    pub(crate) len: u64,
+    pub(crate) version: u32,
+}
+
 // An encoded packet number
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub(crate) enum PacketNumber {