libdd-common 6.0.0

// Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/
// SPDX-License-Identifier: Apache-2.0

use core::{
    pin::Pin,
    task::{Context, Poll},
};

use futures::{Future, FutureExt};
use hyper::rt::ReadBufCursor;
use pin_project::pin_project;

#[pin_project(project=ConnStreamProj)]
#[derive(Debug)]
// `BoxedConn` is internal, and the field anyway not pub.
#[allow(private_interfaces)]
pub enum ConnStream {
    Tcp {
        #[pin]
        transport: TokioIo<tokio::net::TcpStream>,
    },
    Tls {
        #[pin]
        transport: wrapped::BoxedConn,
    },
    #[cfg(unix)]
    Udp {
        #[pin]
        transport: TokioIo<tokio::net::UnixStream>,
    },
    #[cfg(windows)]
    NamedPipe {
        #[pin]
        transport: TokioIo<tokio::net::windows::named_pipe::NamedPipeClient>,
    },
}

/// Wrapped connection, usable from hyper directly, or e.g. a hyper proxy wrapper
mod wrapped {
    use super::*;

    pub trait EstablishedConn:
        hyper::rt::Read + hyper::rt::Write + connect::Connection + Send
    {
    }

    impl<T: hyper::rt::Read + hyper::rt::Write + connect::Connection + Send> EstablishedConn for T {}

    #[derive(Debug)]
    pub struct BoxedConn(Pin<Box<dyn EstablishedConn>>);

    impl core::fmt::Debug for dyn EstablishedConn {
        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
            f.write_str("EstablishedConn")
        }
    }

    impl BoxedConn {
        pub fn new<T: EstablishedConn + 'static>(inner: T) -> Self {
            Self(Box::pin(inner))
        }
    }

    impl hyper::rt::Read for BoxedConn {
        fn poll_read(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: ReadBufCursor<'_>,
        ) -> Poll<std::io::Result<()>> {
            self.get_mut().0.as_mut().poll_read(cx, buf)
        }
    }

    impl hyper::rt::Write for BoxedConn {
        fn poll_write(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<Result<usize, std::io::Error>> {
            self.get_mut().0.as_mut().poll_write(cx, buf)
        }

        fn poll_flush(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<Result<(), std::io::Error>> {
            self.get_mut().0.as_mut().poll_flush(cx)
        }

        fn poll_shutdown(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<Result<(), std::io::Error>> {
            self.get_mut().0.as_mut().poll_shutdown(cx)
        }
    }

    impl connect::Connection for BoxedConn {
        fn connected(&self) -> connect::Connected {
            self.0.connected()
        }
    }
}

pub type ConnStreamError = Box<dyn core::error::Error + Send + Sync>;

use hyper_util::client::legacy::connect::{self, HttpConnector};
use hyper_util::rt::TokioIo;
use tower_service::Service;

impl ConnStream {
    pub async fn from_uds_uri(uri: hyper::Uri) -> Result<ConnStream, ConnStreamError> {
        #[cfg(unix)]
        {
            let path = super::uds::socket_path_from_uri(&uri)?;
            Ok(ConnStream::Udp {
                transport: TokioIo::new(tokio::net::UnixStream::connect(path).await?),
            })
        }
        #[cfg(not(unix))]
        {
            let _ = uri;
            Err(super::errors::Error::UnixSocketUnsupported.into())
        }
    }

    pub async fn from_named_pipe_uri(uri: hyper::Uri) -> Result<ConnStream, ConnStreamError> {
        #[cfg(windows)]
        {
            let path = super::named_pipe::named_pipe_path_from_uri(&uri)?;
            Ok(ConnStream::NamedPipe {
                transport: TokioIo::new(
                    tokio::net::windows::named_pipe::ClientOptions::new()
                        .security_qos_flags(super::named_pipe::ANONYMOUS_IMPERSONATION_QOS)
                        .open(path)?,
                ),
            })
        }
        #[cfg(not(windows))]
        {
            let _ = uri;
            Err(super::errors::Error::WindowsNamedPipeUnsupported.into())
        }
    }

    pub fn from_http_connector_with_uri(
        c: &mut HttpConnector,
        uri: hyper::Uri,
    ) -> impl Future<Output = Result<ConnStream, ConnStreamError>> + 'static {
        c.call(uri).map(|r| match r {
            Ok(t) => Ok(ConnStream::Tcp { transport: t }),
            Err(e) => Err(e.into()),
        })
    }

    #[cfg(feature = "tls-core")]
    pub fn from_https_connector_with_uri<C>(
        c: &mut hyper_rustls::HttpsConnector<C>,
        uri: hyper::Uri,
        require_tls: bool,
    ) -> impl Future<Output = Result<ConnStream, ConnStreamError>> + 'static
    where
        C: Service<hyper::Uri> + 'static,
        C::Response:
            hyper::rt::Read + hyper::rt::Write + connect::Connection + Send + Unpin + 'static,
        C::Future: Send + 'static,
        C::Error: Into<Box<dyn core::error::Error + Send + Sync>>,
    {
        #[allow(clippy::unwrap_used)]
        let stream_fut = c.call(uri.to_string().parse().unwrap());
        async move {
            let stream = stream_fut.await?;
            if require_tls && matches!(stream, hyper_rustls::MaybeHttpsStream::Http(_)) {
                return Err(super::errors::Error::CannotEstablishTlsConnection.into());
            }
            Ok(ConnStream::Tls {
                transport: wrapped::BoxedConn::new(stream),
            })
        }
    }

    pub fn from_tunnel_with_uri<C>(
        c: &mut C,
        uri: hyper::Uri,
    ) -> impl Future<Output = Result<ConnStream, ConnStreamError>> + 'static
    where
        C: Service<hyper::Uri> + 'static,
        C::Response:
            hyper::rt::Read + hyper::rt::Write + connect::Connection + Send + Unpin + 'static,
        C::Future: Send + 'static,
        C::Error: Into<Box<dyn core::error::Error + Send + Sync>>,
    {
        let stream_fut = c.call(uri);
        async move {
            let stream = stream_fut.await.map_err(Into::into)?;
            Ok(ConnStream::Tls {
                transport: wrapped::BoxedConn::new(stream),
            })
        }
    }
}

impl hyper::rt::Read for ConnStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: ReadBufCursor<'_>,
    ) -> Poll<std::io::Result<()>> {
        match self.project() {
            ConnStreamProj::Tcp { transport } => transport.poll_read(cx, buf),
            ConnStreamProj::Tls { transport } => transport.poll_read(cx, buf),
            #[cfg(unix)]
            ConnStreamProj::Udp { transport } => transport.poll_read(cx, buf),
            #[cfg(windows)]
            ConnStreamProj::NamedPipe { transport } => transport.poll_read(cx, buf),
        }
    }
}

