From c9879ba540ff73ea6e8859ad1eadef9e17cff255 Mon Sep 17 00:00:00 2001
From: Cameron Bytheway <bythewc@amazon.com>
Date: Wed, 22 Jan 2025 11:54:00 -0700
Subject: [PATCH] fix(s2n-quic-dc): derive crypto before opening TCP stream

---
 dc/s2n-quic-dc/src/stream/client/tokio.rs | 30 +++++++++++++++++------
 dc/s2n-quic-dc/src/stream/endpoint.rs     | 15 ++++--------
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/dc/s2n-quic-dc/src/stream/client/tokio.rs b/dc/s2n-quic-dc/src/stream/client/tokio.rs
index ceb345594..4cb23a3b4 100644
--- a/dc/s2n-quic-dc/src/stream/client/tokio.rs
+++ b/dc/s2n-quic-dc/src/stream/client/tokio.rs
@@ -9,6 +9,7 @@ use crate::{
         endpoint,
         environment::tokio::{self as env, Environment},
         socket::Protocol,
+        TransportFeatures,
     },
 };
 use std::{io, net::SocketAddr};
@@ -29,12 +30,15 @@ where
     // ensure we have a secret for the peer
     let peer = handshake.await?;
 
+    let (crypto, parameters) = peer.pair(&TransportFeatures::UDP);
+
     let stream = endpoint::open_stream(
         env,
-        peer,
+        peer.map(),
+        crypto,
+        parameters,
         env::UdpUnbound(acceptor_addr.into()),
         subscriber,
-        None,
     )?;
 
     // build the stream inside the application context
@@ -60,7 +64,14 @@ where
     Sub: event::Subscriber,
 {
     // Race TCP handshake with the TLS handshake
-    let (socket, peer) = tokio::try_join!(TcpStream::connect(acceptor_addr), handshake,)?;
+    let handshake = async {
+        let peer = handshake.await?;
+        let (crypto, parameters) = peer.pair(&TransportFeatures::TCP);
+        Ok((peer, crypto, parameters))
+    };
+    // poll the crypto first so the server can read the first packet on accept in the happy path
+    let ((peer, crypto, parameters), socket) =
+        tokio::try_join!(handshake, TcpStream::connect(acceptor_addr))?;
 
     // Make sure TCP_NODELAY is set
     let _ = socket.set_nodelay(true);
@@ -77,14 +88,15 @@ where
 
     let stream = endpoint::open_stream(
         env,
-        peer,
+        peer.map(),
+        crypto,
+        parameters,
         env::TcpRegistered {
             socket,
             peer_addr,
             local_port,
         },
         subscriber,
-        None,
     )?;
 
     // build the stream inside the application context
@@ -114,16 +126,20 @@ where
 {
     let local_port = socket.local_addr()?.port();
     let peer_addr = socket.peer_addr()?.into();
+
+    let (crypto, parameters) = peer.pair(&TransportFeatures::TCP);
+
     let stream = endpoint::open_stream(
         env,
-        peer,
+        peer.map(),
+        crypto,
+        parameters,
         env::TcpRegistered {
             socket,
             peer_addr,
             local_port,
         },
         subscriber,
-        None,
     )?;
 
     // build the stream inside the application context
diff --git a/dc/s2n-quic-dc/src/stream/endpoint.rs b/dc/s2n-quic-dc/src/stream/endpoint.rs
index 3554bde50..1b914f842 100644
--- a/dc/s2n-quic-dc/src/stream/endpoint.rs
+++ b/dc/s2n-quic-dc/src/stream/endpoint.rs
@@ -4,7 +4,7 @@
 use crate::{
     event::{self, api::Subscriber as _, IntoEvent as _},
     msg, packet,
-    path::secret::{self, map, Map},
+    path::secret::{self, Map},
     random::Random,
     stream::{
         application,
@@ -35,21 +35,16 @@ pub struct AcceptError<Peer> {
 #[inline]
 pub fn open_stream<Env, P>(
     env: &Env,
-    entry: map::Peer,
+    map: &Map,
+    crypto: secret::map::Bidirectional,
+    parameters: dc::ApplicationParams,
     peer: P,
     subscriber: Env::Subscriber,
-    parameter_override: Option<&dyn Fn(dc::ApplicationParams) -> dc::ApplicationParams>,
 ) -> Result<application::Builder<Env::Subscriber>>
 where
     Env: Environment,
     P: Peer<Env>,
 {
-    let (crypto, mut parameters) = entry.pair(&peer.features());
-
-    if let Some(o) = parameter_override {
-        parameters = o(parameters);
-    }
-
     let key_id = crypto.credentials.key_id;
     let stream_id = packet::stream::Id {
         key_id,
@@ -74,7 +69,7 @@ where
         stream_id,
         None,
         crypto,
-        entry.map(),
+        map,
         parameters,
         None,
         None,