From 2d21680cf881a5e3f91f6567c378ca569e81b39d Mon Sep 17 00:00:00 2001
From: Phoenix Kahlo <kahlo.phoenix@gmail.com>
Date: Sat, 14 Dec 2024 13:42:27 -0600
Subject: [PATCH] proto: Make Connection externally use SideArgs

Server/client-specific args to proto::Connection::new are now passed
through a new SideArgs enum.
---
 quinn-proto/src/connection/mod.rs | 64 +++++++++++++++++++++++++------
 quinn-proto/src/endpoint.rs       | 28 ++++++--------
 2 files changed, 63 insertions(+), 29 deletions(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index fd4a7943d..f99c0f9f3 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -238,12 +238,10 @@ pub struct Connection {
 impl Connection {
     pub(crate) fn new(
         endpoint_config: Arc<EndpointConfig>,
-        server_config: Option<Arc<ServerConfig>>,
         config: Arc<TransportConfig>,
         init_cid: ConnectionId,
         loc_cid: ConnectionId,
         rem_cid: ConnectionId,
-        pref_addr_cid: Option<ConnectionId>,
         remote: SocketAddr,
         local_ip: Option<IpAddr>,
         crypto: Box<dyn crypto::Session>,
@@ -252,17 +250,11 @@ impl Connection {
         version: u32,
         allow_mtud: bool,
         rng_seed: [u8; 32],
-        path_validated: bool,
+        side_args: SideArgs,
     ) -> Self {
-        let connection_side = if let Some(server_config) = server_config.clone() {
-            ConnectionSide::Server { server_config }
-        } else {
-            assert!(pref_addr_cid.is_none());
-            assert!(path_validated);
-            ConnectionSide::Client {
-                token: Bytes::new(),
-            }
-        };
+        let pref_addr_cid = side_args.pref_addr_cid();
+        let path_validated = side_args.path_validated();
+        let connection_side = ConnectionSide::from(side_args);
         let side = connection_side.side();
         let initial_space = PacketSpace {
             crypto: Some(crypto.initial_keys(&init_cid, side)),
@@ -3655,6 +3647,54 @@ impl ConnectionSide {
     }
 }
 
+impl From<SideArgs> for ConnectionSide {
+    fn from(side: SideArgs) -> Self {
+        match side {
+            SideArgs::Client => Self::Client {
+                token: Bytes::new(),
+            },
+            SideArgs::Server {
+                server_config,
+                pref_addr_cid: _,
+                path_validated: _,
+            } => Self::Server { server_config },
+        }
+    }
+}
+
+/// Parameters to `Connection::new` specific to it being client-side or server-side
+pub(crate) enum SideArgs {
+    Client,
+    Server {
+        server_config: Arc<ServerConfig>,
+        pref_addr_cid: Option<ConnectionId>,
+        path_validated: bool,
+    },
+}
+
+impl SideArgs {
+    pub(crate) fn side(&self) -> Side {
+        match *self {
+            Self::Client { .. } => Side::Client,
+            Self::Server { .. } => Side::Server,
+        }
+    }
+
+    pub(crate) fn pref_addr_cid(&self) -> Option<ConnectionId> {
+        match *self {
+            Self::Client { .. } => None,
+            Self::Server { pref_addr_cid, .. } => pref_addr_cid,
+        }
+    }
+
+    pub(crate) fn path_validated(&self) -> bool {
+        match *self {
+            Self::Client { .. } => true,
+            Self::Server { path_validated, .. } => path_validated,
+        }
+    }
+}
+
 /// Reasons why a connection might be lost
 #[derive(Debug, Error, Clone, PartialEq, Eq)]
 pub enum ConnectionError {
diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
index a7de52eef..48de93807 100644
--- a/quinn-proto/src/endpoint.rs
+++ b/quinn-proto/src/endpoint.rs
@@ -18,7 +18,7 @@ use crate::{
     cid_generator::ConnectionIdGenerator,
     coding::BufMutExt,
     config::{ClientConfig, EndpointConfig, ServerConfig},
-    connection::{Connection, ConnectionError},
+    connection::{Connection, ConnectionError, SideArgs},
     crypto::{self, Keys, UnsupportedVersion},
     frame,
     packet::{
@@ -423,16 +423,14 @@ impl Endpoint {
             remote_id,
             loc_cid,
             remote_id,
-            None,
             FourTuple {
                 remote,
                 local_ip: None,
             },
             now,
             tls,
-            None,
             config.transport,
-            true,
+            SideArgs::Client,
         );
         Ok((ch, conn))
     }
@@ -660,13 +658,15 @@ impl Endpoint {
             dst_cid,
             loc_cid,
             src_cid,
-            pref_addr_cid,
             incoming.addresses,
             incoming.received_at,
             tls,
-            Some(server_config),
             transport_config,
-            remote_address_validated,
+            SideArgs::Server {
+                server_config,
+                pref_addr_cid,
+                path_validated: remote_address_validated,
+            },
         );
         self.index.insert_initial(dst_cid, ch);
 
@@ -829,28 +829,22 @@ impl Endpoint {
         init_cid: ConnectionId,
         loc_cid: ConnectionId,
         rem_cid: ConnectionId,
-        pref_addr_cid: Option<ConnectionId>,
         addresses: FourTuple,
         now: Instant,
         tls: Box<dyn crypto::Session>,
-        server_config: Option<Arc<ServerConfig>>,
         transport_config: Arc<TransportConfig>,
-        path_validated: bool,
+        side_args: SideArgs,
     ) -> Connection {
         let mut rng_seed = [0; 32];
         self.rng.fill_bytes(&mut rng_seed);
-        let side = match server_config.is_some() {
-            true => Side::Server,
-            false => Side::Client,
-        };
+        let side = side_args.side();
+        let pref_addr_cid = side_args.pref_addr_cid();
         let conn = Connection::new(
             self.config.clone(),
-            server_config,
             transport_config,
             init_cid,
             loc_cid,
             rem_cid,
-            pref_addr_cid,
             addresses.remote,
             addresses.local_ip,
             tls,
@@ -859,7 +853,7 @@ impl Endpoint {
             version,
             self.allow_mtud,
             rng_seed,
-            path_validated,
+            side_args,
         );
 
         let mut cids_issued = 0;