apollo-router 2.13.1

use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;

use derivative::Derivative;
use futures::channel::oneshot;
use futures::prelude::*;
use itertools::Itertools;
use multimap::MultiMap;
use tokio::sync::mpsc;

use super::router::ApolloRouterError;
use crate::configuration::Configuration;
use crate::configuration::ListenAddr;
use crate::router_factory::Endpoint;
use crate::router_factory::RouterFactory;
use crate::uplink::license_enforcement::LicenseState;

/// Factory for creating the http server component.
///
/// This trait enables us to test that `StateMachine` correctly recreates the http server when
/// necessary e.g. when listen address changes.
pub(crate) trait HttpServerFactory {
    type Future: Future<Output = Result<HttpServerHandle, ApolloRouterError>> + Send;

    #[allow(clippy::too_many_arguments)]
    fn create<RF>(
        &self,
        service_factory: RF,
        configuration: Arc<Configuration>,
        main_listener: Option<Listener>,
        previous_listeners: ExtraListeners,
        extra_endpoints: MultiMap<ListenAddr, Endpoint>,
        license: Arc<LicenseState>,
        all_connections_stopped_sender: mpsc::Sender<()>,
    ) -> Self::Future
    where
        RF: RouterFactory;
}

type ExtraListeners = Vec<(ListenAddr, Listener)>;

/// A handle with with a client can shut down the server gracefully.
/// This relies on the underlying server implementation doing the right thing.
/// There are various ways that a user could prevent this working, including holding open connections
/// and sending huge requests. There is potential work needed for hardening.
#[derive(Derivative)]
#[derivative(Debug)]
pub(crate) struct HttpServerHandle {
    /// Sender to use to notify of shutdown
    main_shutdown_sender: oneshot::Sender<()>,

    /// Sender to use to notify extras of shutdown
    extra_shutdown_sender: oneshot::Sender<()>,

    /// Future to wait on for graceful shutdown
    #[derivative(Debug = "ignore")]
    main_future: Pin<Box<dyn Future<Output = Result<Listener, ApolloRouterError>> + Send>>,

    /// More futures to wait on for graceful shutdown
    #[derivative(Debug = "ignore")]
    extra_futures: Pin<Box<dyn Future<Output = Result<ExtraListeners, ApolloRouterError>> + Send>>,

    /// The listen addresses that the server is actually listening on.
    /// This includes the `graphql_listen_address` as well as any other address a plugin listens on.
    /// If a socket address specified port zero the OS will assign a random free port.
    listen_addresses: Vec<ListenAddr>,

    /// The listen addresses that the graphql server is actually listening on.
    /// If a socket address specified port zero the OS will assign a random free port.
    graphql_listen_address: Option<ListenAddr>,

    /// copied into every client session, to track if there are still running sessions when shutting down
    all_connections_stopped_sender: mpsc::Sender<()>,
}

impl HttpServerHandle {
    pub(crate) fn new(
        main_shutdown_sender: oneshot::Sender<()>,
        extra_shutdown_sender: oneshot::Sender<()>,
        main_future: Pin<
            Box<dyn Future<Output = Result<Listener, ApolloRouterError>> + Send + 'static>,
        >,
        extra_futures: Pin<
            Box<dyn Future<Output = Result<ExtraListeners, ApolloRouterError>> + Send + 'static>,
        >,
        graphql_listen_address: Option<ListenAddr>,
        listen_addresses: Vec<ListenAddr>,
        all_connections_stopped_sender: mpsc::Sender<()>,
    ) -> Self {
        Self {
            main_shutdown_sender,
            extra_shutdown_sender,
            main_future,
            extra_futures,
            graphql_listen_address,
            listen_addresses,
            all_connections_stopped_sender,
        }
    }

    #[cfg(unix)]
    pub(crate) async fn shutdown(mut self) -> Result<(), ApolloRouterError> {
        let listen_addresses = std::mem::take(&mut self.listen_addresses);

        let (_main_listener, _extra_listener) = self.wait_for_servers().await?;

        // listen_addresses includes the main graphql_address
        for listen_address in listen_addresses {
            if let ListenAddr::UnixSocket(path) = listen_address {
                let _ = tokio::fs::remove_file(path).await;
            }
        }
        Ok(())
    }

    #[cfg(not(unix))]
    pub(crate) async fn shutdown(self) -> Result<(), ApolloRouterError> {
        let (_main_listener, _extra_listener) = self.wait_for_servers().await?;

        Ok(())
    }

