s2n-quic-dc 0.83.0

// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

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},
        shared::ArcShared,
        socket,
    },
};
use core::{fmt, time::Duration};
use s2n_quic_core::{buffer, time::Timestamp};
use std::{io, net::SocketAddr};

pub struct Builder<Sub: event::Subscriber> {
    pub read: recv::Builder<Sub>,
    pub write: send::Builder<Sub>,
    pub shared: ArcShared<Sub>,
    pub sockets: Box<dyn socket::application::Builder>,
    /// Timestamp of accept(2)/connect(2) syscall -- either end of that call or start, depending on
    /// the details of our internal implementation.
    pub kernel_start_time: Timestamp,
    /// Timestamp of dcQUIC enqueuing the stream to the application (for server sockets)
    pub app_queue_time: Option<Timestamp>,
}

/// Carries timestamps of events before the stream is returned to the application in accept().
#[non_exhaustive]
pub struct AcceptInfo {
    /// How long the stream spent inside the dcQUIC acceptor before being enqueued for the
    /// application.
    pub dc_quic_accept_time: Duration,
    /// How long the stream spent enqueued for the application.
    pub app_queue_sojourn_time: Duration,
}

impl<Sub> Builder<Sub>
where
    Sub: event::Subscriber,
{
    /// Builds the stream and emits an event indicating that the stream was built
    #[inline]
    #[allow(
        clippy::unwrap_in_result,
        reason = "app_queue_time is always set by accept_stream before accept() is called on a server stream"
    )]
    pub(crate) fn accept(self) -> io::Result<(Stream<Sub>, AcceptInfo)> {
        let kernel_start_time = self.kernel_start_time;
        let app_queue_time = self.app_queue_time.expect("set by accept_stream");
        let remote_address = self.shared.remote_addr();
        let remote_address = &remote_address;
        let creds = self.shared.credentials();
        let credential_id = &*creds.id;
        let stream_id = creds.key_id.as_u64();
        let now = self.shared.common.clock.get_time();
        let total_sojourn_time = now.saturating_duration_since(self.kernel_start_time);
        let queue_sojourn_time = now.saturating_duration_since(app_queue_time);

        self.shared
            .endpoint_publisher(now)
            .on_acceptor_stream_dequeued(event::builder::AcceptorStreamDequeued {
                remote_address,
                credential_id,
                stream_id,
                sojourn_time: total_sojourn_time,
                queue_sojourn_time,
            });

        self.build().map(|stream| {
            (
                stream,
                AcceptInfo {
                    dc_quic_accept_time: app_queue_time
                        .saturating_duration_since(kernel_start_time),
                    app_queue_sojourn_time: now.saturating_duration_since(app_queue_time),
                },
            )
        })
    }

    #[inline]
    pub(crate) fn connect(self) -> io::Result<Stream<Sub>> {
        self.build()
    }

    #[inline]
    pub(crate) fn build(self) -> io::Result<Stream<Sub>> {
        let Self {
            read,
            write,
            shared,
            sockets,
            kernel_start_time: _,
            app_queue_time: _,
        } = self;

        // TODO emit event

        let sockets = sockets.build()?;
        let read = read.build(shared.clone(), sockets.clone());
        let write = write.build(shared, sockets);
        Ok(Stream { read, write })
    }

    /// Emits an event indicating that the stream was pruned
    #[inline]
    pub(crate) fn prune(self, reason: event::builder::AcceptorStreamPruneReason) {
        let now = self.shared.clock.get_time();
        let remote_address = self.shared.remote_addr();
        let remote_address = &remote_address;
        let creds = self.shared.credentials();
        let credential_id = &*creds.id;
        let stream_id = creds.key_id.as_u64();
        let sojourn_time = now
            .saturating_duration_since(self.app_queue_time.expect("only called on server streams"));

        self.shared
            .endpoint_publisher(now)
            .on_acceptor_stream_pruned(event::builder::AcceptorStreamPruned {
                remote_address,
                credential_id,
                stream_id,
                sojourn_time,
                reason,
            });

        let publisher = self.shared.publisher_with_timestamp(now);

        self.shared.receiver.on_prune(&publisher);
        self.shared.sender.on_prune();
    }
}

pub struct Stream<Sub>
where
    Sub: event::Subscriber,
{
    read: Reader<Sub>,
    write: Writer<Sub>,
}

