From b739b0f79860777d8292caca75d56b7fdb5586c1 Mon Sep 17 00:00:00 2001
From: Mark Rousskov <thismark@amazon.com>
Date: Mon, 27 Jan 2025 18:25:16 +0000
Subject: [PATCH] Revert "fix(s2n-quic-dc): derive crypto before opening TCP
 stream (#2451)"

In further testing, this introduces too much reordering under high
concurrency for our dedup tracking to mitigate. connect() can "complete"
in reverse order of starts (e.g., due to queueing of the wake events)
which leads to huge gaps (up to concurrency) in the chosen key IDs.

We can likely get the best of both worlds by deriving crypto beforehand
but only associating it with a stream just-in-time as connect()s
complete (keeping a per-peer pool just for opening connections) but
that's a more invasive change, for now just revert this.

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

diff --git a/dc/s2n-quic-dc/src/stream/client/tokio.rs b/dc/s2n-quic-dc/src/stream/client/tokio.rs
index 4cb23a3b4..ceb345594 100644
--- a/dc/s2n-quic-dc/src/stream/client/tokio.rs
+++ b/dc/s2n-quic-dc/src/stream/client/tokio.rs
@@ -9,7 +9,6 @@ use crate::{
         endpoint,
         environment::tokio::{self as env, Environment},
         socket::Protocol,
-        TransportFeatures,
     },
 };
 use std::{io, net::SocketAddr};
@@ -30,15 +29,12 @@ 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.map(),
-        crypto,
-        parameters,
+        peer,
         env::UdpUnbound(acceptor_addr.into()),
         subscriber,
+        None,
     )?;
 
     // build the stream inside the application context
@@ -64,14 +60,7 @@ where
     Sub: event::Subscriber,
 {
     // Race TCP handshake with the TLS 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))?;
+    let (socket, peer) = tokio::try_join!(TcpStream::connect(acceptor_addr), handshake,)?;
 
     // Make sure TCP_NODELAY is set
     let _ = socket.set_nodelay(true);
@@ -88,15 +77,14 @@ where
 
     let stream = endpoint::open_stream(
         env,
-        peer.map(),
-        crypto,
-        parameters,
+        peer,
         env::TcpRegistered {
             socket,
             peer_addr,
             local_port,
         },
         subscriber,
+        None,
     )?;
 
     // build the stream inside the application context
@@ -126,20 +114,16 @@ 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.map(),
-        crypto,
-        parameters,
+        peer,
         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 1b914f842..3554bde50 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},
+    path::secret::{self, map, Map},
     random::Random,
     stream::{
         application,
@@ -35,16 +35,21 @@ pub struct AcceptError<Peer> {
 #[inline]
 pub fn open_stream<Env, P>(
     env: &Env,
-    map: &Map,
-    crypto: secret::map::Bidirectional,
-    parameters: dc::ApplicationParams,
+    entry: map::Peer,
     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,
@@ -69,7 +74,7 @@ where
         stream_id,
         None,
         crypto,
-        map,
+        entry.map(),
         parameters,
         None,
         None,