    pub(crate) async fn restart<RF, SF>(
        self,
        factory: &SF,
        router: RF,
        configuration: Arc<Configuration>,
        web_endpoints: MultiMap<ListenAddr, Endpoint>,
        license: Arc<LicenseState>,
    ) -> Result<Self, ApolloRouterError>
    where
        SF: HttpServerFactory,
        RF: RouterFactory,
    {
        let all_connections_stopped_sender = self.all_connections_stopped_sender.clone();

        // when the server receives the shutdown signal, it stops accepting new
        // connections, and returns the TCP listener, to reuse it in the next server
        // it is necessary to keep the queue of new TCP sockets associated with
        // the listeners instead of dropping them
        let (main_listener, extra_listeners) = self.wait_for_servers().await?;

        tracing::debug!("previous server stopped");

        // we give the listeners to the new configuration, they'll clean up whatever needs to
        let handle = factory
            .create(
                router,
                configuration,
                Some(main_listener),
                extra_listeners,
                web_endpoints,
                license,
                all_connections_stopped_sender,
            )
            .await?;
        tracing::debug!(
            "restarted on {}",
            handle
                .listen_addresses()
                .iter()
                .map(std::string::ToString::to_string)
                .join(" - ")
        );

        Ok(handle)
    }

    pub(crate) fn listen_addresses(&self) -> &[ListenAddr] {
        self.listen_addresses.as_slice()
    }

    pub(crate) fn graphql_listen_address(&self) -> &Option<ListenAddr> {
        &self.graphql_listen_address
    }

    async fn wait_for_servers(self) -> Result<(Listener, ExtraListeners), ApolloRouterError> {
        if let Err(_err) = self.main_shutdown_sender.send(()) {
            tracing::error!("Failed to notify http thread of shutdown")
        };
        let main_listener = self.main_future.await?;

        if let Err(_err) = self.extra_shutdown_sender.send(()) {
            tracing::error!("Failed to notify http thread of shutdown")
        };
        let extra_listeners = self.extra_futures.await?;
        Ok((main_listener, extra_listeners))
    }
}

pub(crate) enum Listener {
    Tcp(tokio::net::TcpListener),
    #[cfg(unix)]
    Unix(tokio::net::UnixListener),
    Tls {
        listener: tokio::net::TcpListener,
        acceptor: tokio_rustls::TlsAcceptor,
    },
}

// Though there is a large difference in variant sizes, only a few instances of this type will
// exist ever, so it's not a big deal.
#[allow(clippy::large_enum_variant)]
pub(crate) enum NetworkStream {
    Tcp(tokio::net::TcpStream),
    #[cfg(unix)]
    Unix(tokio::net::UnixStream),
    /// A TCP stream that needs TLS handshake. This variant is used to prevent blocking
    /// the accept loop during TLS handshake, which could be exploited for DoS attacks.
    Tls {
        stream: tokio::net::TcpStream,
        acceptor: tokio_rustls::TlsAcceptor,
    },
}

impl Listener {
    pub(crate) async fn new_from_socket_addr(
        address: SocketAddr,
        tls_acceptor: Option<tokio_rustls::TlsAcceptor>,
    ) -> Result<Self, ApolloRouterError> {
        let listener = tokio::net::TcpListener::bind(address)
            .await
            .map_err(ApolloRouterError::ServerCreationError)?;
        match tls_acceptor {
            None => Ok(Listener::Tcp(listener)),
            Some(acceptor) => Ok(Listener::Tls { listener, acceptor }),
        }
    }

    pub(crate) fn new_from_listener(
        listener: tokio::net::TcpListener,
        tls_acceptor: Option<tokio_rustls::TlsAcceptor>,
    ) -> Self {
        match tls_acceptor {
            None => Listener::Tcp(listener),
            Some(acceptor) => Listener::Tls { listener, acceptor },
        }
    }

    pub(crate) fn local_addr(&self) -> std::io::Result<ListenAddr> {
        match self {
            Listener::Tcp(listener) => listener.local_addr().map(Into::into),
            #[cfg(unix)]
            Listener::Unix(listener) => listener.local_addr().map(|addr| {
                ListenAddr::UnixSocket(
                    addr.as_pathname()
                        .map(ToOwned::to_owned)
                        .unwrap_or_default(),
                )
            }),
            Listener::Tls { listener, .. } => listener.local_addr().map(Into::into),
        }
    }

