fix(s2n-quic-transport): Fixes s2n-quic TLS state incongruity when s2n-tls returns pending frequently (#2574)

Author: maddeleine

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

@@ -17,6 +17,9 @@ pub mod testing;
 #[cfg(all(feature = "alloc", any(test, feature = "testing")))]
 pub mod null;
 
+#[cfg(feature = "alloc")]
+pub mod slow_tls;
+
 /// Holds all application parameters which are exchanged within the TLS handshake.
 #[derive(Debug)]
 pub struct ApplicationParameters<'a> {

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

@@ -0,0 +1,217 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+use crate::{
+    application,
+    crypto::{tls, CryptoSuite},
+    transport,
+};
+use alloc::{boxed::Box, vec::Vec};
+use core::{any::Any, task::Poll};
+
+const DEFER_COUNT: u8 = 3;
+
+pub struct SlowEndpoint<E: tls::Endpoint> {
+    endpoint: E,
+}
+
+impl<E: tls::Endpoint> SlowEndpoint<E> {
+    pub fn new(endpoint: E) -> Self {
+        SlowEndpoint { endpoint }
+    }
+}
+
+impl<E: tls::Endpoint> tls::Endpoint for SlowEndpoint<E> {
+    type Session = SlowSession<<E as tls::Endpoint>::Session>;
+
+    fn new_server_session<Params: s2n_codec::EncoderValue>(
+        &mut self,
+        transport_parameters: &Params,
+    ) -> Self::Session {
+        let inner_session = self.endpoint.new_server_session(transport_parameters);
+        SlowSession {
+            defer: DEFER_COUNT,
+            inner_session,
+        }
+    }
+
+    fn new_client_session<Params: s2n_codec::EncoderValue>(
+        &mut self,
+        transport_parameters: &Params,
+        server_name: application::ServerName,
+    ) -> Self::Session {
+        let inner_session = self
+            .endpoint
+            .new_client_session(transport_parameters, server_name);
+        SlowSession {
+            defer: DEFER_COUNT,
+            inner_session,
+        }
+    }
+
+    fn max_tag_length(&self) -> usize {
+        self.endpoint.max_tag_length()
+    }
+}
+
+// SlowSession is a test TLS provider that is slow, namely, for each call to poll,
+// it returns Poll::Pending several times before actually polling the real TLS library.
+// This is used in an integration test to assert that our code is correct in the event
+// of any random pendings/wakeups that might occur when negotiating TLS.
+#[derive(Debug)]
+pub struct SlowSession<S: tls::Session> {
+    defer: u8,
+    inner_session: S,
+}
+
+impl<S: tls::Session> tls::Session for SlowSession<S> {
+    #[inline]
+    fn poll<W>(&mut self, context: &mut W) -> Poll<Result<(), transport::Error>>
+    where
+        W: tls::Context<Self>,
+    {
+        // Self-wake and return Pending if defer is non-zero
+        if let Some(d) = self.defer.checked_sub(1) {
+            self.defer = d;
+            context.waker().wake_by_ref();
+            return Poll::Pending;
+        }
+
+        // Otherwise we'll call the function to actually make progress
+        // in the TLS handshake and set up to defer again the next time
+        // we're here.
+        self.defer = DEFER_COUNT;
+        self.inner_session.poll(&mut SlowContext(context))
+    }
+}
+
+impl<S: tls::Session> CryptoSuite for SlowSession<S> {
+    type HandshakeKey = <S as CryptoSuite>::HandshakeKey;
+    type HandshakeHeaderKey = <S as CryptoSuite>::HandshakeHeaderKey;
+    type InitialKey = <S as CryptoSuite>::InitialKey;
+    type InitialHeaderKey = <S as CryptoSuite>::InitialHeaderKey;
+    type ZeroRttKey = <S as CryptoSuite>::ZeroRttKey;
+    type ZeroRttHeaderKey = <S as CryptoSuite>::ZeroRttHeaderKey;
+    type OneRttKey = <S as CryptoSuite>::OneRttKey;
+    type OneRttHeaderKey = <S as CryptoSuite>::OneRttHeaderKey;
+    type RetryKey = <S as CryptoSuite>::RetryKey;
+}
+
+struct SlowContext<'a, Inner>(&'a mut Inner);
+
+impl<I, S: tls::Session> tls::Context<S> for SlowContext<'_, I>
+where
+    I: tls::Context<SlowSession<S>>,
+{
+    fn on_client_application_params(
+        &mut self,
+        client_params: tls::ApplicationParameters,
+        server_params: &mut Vec<u8>,
+    ) -> Result<(), transport::Error> {
+        self.0
+            .on_client_application_params(client_params, server_params)
+    }
+
+    fn on_handshake_keys(
+        &mut self,
+        key: <S as CryptoSuite>::HandshakeKey,
+        header_key: <S as CryptoSuite>::HandshakeHeaderKey,
+    ) -> Result<(), transport::Error> {
+        self.0.on_handshake_keys(key, header_key)
+    }
+
+    fn on_zero_rtt_keys(
+        &mut self,
+        key: <S>::ZeroRttKey,
+        header_key: <S>::ZeroRttHeaderKey,
+        application_parameters: tls::ApplicationParameters,
+    ) -> Result<(), transport::Error> {
+        self.0
+            .on_zero_rtt_keys(key, header_key, application_parameters)
+    }
+
+    fn on_one_rtt_keys(
+        &mut self,
+        key: <S>::OneRttKey,
+        header_key: <S>::OneRttHeaderKey,
+        application_parameters: tls::ApplicationParameters,
+    ) -> Result<(), transport::Error> {
+        self.0
+            .on_one_rtt_keys(key, header_key, application_parameters)
+    }
+
+    fn on_server_name(
+        &mut self,
+        server_name: application::ServerName,
+    ) -> Result<(), transport::Error> {
+        self.0.on_server_name(server_name)
+    }
+
+    fn on_application_protocol(
+        &mut self,
+        application_protocol: tls::Bytes,
+    ) -> Result<(), transport::Error> {
+        self.0.on_application_protocol(application_protocol)
+    }
+
+    fn on_handshake_complete(&mut self) -> Result<(), transport::Error> {
+        self.0.on_handshake_complete()
+    }
+
+    fn on_tls_exporter_ready(
+        &mut self,
+        session: &impl tls::TlsSession,
+    ) -> Result<(), transport::Error> {
+        self.0.on_tls_exporter_ready(session)
+    }
+
+    fn receive_initial(&mut self, max_len: Option<usize>) -> Option<tls::Bytes> {
+        self.0.receive_initial(max_len)
+    }
+
+    fn receive_handshake(&mut self, max_len: Option<usize>) -> Option<tls::Bytes> {
+        self.0.receive_handshake(max_len)
+    }
+
+    fn receive_application(&mut self, max_len: Option<usize>) -> Option<tls::Bytes> {
+        self.0.receive_application(max_len)
+    }
+
+    fn can_send_initial(&self) -> bool {
+        self.0.can_send_initial()
+    }
+
+    fn send_initial(&mut self, transmission: tls::Bytes) {
+        self.0.send_initial(transmission);
+    }
+
+    fn can_send_handshake(&self) -> bool {
+        self.0.can_send_handshake()
+    }
+
+    fn send_handshake(&mut self, transmission: tls::Bytes) {
+        self.0.send_handshake(transmission);
+    }
+
+    fn can_send_application(&self) -> bool {
+        self.0.can_send_application()
+    }
+
+    fn send_application(&mut self, transmission: tls::Bytes) {
+        self.0.send_application(transmission)
+    }
+
+    fn waker(&self) -> &core::task::Waker {
+        self.0.waker()
+    }
+
+    fn on_key_exchange_group(
+        &mut self,
+        named_group: tls::NamedGroup,
+    ) -> Result<(), transport::Error> {
+        self.0.on_key_exchange_group(named_group)
+    }
+
+    fn on_tls_context(&mut self, context: Box<dyn Any + Send>) {
+        self.0.on_tls_context(context)
+    }
+}

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