impl connect::Connection for ConnStream {
    fn connected(&self) -> connect::Connected {
        match self {
            Self::Tcp { transport } => transport.connected(),
            Self::Tls { transport } => transport.connected(),
            #[cfg(unix)]
            Self::Udp { transport: _ } => connect::Connected::new(),
            #[cfg(windows)]
            Self::NamedPipe { transport: _ } => connect::Connected::new(),
        }
    }
}

impl hyper::rt::Write for ConnStream {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, std::io::Error>> {
        match self.project() {
            ConnStreamProj::Tcp { transport } => transport.poll_write(cx, buf),
            ConnStreamProj::Tls { transport } => transport.poll_write(cx, buf),
            #[cfg(unix)]
            ConnStreamProj::Udp { transport } => transport.poll_write(cx, buf),
            #[cfg(windows)]
            ConnStreamProj::NamedPipe { transport } => transport.poll_write(cx, buf),
        }
    }

    fn poll_shutdown(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), std::io::Error>> {
        match self.project() {
            ConnStreamProj::Tcp { transport } => transport.poll_shutdown(cx),
            ConnStreamProj::Tls { transport } => transport.poll_shutdown(cx),
            #[cfg(unix)]
            ConnStreamProj::Udp { transport } => transport.poll_shutdown(cx),
            #[cfg(windows)]
            ConnStreamProj::NamedPipe { transport } => transport.poll_shutdown(cx),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
        match self.project() {
            ConnStreamProj::Tcp { transport } => transport.poll_flush(cx),
            ConnStreamProj::Tls { transport } => transport.poll_flush(cx),
            #[cfg(unix)]
            ConnStreamProj::Udp { transport } => transport.poll_flush(cx),
            #[cfg(windows)]
            ConnStreamProj::NamedPipe { transport } => transport.poll_flush(cx),
        }
    }
}

#[cfg(all(test, windows))]
mod windows_named_pipe_tests {
    use super::ConnStream;
    use crate::connector::named_pipe::named_pipe_path_to_uri;
    use std::path::Path;
    use tokio::net::windows::named_pipe::ServerOptions;

    /// Verifies that `from_named_pipe_uri` opens a client connection successfully
    /// when impersonation is disabled (Anonymous QoS). The server accepting the
    /// connection confirms the `SECURITY_SQOS_PRESENT | SECURITY_ANONYMOUS` flags
    /// produce a usable transport.
    #[tokio::test]
    async fn from_named_pipe_uri_connects_with_anonymous_qos() {
        let pipe_name = format!(
            r"\\.\pipe\libdd_common_conn_stream_test_{}_{}",
            std::process::id(),
            rand::random::<u64>()
        );

        let server = ServerOptions::new()
            .first_pipe_instance(true)
            .create(&pipe_name)
            .expect("failed to create named pipe server");

        let server_task = tokio::spawn(async move {
            server.connect().await.expect("server failed to accept");
        });

        let uri = named_pipe_path_to_uri(Path::new(&pipe_name)).expect("failed to build uri");
        let conn = ConnStream::from_named_pipe_uri(uri).await;
        assert!(
            conn.is_ok(),
            "expected named pipe client to connect with Anonymous QoS: {:?}",
            conn.err()
        );

        server_task.await.expect("server task panicked");
    }
}