    pub(crate) async fn accept(&mut self) -> std::io::Result<NetworkStream> {
        match self {
            Listener::Tcp(listener) => listener
                .accept()
                .await
                .map(|(stream, _)| NetworkStream::Tcp(stream)),
            #[cfg(unix)]
            Listener::Unix(listener) => listener
                .accept()
                .await
                .map(|(stream, _)| NetworkStream::Unix(stream)),
            Listener::Tls { listener, acceptor } => {
                let (stream, _) = listener.accept().await?;

                // Return the TCP stream immediately with the acceptor, so the TLS handshake
                // can be performed in a spawned task with a timeout. This prevents blocking
                // the accept loop, which could be exploited for DoS attacks.
                Ok(NetworkStream::Tls {
                    stream,
                    acceptor: acceptor.clone(),
                })
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::net::SocketAddr;
    use std::str::FromStr;
    use std::time::Duration;

    use futures::channel::oneshot;
    use rustls::ServerConfig;
    use test_log::test;
    use tokio::net::TcpStream;
    use tokio::time::timeout;
    use tokio_rustls::TlsAcceptor;

    use super::*;
    use crate::configuration::load_certs;
    use crate::configuration::load_key;

    #[test(tokio::test)]
    // TODO [igni]: add a check with extra endpoints
    async fn sanity() {
        let (shutdown_sender, shutdown_receiver) = oneshot::channel();
        let (extra_shutdown_sender, extra_shutdown_receiver) = oneshot::channel();
        let listener = Listener::Tcp(tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap());
        let (all_connections_stopped_sender, _) = mpsc::channel::<()>(1);

        HttpServerHandle::new(
            shutdown_sender,
            extra_shutdown_sender,
            futures::future::ready(Ok(listener)).boxed(),
            futures::future::ready(Ok(vec![])).boxed(),
            Some(SocketAddr::from_str("127.0.0.1:0").unwrap().into()),
            Default::default(),
            all_connections_stopped_sender,
        )
        .shutdown()
        .await
        .expect("Should have waited for shutdown");

        shutdown_receiver
            .await
            .expect("Should have been send notification to shutdown");
        extra_shutdown_receiver
            .await
            .expect("Should have been send notification to shutdown");
    }

    #[test(tokio::test)]
    #[cfg(unix)]
    // TODO [igni]: add a check with extra endpoints
    async fn sanity_unix() {
        let temp_dir = tempfile::tempdir().unwrap();
        let sock = temp_dir.as_ref().join("sock");
        let (shutdown_sender, shutdown_receiver) = oneshot::channel();
        let (extra_shutdown_sender, extra_shutdown_receiver) = oneshot::channel();
        let listener = Listener::Unix(tokio::net::UnixListener::bind(&sock).unwrap());
        let (all_connections_stopped_sender, _) = mpsc::channel::<()>(1);

        HttpServerHandle::new(
            shutdown_sender,
            extra_shutdown_sender,
            futures::future::ready(Ok(listener)).boxed(),
            futures::future::ready(Ok(vec![])).boxed(),
            Some(ListenAddr::UnixSocket(sock)),
            Default::default(),
            all_connections_stopped_sender,
        )
        .shutdown()
        .await
        .expect("Should have waited for shutdown");

        shutdown_receiver
            .await
            .expect("Should have sent notification to shutdown");

        extra_shutdown_receiver
            .await
            .expect("Should have sent notification to shutdown");
    }

    #[tokio::test]
    async fn tls_handshake_timeout_does_not_block_accept_loop() {
        // Load test certificates
        let certificate_pem = include_str!("services/http/testdata/server_self_signed.crt");
        let key_pem = include_str!("services/http/testdata/server.key");

        let certificates = load_certs(certificate_pem).unwrap();
        let key = load_key(key_pem).unwrap();

        // Create TLS listener
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let socket_addr = tcp_listener.local_addr().unwrap();

        let mut tls_config = ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(certificates, key)
            .expect("built our tls config");
        tls_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
        let tls_config = std::sync::Arc::new(tls_config);
        let acceptor = TlsAcceptor::from(tls_config);

        let mut listener = Listener::Tls {
            listener: tcp_listener,
            acceptor,
        };

        //  Connect with plain TCP (simulating `nc`) - should timeout and not block
        let _plain_tcp_stream = TcpStream::connect(socket_addr).await.unwrap();

        // Accept the connection - this should return immediately with NetworkStream::Tls
        let network_stream = listener.accept().await.unwrap();

        match network_stream {
            NetworkStream::Tls { stream, .. } => {
                // The stream should be accepted immediately, handshake happens later
                // Just verify we got a valid stream (can get local/peer addresses)
                assert!(stream.local_addr().is_ok());
                assert!(stream.peer_addr().is_ok());
            }
            _ => panic!("Expected NetworkStream::Tls variant"),
        }

        // Verify we can still accept new connections (accept loop not blocked)
        let accept_future = listener.accept();
        let accept_result = timeout(Duration::from_millis(100), accept_future).await;

        // The accept should timeout because no new connection was made,
        // but importantly it should NOT hang indefinitely
        assert!(
            accept_result.is_err(),
            "Accept should timeout when no connection is made"
        );
    }
}