1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

use hyper::server::conn::http1;
use hyper_util::rt::TokioIo;
use std::net::SocketAddr;
use thiserror::Error;

/// Choose to accept either a HTTP or HTTPS connection
///
/// Created by calling the `build` method on an `AcceptorBuilder`
// Use a struct instead of the enum directly to avoid users constructing/matching on enum variants
pub struct HttpOrHttpsAcceptor(pub(crate) AcceptorInner);

pub enum AcceptorInner {
    Http(tokio::net::TcpListener),
    Https(tls_listener::TlsListener<tokio::net::TcpListener, tokio_rustls::TlsAcceptor>),
}

impl HttpOrHttpsAcceptor {
    /// Accepts every connection using the service provided, never completes.
    /// Ignores any connection errors produced by the `accept` method.
    pub async fn serve<S>(&mut self, service: S)
    where
        S: hyper::service::HttpService<hyper::body::Incoming> + Clone + Send + 'static,
        S::Future: Send,
        S::ResBody: Send + 'static,
        <S::ResBody as hyper::body::Body>::Error: std::error::Error + Send + Sync + 'static,
        <S::ResBody as hyper::body::Body>::Data: Send,
    {
        loop {
            // Ignore result here
            let _ = self.accept(service.clone()).await;
        }
    }

    /// Accepts a singular connection and spawns it onto the tokio runtime.
    /// Returns the address of the connected client.
    ///
    /// # Errors
    /// If the TCP connection or TLS handshake (HTTPS only) fails
    // Function won't panic, however tokio worker might
    #[allow(clippy::missing_panics_doc)]
    pub async fn accept<S>(&mut self, service: S) -> Result<SocketAddr, AcceptorError>
    where
        S: hyper::service::HttpService<hyper::body::Incoming> + Send + 'static,
        S::Future: Send,
        S::ResBody: Send + 'static,
        <S::ResBody as hyper::body::Body>::Error: std::error::Error + Send + Sync + 'static,
        <S::ResBody as hyper::body::Body>::Data: Send,
    {
        let conn_builder = http1::Builder::new();

        match &mut self.0 {
            AcceptorInner::Http(listener) => {
                let (conn, peer_addr) =
                    listener.accept().await.map_err(AcceptorError::TcpConnect)?;

                let conn = TokioIo::new(conn);

                // Allow upgrades so websockets can be used in client code
                let conn = conn_builder.serve_connection(conn, service).with_upgrades();

                // Ignore result here, it will only show up if client terminates connection without sending a request
                tokio::spawn(conn);

                Ok(peer_addr)
            }
            AcceptorInner::Https(listener) => {
                let (conn, peer_addr) = loop {
                    match listener.accept().await {
                        Err(tls_listener::Error::ListenerError(e)) => {
                            return Err(AcceptorError::TcpConnect(e))
                        }
                        Err(tls_listener::Error::TlsAcceptError { error, .. }) => {
                            return Err(AcceptorError::TcpConnect(error))
                        }
                        // Ignore handshake timeout errors, just try to get another connection
                        Err(_) => continue,
                        Ok(conn_and_addr) => break conn_and_addr,
                    }
                };

                let conn = TokioIo::new(conn);

                // Allow upgrades so websockets can be used in client code
                let conn = conn_builder.serve_connection(conn, service).with_upgrades();

                // Ignore result here, it will only show up if client terminates connection without sending a request
                tokio::spawn(conn);

                Ok(peer_addr)
            }
        }
    }
}

/// Error when accepting connections
#[derive(Error, Debug)]
pub enum AcceptorError {
    /// Failed to connect to client over TCP
    #[error("TCP connection to client failed")]
    TcpConnect(#[source] std::io::Error),
    /// Failed to make TLS handshake with client
    #[error("TLS handshake with client failed")]
    TlsHandshake(#[source] std::io::Error),
}