@@ -1141,6 +1141,16 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
         // check if crypto progress can be made
         self.update_crypto_state(timestamp, subscriber, datagram, dc, conn_limits)?;
 
+        if self.space_manager.handshake().is_some() && self.space_manager.is_handshake_confirmed() {
+            let mut publisher = self.event_context.publisher(timestamp, subscriber);
+
+            //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.2
+            //# An endpoint MUST discard its handshake keys when the TLS handshake is
+            //# confirmed (Section 4.1.2).
+            self.space_manager
+                .discard_handshake(&mut self.path_manager, &mut publisher);
+        }
+
         // return an error if the application set one
         self.error?;
 

quic/s2n-quic/src/tests.rs

@@ -27,6 +27,7 @@ mod recorder;
 mod connection_limits;
 mod resumption;
 mod setup;
+mod slow_tls;
 use setup::*;
 
 mod blackhole;

quic/s2n-quic/src/tests/setup.rs

@@ -242,6 +242,28 @@ mod mtls {
     }
 }
 
+mod slow_tls {
+    use crate::provider::tls::Provider;
+    use s2n_quic_core::crypto::tls::{slow_tls::SlowEndpoint, Endpoint};
+    pub struct SlowTlsProvider<E: Endpoint> {
+        pub endpoint: E,
+    }
+
+    impl<E: Endpoint> Provider for SlowTlsProvider<E> {
+        type Server = SlowEndpoint<E>;
+        type Client = SlowEndpoint<E>;
+        type Error = String;
+
+        fn start_server(self) -> Result<Self::Server, Self::Error> {
+            Ok(SlowEndpoint::new(self.endpoint))
+        }
+
+        fn start_client(self) -> Result<Self::Client, Self::Error> {
+            Ok(SlowEndpoint::new(self.endpoint))
+        }
+    }
+}
+
 #[cfg(feature = "s2n-quic-tls")]
 mod resumption {
     use super::*;
@@ -326,3 +348,6 @@ pub use mtls::*;
 
 #[cfg(feature = "s2n-quic-tls")]
 pub use resumption::*;
+
+#[cfg(not(feature = "provider-tls-fips"))]
+pub use slow_tls::SlowTlsProvider;

quic/s2n-quic/src/tests/slow_tls.rs

@@ -0,0 +1,47 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#[test]
+#[cfg(not(feature = "provider-tls-fips"))]
+fn slow_default_tls() {
+    use super::*;
+    use crate::provider::tls::default;
+    use s2n_quic_core::crypto::tls::testing::certificates::{CERT_PEM, KEY_PEM};
+
+    let model = Model::default();
+
+    let server_endpoint = default::Server::builder()
+        .with_certificate(CERT_PEM, KEY_PEM)
+        .unwrap()
+        .build()
+        .unwrap();
+    let slow_server = SlowTlsProvider {
+        endpoint: server_endpoint,
+    };
+
+    let client_endpoint = default::Client::builder()
+        .with_certificate(CERT_PEM)
+        .unwrap()
+        .build()
+        .unwrap();
+    let slow_client = SlowTlsProvider {
+        endpoint: client_endpoint,
+    };
+
+    test(model, |handle| {
+        let server = Server::builder()
+            .with_io(handle.builder().build()?)?
+            .with_tls(slow_server)?
+            .start()?;
+
+        let client = Client::builder()
+            .with_io(handle.builder().build().unwrap())?
+            .with_tls(slow_client)?
+            .start()?;
+        let addr = start_server(server)?;
+        start_client(client, addr, Data::new(1000))?;
+
+        Ok(addr)
+    })
+    .unwrap();
+}