feat(s2n-quic-dc): Add application_data to stream API for peer identity extraction (#3083)

Author: jbbbriant

dc/s2n-quic-dc/src/path/secret/map.rs

@@ -177,15 +177,20 @@ impl Map {
         queue_id: Option<VarInt>,
         features: &TransportFeatures,
         control_out: &mut Vec<u8>,
-    ) -> Option<(entry::Bidirectional, dc::ApplicationParams)> {
+    ) -> Option<(
+        entry::Bidirectional,
+        dc::ApplicationParams,
+        Option<entry::ApplicationData>,
+    )> {
         let entry = self
             .store
             .pre_authentication(credentials, queue_id, control_out)?;
 
         let params = entry.parameters();
         let keys = entry.bidi_remote(self.clone(), credentials, queue_id, features);
 
-        Some((keys, params))
+        let application_data = entry.application_data().clone();
+        Some((keys, params, application_data))
     }
 
     pub fn secret_for_credentials(

dc/s2n-quic-dc/src/stream/application.rs

@@ -4,6 +4,7 @@
 use crate::{
     credentials::Id,
     event::{self, EndpointPublisher as _},
+    path::secret::map::ApplicationData,
     stream::{
         recv::application::{self as recv, Reader},
         send::application::{self as send, Writer},
@@ -177,6 +178,12 @@ where
         self.read.path_secret_id()
     }
 
+    /// Returns the application data associated with the path secret, if available.
+    #[inline]
+    pub fn path_application_data(&self) -> Option<&ApplicationData> {
+        self.read.path_application_data()
+    }
+
     /// Returns the validated peer certificate chain, if available.
     ///
     /// Currently this is only available for TLS streams, but in the future it may be opt-in

dc/s2n-quic-dc/src/stream/endpoint.rs

@@ -4,7 +4,7 @@
 use crate::{
     event::{self, api::Subscriber as _, IntoEvent as _},
     packet,
-    path::secret::{self, map, Map},
+    path::secret::{self, map, map::ApplicationData, Map},
     random::Random,
     stream::{
         self, application,
@@ -75,6 +75,7 @@ where
         endpoint::Type::Client,
         subscriber,
         subscriber_ctx,
+        None,
     )
 }
 
@@ -84,10 +85,17 @@ pub fn derive_stream_credentials(
     map: &Map,
     features: &stream::TransportFeatures,
     secret_control: &mut Vec<u8>,
-) -> Result<(secret::map::Bidirectional, dc::ApplicationParams), io::Error> {
+) -> Result<
+    (
+        secret::map::Bidirectional,
+        dc::ApplicationParams,
+        Option<ApplicationData>,
+    ),
+    io::Error,
+> {
     let credentials = &packet.credentials;
 
-    let Some((crypto, parameters)) = map.pair_for_credentials(
+    let Some((crypto, parameters, application_data)) = map.pair_for_credentials(
         credentials,
         packet.source_queue_id,
         features,
@@ -100,7 +108,7 @@ pub fn derive_stream_credentials(
         return Err(error);
     };
 
-    Ok((crypto, parameters))
+    Ok((crypto, parameters, application_data))
 }
 
 #[inline]
@@ -115,6 +123,7 @@ pub fn accept_stream<Env, P>(
     crypto: secret::map::Bidirectional,
     mut parameters: dc::ApplicationParams,
     secret_control: Vec<u8>,
+    application_data: Option<ApplicationData>,
 ) -> Result<application::Builder<Env::Subscriber>, AcceptError>
 where
     Env: Environment,
@@ -151,6 +160,7 @@ where
         endpoint::Type::Server,
         subscriber,
         subscriber_ctx,
+        application_data,
     );
 
     match res {
@@ -180,6 +190,7 @@ fn build_stream<Env, P>(
     endpoint_type: endpoint::Type,
     subscriber: Env::Subscriber,
     subscriber_ctx: <Env::Subscriber as event::Subscriber>::ConnectionContext,
+    application_data: Option<ApplicationData>,
 ) -> Result<application::Builder<Env::Subscriber>>
 where
     Env: Environment,
@@ -292,6 +303,7 @@ where
     let shared = Arc::new(shared::Shared {
         receiver: reader,
         sender: writer.0,
+        application_data,
         common,
         crypto,
     });

dc/s2n-quic-dc/src/stream/recv/application.rs

@@ -6,6 +6,7 @@ use crate::{
     credentials::Id,
     event::{self, ConnectionPublisher as _},
     msg,
+    path::secret::map::ApplicationData,
     stream::{
         recv, runtime,
         shared::{AcceptState, ArcShared, ShutdownKind},
@@ -140,6 +141,11 @@ where
         &self.0.shared.credentials().id
     }
 
+    #[inline]
+    pub fn path_application_data(&self) -> Option<&ApplicationData> {
+        self.0.shared.application_data()
+    }
+
     #[inline]
     pub fn protocol(&self) -> socket::Protocol {
         self.0.sockets.protocol()

dc/s2n-quic-dc/src/stream/server/tokio/tcp/worker.rs

@@ -411,7 +411,7 @@ where
 
             let mut secret_control = vec![];
 
-            let (crypto, parameters) = match endpoint::derive_stream_credentials(
+            let (crypto, parameters, application_data) = match endpoint::derive_stream_credentials(
                 &initial_packet,
                 &context.secrets,
                 &TransportFeatures::TCP,
@@ -463,6 +463,7 @@ where
                 crypto,
                 parameters,
                 secret_control,
+                application_data,
             ) {
                 Ok(stream) => stream,
                 Err(error) => {

dc/s2n-quic-dc/src/stream/server/tokio/udp.rs

@@ -119,7 +119,7 @@ where
         let peer = env::udp::Owned(remote_addr, recv_buffer);
 
         let mut secret_control = vec![];
-        let (crypto, parameters) = match endpoint::derive_stream_credentials(
+        let (crypto, parameters, application_data) = match endpoint::derive_stream_credentials(
             &packet,
             &self.secrets,
             &TransportFeatures::UDP,
@@ -148,6 +148,7 @@ where
             crypto,
             parameters,
             secret_control,
+            application_data,
         ) {
             Ok(stream) => stream,
             Err(error) => {

dc/s2n-quic-dc/src/stream/server/tokio/uds.rs

@@ -163,6 +163,9 @@ where
             crypto,
             decoded_packet.application_params().clone(),
             secret_control,
+            // application_data is not available for UDS streams because the in-application
+            // map does not contain the credential ID used by the forwarding process.
+            None,
         ) {
             Ok(stream) => stream,
             Err(error) => {

dc/s2n-quic-dc/src/stream/server/udp.rs

@@ -145,7 +145,7 @@ where
         };
 
         let mut secret_control = vec![];
-        let (crypto, parameters) = match endpoint::derive_stream_credentials(
+        let (crypto, parameters, application_data) = match endpoint::derive_stream_credentials(
             &self.packet,
             &self.secrets,
             &TransportFeatures::UDP,
@@ -177,6 +177,7 @@ where
             crypto,
             parameters,
             secret_control,
+            application_data,
         ) {
             Ok(stream) => stream,
             Err(error) => {

dc/s2n-quic-dc/src/stream/shared.rs

@@ -6,6 +6,7 @@ use crate::{
     credentials::Credentials,
     event::{self, IntoEvent as _},
     packet::stream,
+    path::secret::map::ApplicationData,
     stream::{
         recv::shared as recv,
         send::{application, shared as send},
@@ -73,6 +74,7 @@ where
     pub receiver: recv::State,
     pub sender: send::State,
     pub crypto: Crypto,
+    pub application_data: Option<ApplicationData>,
     pub common: Common<Subscriber, Clk>,
 }
 
@@ -188,6 +190,11 @@ where
             &*self.common.fixed.credentials.get()
         }
     }
+
+    #[inline]
+    pub fn application_data(&self) -> Option<&ApplicationData> {
+        self.application_data.as_ref()
+    }
 }
 
 impl<Sub, C> ops::Deref for Shared<Sub, C>

dc/s2n-quic-dc/src/stream/tests/api.rs

@@ -2,6 +2,14 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::testing::dcquic::tcp::*;
+use crate::{
+    path::secret::{stateless_reset::Signer, Map},
+    psk::{client, server},
+    stream::{socket::Protocol, testing::bind_pair},
+    testing::{init_tracing, server_name, NoopSubscriber, TestTlsProvider},
+};
+use s2n_quic_core::time::StdClock;
+use std::sync::Arc;
 
 #[tokio::test]
 async fn context_accessible() {
@@ -14,3 +22,94 @@ async fn context_accessible() {
     read.query_event_context(|_: &()| ()).unwrap();
     write.query_event_context(|_: &()| ()).unwrap();
 }
+
+#[tokio::test]
+async fn path_application_data_none_by_default() {
+    let context = Context::new().await;
+    let (_client, server) = context.pair().await;
+
+    // Without a registered make_application_data callback, the server stream
+    // should have no application_data.
+    assert!(server.path_application_data().is_none());
+}
+
+#[tokio::test]
+async fn path_application_data_available_on_server_stream() {
+    init_tracing();
+
+    let tls = TestTlsProvider {};
+    let sub = NoopSubscriber {};
+
+    // Build server Map with application_data callback
+    let server_map = Map::new(
+        Signer::new(b"default"),
+        50_000,
+        false,
+        StdClock::default(),
+        sub.clone(),
+    );
+    server_map.register_make_application_data(Box::new(|_session| {
+        let data: Arc<dyn std::any::Any + Send + Sync> = Arc::new(42u64);
+        Ok(Some(data))
+    }));
+
+    let server_prov = server::Provider::builder()
+        .start(
+            "[::1]:0".parse().unwrap(),
+            tls.clone(),
+            sub.clone(),
+            server_map,
+        )
+        .await
+        .unwrap();
+
+    let client_map = Map::new(
+        Signer::new(b"default"),
+        50_000,
+        false,
+        StdClock::default(),
+        sub.clone(),
+    );
+    let client_prov = client::Provider::builder()
+        .with_success_jitter(std::time::Duration::ZERO)
+        .start(
+            "[::]:0".parse().unwrap(),
+            client_map,
+            tls,
+            sub.clone(),
+            server_name(),
+        )
+        .unwrap();
+
+    let (client, server) = bind_pair(
+        Protocol::Tcp,
+        "127.0.0.1:0".parse().unwrap(),
+        client_prov,
+        server_prov,
+    );
+
+    let acceptor_addr = server.acceptor_addr().expect("acceptor_addr");
+    let handshake_addr = server.handshake_addr().expect("handshake_addr");
+
+    let (client_stream, server_stream) = tokio::join!(
+        async {
+            client
+                .connect(handshake_addr, acceptor_addr, server_name())
+                .await
+                .unwrap()
+        },
+        async { server.accept().await.unwrap().0 }
+    );
+
+    // Client doesn't have application_data
+    assert!(client_stream.path_application_data().is_none());
+
+    // Server stream should have the application_data we registered
+    let app_data = server_stream
+        .path_application_data()
+        .expect("should have application_data");
+    let value = app_data
+        .downcast_ref::<u64>()
+        .expect("should downcast to u64");
+    assert_eq!(*value, 42u64);
+}