impl<Sub> fmt::Debug for Stream<Sub>
where
    Sub: event::Subscriber,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut s = f.debug_struct("Stream");

        for (name, addr) in [
            ("peer_addr", self.peer_addr()),
            ("local_addr", self.local_addr()),
        ] {
            if let Ok(addr) = addr {
                s.field(name, &addr);
            }
        }

        s.finish()
    }
}

impl<Sub> Stream<Sub>
where
    Sub: event::Subscriber,
{
    #[inline]
    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        self.read.peer_addr()
    }

    #[inline]
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.read.local_addr()
    }

    #[inline]
    pub fn path_secret_id(&self) -> &Id {
        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
    /// exposed for dcQUIC streams (at the cost of memory usage).
    #[inline]
    pub fn peer_cert_chain(&self) -> Option<&crate::stream::tls::CertificateChain> {
        self.read.peer_cert_chain()
    }

    #[inline]
    pub fn protocol(&self) -> socket::Protocol {
        self.read.protocol()
    }

    #[inline]
    pub fn set_read_mode(&mut self, read_mode: recv::ReadMode) -> &mut Self {
        self.read.set_read_mode(read_mode);
        self
    }

    #[inline]
    pub fn set_ack_mode(&mut self, ack_mode: recv::AckMode) -> &mut Self {
        self.read.set_ack_mode(ack_mode);
        self
    }

    #[inline]
    pub async fn write_from(
        &mut self,
        buf: &mut impl buffer::reader::storage::Infallible,
    ) -> io::Result<usize> {
        self.write.write_from(buf).await
    }

    #[inline]
    pub async fn write_all_from(
        &mut self,
        buf: &mut impl buffer::reader::storage::Infallible,
    ) -> io::Result<usize> {
        self.write.write_all_from(buf).await
    }

    #[inline]
    pub async fn write_from_fin(
        &mut self,
        buf: &mut impl buffer::reader::storage::Infallible,
    ) -> io::Result<usize> {
        self.write.write_from_fin(buf).await
    }

    #[inline]
    pub async fn write_all_from_fin(
        &mut self,
        buf: &mut impl buffer::reader::storage::Infallible,
    ) -> io::Result<usize> {
        self.write.write_all_from_fin(buf).await
    }

    #[inline]
    pub async fn read_into(
        &mut self,
        out_buf: &mut impl buffer::writer::Storage,
    ) -> io::Result<usize> {
        self.read.read_into(out_buf).await
    }

    #[inline]
    pub fn split(&mut self) -> (&mut Reader<Sub>, &mut Writer<Sub>) {
        (&mut self.read, &mut self.write)
    }

    #[inline]
    pub fn into_split(self) -> (Reader<Sub>, Writer<Sub>) {
        (self.read, self.write)
    }

    pub fn query_event_context<C: 'static, R>(&self, query: impl FnOnce(&C) -> R) -> Option<R> {
        self.read.query_event_context(query)
    }
}

#[cfg(feature = "tokio")]
mod tokio_impl {
    use super::{event, Stream};
    use core::{
        pin::Pin,
        task::{Context, Poll},
    };
    use tokio::io::{self, AsyncRead, AsyncWrite, ReadBuf};

    impl<Sub> AsyncRead for Stream<Sub>
    where
        Sub: event::Subscriber,
    {
        #[inline]
        fn poll_read(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &mut ReadBuf<'_>,
        ) -> Poll<io::Result<()>> {
            Pin::new(&mut self.read).poll_read(cx, buf)
        }
    }

    impl<Sub> AsyncWrite for Stream<Sub>
    where
        Sub: event::Subscriber,
    {
        #[inline]
        fn poll_write(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<io::Result<usize>> {
            Pin::new(&mut self.write).poll_write(cx, buf)
        }

        #[inline]
        fn poll_write_vectored(
            mut self: Pin<&mut Self>,
            cx: &mut Context,
            buf: &[std::io::IoSlice],
        ) -> Poll<io::Result<usize>> {
            Pin::new(&mut self.write).poll_write_vectored(cx, buf)
        }

        #[inline]
        fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
            Pin::new(&mut self.write).poll_flush(cx)
        }

        #[inline]
        fn poll_shutdown(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<Result<(), io::Error>> {
            Pin::new(&mut self.write).poll_shutdown(cx)
        }

        #[inline(always)]
        fn is_write_vectored(&self) -> bool {
            self.write.is_write_vectored()
        }
    }
}