access application context from tls

Author: taikulawo

quic/s2n-quic-core/src/crypto/tls.rs

@@ -5,7 +5,7 @@
 use alloc::vec::Vec;
 #[cfg(feature = "alloc")]
 pub use bytes::{Bytes, BytesMut};
-use core::fmt::Debug;
+use core::{any::Any, fmt::Debug};
 use zerocopy::{AsBytes, FromBytes, FromZeroes, Unaligned};
 
 mod error;
@@ -130,6 +130,8 @@ pub trait Context<Crypto: crate::crypto::CryptoSuite> {
     //# peer's Finished message.
     fn on_handshake_complete(&mut self) -> Result<(), crate::transport::Error>;
 
+    /// Transfer application context from TLS connection to quic connection
+    fn on_application_context(&mut self, _context: Option<Box<dyn Any + Send + Sync>>) {}
     fn on_tls_exporter_ready(
         &mut self,
         session: &impl TlsSession,

quic/s2n-quic-tls/src/session.rs

@@ -16,6 +16,7 @@ use s2n_tls::{
     enums::{Blinding, Mode},
     error::{Error, ErrorType},
 };
+use std::any::Any;
 
 #[derive(Debug)]
 pub struct Session {
@@ -155,6 +156,10 @@ impl tls::Session for Session {
                     context.on_tls_exporter_ready(self)?;
                     self.handshake_complete = true;
                 }
+                // TODO Add new s2n-tls new api, take and put in quic::connection
+                let ctx: Option<Box<dyn Any + Send + Sync>> =
+                    self.connection.application_context_owned();
+                context.on_application_context(ctx);
                 Poll::Ready(Ok(()))
             }
             Poll::Ready(Err(e)) => Poll::Ready(Err(e

quic/s2n-quic-transport/src/connection/api_provider.rs

@@ -21,6 +21,7 @@ use s2n_quic_core::{
     query::{Query, QueryMut},
     stream::{ops, StreamId, StreamType},
 };
+use std::any::Any;
 
 /// A dynamically dispatched connection API
 pub(crate) type ConnectionApi = Arc<dyn ConnectionApiProvider>;
@@ -58,6 +59,10 @@ pub(crate) trait ConnectionApiProvider: Sync + Send {
 
     fn application_protocol(&self) -> Result<Bytes, connection::Error>;
 
+    fn take_application_context(
+        &mut self,
+    ) -> Result<Option<Box<dyn Any + Send + Sync>>, connection::Error>;
+
     fn id(&self) -> u64;
 
     fn ping(&self) -> Result<(), connection::Error>;

quic/s2n-quic-transport/src/connection/connection_container.rs

@@ -41,6 +41,7 @@ use s2n_quic_core::{
     transport,
 };
 use smallvec::SmallVec;
+use std::any::Any;
 
 // Intrusive list adapter for managing the list of `done` connections
 intrusive_adapter!(DoneConnectionsAdapter<C, L> = Arc<ConnectionNode<C, L>>: ConnectionNode<C, L> {
@@ -310,6 +311,12 @@ impl<C: connection::Trait, L: connection::Lock<C>> ConnectionApiProvider for Con
         self.api_read_call(|conn| Ok(conn.application_protocol()))
     }
 
+    fn take_application_context(
+        &mut self,
+    ) -> Result<Option<Box<dyn Any + Send + Sync>>, connection::Error> {
+        self.api_write_call(|conn| Ok(conn.take_application_context()))
+    }
+
     fn id(&self) -> u64 {
         self.internal_connection_id.into()
     }

quic/s2n-quic-transport/src/connection/connection_container/tests.rs

@@ -62,7 +62,9 @@ impl connection::Trait for TestConnection {
     fn new(_params: connection::Parameters<Self::Config>) -> Result<Self, connection::Error> {
         Ok(Self::default())
     }
-
+    fn take_application_context(&mut self) -> Option<Box<dyn std::any::Any + Send + Sync>>{
+        todo!()
+    }
     fn internal_connection_id(&self) -> InternalConnectionId {
         todo!()
     }

quic/s2n-quic-transport/src/connection/connection_impl.rs

@@ -61,6 +61,7 @@ use s2n_quic_core::{
     time::{timer, Timestamp},
     transport,
 };
+use std::any::Any;
 
 /// Possible states for handing over a connection from the endpoint to the
 /// application.
@@ -1944,6 +1945,10 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
         self.space_manager.application_protocol.clone()
     }
 
+    fn take_application_context(&mut self) -> Option<Box<dyn Any + Send + Sync>> {
+        self.space_manager.application_context.take()
+    }
+
     fn ping(&mut self) -> Result<(), connection::Error> {
         self.error?;
 

quic/s2n-quic-transport/src/connection/connection_trait.rs

@@ -36,6 +36,7 @@ use s2n_quic_core::{
     query,
     time::Timestamp,
 };
+use std::any::Any;
 
 /// A trait which represents an internally used `Connection`
 pub trait ConnectionTrait: 'static + Send + Sized {
@@ -505,6 +506,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
             >,
             &path::Path<Self::Config>,
         );
+    fn take_application_context(&mut self) -> Option<Box<dyn Any + Send + Sync>>;
 }
 
 /// A lock that synchronizes connection state between the QUIC endpoint thread and application

quic/s2n-quic-transport/src/space/mod.rs

@@ -16,6 +16,7 @@ use core::{
     ops::RangeInclusive,
     task::{Poll, Waker},
 };
+use std::any::Any;
 use s2n_codec::DecoderBufferMut;
 use s2n_quic_core::{
     application::ServerName,
@@ -63,6 +64,7 @@ pub struct PacketSpaceManager<Config: endpoint::Config> {
     retry_cid: Option<Box<PeerId>>,
     initial: Option<Box<InitialSpace<Config>>>,
     handshake: Option<Box<HandshakeSpace<Config>>>,
+    pub application_context: Option<Box<dyn Any + Send + Sync>>,
     application: Option<Box<ApplicationSpace<Config>>>,
     zero_rtt_crypto:
         Option<Box<<<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::ZeroRttKey>>,
@@ -124,6 +126,7 @@ impl<Config: endpoint::Config> PacketSpaceManager<Config> {
                 session,
                 initial_cid,
             }),
+            application_context: None,
             retry_cid: None,
             initial: Some(Box::new(InitialSpace::new(
                 initial_key,
@@ -258,6 +261,7 @@ impl<Config: endpoint::Config> PacketSpaceManager<Config> {
                 handshake: &mut self.handshake,
                 application: &mut self.application,
                 zero_rtt_crypto: &mut self.zero_rtt_crypto,
+                application_context: &mut self.application_context,
                 path_manager,
                 handshake_status: &mut self.handshake_status,
                 local_id_registry,
@@ -303,6 +307,7 @@ impl<Config: endpoint::Config> PacketSpaceManager<Config> {
                 retry_cid: self.retry_cid.as_deref(),
                 initial: &mut self.initial,
                 handshake: &mut self.handshake,
+                application_context: &mut self.application_context,
                 application: &mut self.application,
                 zero_rtt_crypto: &mut self.zero_rtt_crypto,
                 path_manager,

quic/s2n-quic-transport/src/space/session_context.rs

@@ -41,6 +41,7 @@ use s2n_quic_core::{
         Error,
     },
 };
+use std::any::Any;
 
 pub struct SessionContext<'a, Config: endpoint::Config, Pub: event::ConnectionPublisher> {
     pub now: Timestamp,
@@ -62,6 +63,7 @@ pub struct SessionContext<'a, Config: endpoint::Config, Pub: event::ConnectionPu
     pub publisher: &'a mut Pub,
     pub datagram: &'a mut Config::DatagramEndpoint,
     pub dc: &'a mut Config::DcEndpoint,
+    pub application_context: &'a mut Option<Box<dyn Any + Send + Sync>>,
 }
 
 impl<Config: endpoint::Config, Pub: event::ConnectionPublisher> SessionContext<'_, Config, Pub> {
@@ -697,4 +699,7 @@ impl<Config: endpoint::Config, Pub: event::ConnectionPublisher>
 
         Ok(())
     }
+    fn on_application_context(&mut self, context: Option<Box<dyn Any + Send + Sync>>) {
+        *self.application_context = context;
+    }
 }