dragoman 0.6.2

//! Hand TLS connections for other hostnames to other servers on this machine.
//!
//! dragoman owns port 443. With `--sni-route <host>=<address>`, a connection
//! whose TLS ClientHello names `<host>` (its SNI) is spliced, still encrypted,
//! to `<address>`, and the server there terminates TLS itself. Every other
//! connection — no SNI, a name with no route, or one that is not TLS at all —
//! is dragoman's, as without routes. The ClientHello is only peeked at, never
//! consumed, so whichever side takes the connection sees it whole.
//!
//! A request for a routed host can still reach dragoman's HTTP stack when a
//! client reuses a connection opened for another name; [`misdirected`] answers
//! it with 421, which tells the client to retry on a connection of its own.

use std::collections::HashMap;
use std::future::Future;
use std::io;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use axum::extract::{Request, State};
use axum::http::{StatusCode, header};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use tokio::net::TcpStream;

/// The largest ClientHello looked for: one TLS record of the largest size.
const MAX_HELLO: usize = 5 + 16 * 1024;

/// How long a client may take to send its whole ClientHello. dragoman's own
/// handshake deadline usually ends the wait sooner.
const HELLO_TIMEOUT: Duration = Duration::from_secs(10);

/// Hostname → backend address.
#[derive(Debug, Default, Clone)]
pub struct Routes(HashMap<String, SocketAddr>);

impl Routes {
    pub fn new(routes: impl IntoIterator<Item = (String, SocketAddr)>) -> Self {
        Routes(routes.into_iter().collect())
    }

    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// The backend for `host`, compared without case and a trailing dot.
    pub fn get(&self, host: &str) -> Option<SocketAddr> {
        self.0.get(&normalize(host)).copied()
    }

    pub fn hosts(&self) -> impl Iterator<Item = &str> {
        self.0.keys().map(String::as_str)
    }
}

fn normalize(host: &str) -> String {
    host.trim_end_matches('.').to_ascii_lowercase()
}

/// Parse a `--sni-route` value, `<host>=<address>` (e.g.
/// `defter.commonmeta.org=127.0.0.1:8443`).
pub fn parse_route(value: &str) -> Result<(String, SocketAddr), String> {
    let (host, addr) = value
        .split_once('=')
        .ok_or_else(|| format!("expected <host>=<address>, got '{value}'"))?;
    let host = normalize(host.trim());
    let valid = !host.is_empty()
        && host.split('.').all(|label| {
            !label.is_empty()
                && label
                    .bytes()
                    .all(|b| b.is_ascii_alphanumeric() || b == b'-')
        });
    if !valid {
        return Err(format!("'{host}' is not a hostname"));
    }
    let addr = addr
        .trim()
        .parse()
        .map_err(|e| format!("'{addr}' is not an address like 127.0.0.1:8443: {e}"))?;
    Ok((host, addr))
}

// ── ClientHello ──────────────────────────────────────────────────────────────

/// What the first bytes of a connection say about its server name.
#[derive(Debug, PartialEq, Eq)]
pub enum Hello {
    /// The first TLS record has not fully arrived yet.
    Incomplete,
    /// The ClientHello is all here: the name it asks for, if any. Also `None`
    /// for a connection that is not TLS or a ClientHello this does not read.
    Complete(Option<String>),
}

/// The SNI host name of the ClientHello at the start of `buf`.
pub fn server_name(buf: &[u8]) -> Hello {
    if buf.len() < 5 {
        // Not even a record header; but a first byte that cannot start a
        // handshake record is not TLS, whatever follows.
        return match buf.first() {
            Some(&first) if first != 0x16 => Hello::Complete(None),
            _ => Hello::Incomplete,
        };
    }
    if buf[0] != 0x16 || buf[1] != 0x03 {
        return Hello::Complete(None);
    }
    let length = u16::from_be_bytes([buf[3], buf[4]]) as usize;
    if 5 + length > MAX_HELLO {
        return Hello::Complete(None);
    }
    let Some(record) = buf.get(5..5 + length) else {
        return Hello::Incomplete;
    };
    Hello::Complete(parse_client_hello(record))
}

/// A cursor over TLS's length-prefixed fields; `None` past the end.
struct Reader<'a>(&'a [u8]);

impl<'a> Reader<'a> {
    fn take(&mut self, n: usize) -> Option<&'a [u8]> {
        if self.0.len() < n {
            return None;
        }
        let (head, rest) = self.0.split_at(n);
        self.0 = rest;
        Some(head)
    }

    fn u8(&mut self) -> Option<usize> {
        self.take(1).map(|b| b[0] as usize)
    }

    fn u16(&mut self) -> Option<usize> {
        self.take(2)
            .map(|b| u16::from_be_bytes([b[0], b[1]]) as usize)
    }

    fn u24(&mut self) -> Option<usize> {
        self.take(3)
            .map(|b| u32::from_be_bytes([0, b[0], b[1], b[2]]) as usize)
    }

    /// A field with a `len`-byte length prefix.
    fn vec(&mut self, len: fn(&mut Self) -> Option<usize>) -> Option<Reader<'a>> {
        let n = len(self)?;
        self.take(n).map(Reader)
    }
}

/// The server name in a handshake record holding a ClientHello (RFC 8446
/// §4.1.2, RFC 6066 §3). A ClientHello split over several records is not read:
/// it falls to dragoman, as one without SNI would.
fn parse_client_hello(record: &[u8]) -> Option<String> {
    let mut r = Reader(record);
    if r.u8()? != 0x01 {
        return None; // not a ClientHello
    }
    let mut hello = r.vec(Reader::u24)?;
    hello.take(2 + 32)?; // legacy_version, random
    hello.vec(Reader::u8)?; // legacy_session_id
    hello.vec(Reader::u16)?; // cipher_suites
    hello.vec(Reader::u8)?; // legacy_compression_methods
    let mut extensions = hello.vec(Reader::u16)?;
    while !extensions.0.is_empty() {
        let kind = extensions.u16()?;
        let mut data = extensions.vec(Reader::u16)?;
        if kind != 0x0000 {
            continue; // not server_name
        }
        let mut names = data.vec(Reader::u16)?;
        while !names.0.is_empty() {
            let name_type = names.u8()?;
            let name = names.vec(Reader::u16)?;
            if name_type == 0 {
                let name = std::str::from_utf8(name.0).ok()?;
                return Some(normalize(name));
            }
        }
        return None;
    }
    None
}

/// Peek at `stream` until its ClientHello has arrived, and return its name.
async fn peek_server_name(stream: &TcpStream) -> io::Result<Option<String>> {
    let mut buf = vec![0u8; MAX_HELLO];
    let mut seen = 0;
    loop {
        let n = stream.peek(&mut buf).await?;
        if n == 0 {
            return Ok(None);
        }
        match server_name(&buf[..n]) {
            Hello::Complete(name) => return Ok(name),
            Hello::Incomplete if n == buf.len() => return Ok(None),
            // `peek` returns at once while bytes wait unread, so give the rest
            // of a split ClientHello a moment to arrive.
            Hello::Incomplete if n == seen => tokio::time::sleep(Duration::from_millis(5)).await,
            Hello::Incomplete => seen = n,
        }
    }
}

// ── acceptor ─────────────────────────────────────────────────────────────────

/// Wraps dragoman's TLS acceptor: a connection for a routed host goes to its
/// backend, any other to `inner`.
#[derive(Clone)]
pub struct SniAcceptor<A> {
    inner: A,
    routes: Arc<Routes>,
}

impl<A> SniAcceptor<A> {
    pub fn new(inner: A, routes: Arc<Routes>) -> Self {
        Self { inner, routes }
    }
}

impl<A, S> axum_server::accept::Accept<TcpStream, S> for SniAcceptor<A>
where
    A: axum_server::accept::Accept<TcpStream, S> + Clone + Send + Sync + 'static,
    A::Future: Send + 'static,
    S: Send + 'static,
{
    type Stream = A::Stream;
    type Service = A::Service;
    type Future = Pin<Box<dyn Future<Output = io::Result<(A::Stream, A::Service)>> + Send>>;

    fn accept(&self, stream: TcpStream, service: S) -> Self::Future {
        if self.routes.is_empty() {
            return Box::pin(self.inner.accept(stream, service));
        }
        let inner = self.inner.clone();
        let routes = self.routes.clone();
        Box::pin(async move {
            let name = tokio::time::timeout(HELLO_TIMEOUT, peek_server_name(&stream))
                .await
                .map_err(|_| io::Error::other("no ClientHello in time"))??;
            let route = name.and_then(|name| routes.get(&name).map(|backend| (name, backend)));
            let Some((host, backend)) = route else {
                return inner.accept(stream, service).await;
            };
            // The splice lives as long as the connection, so it runs on its
            // own rather than under the handshake deadline this accept is
            // held to; ending the accept with an error tells the server the
            // connection is not its to serve.
            tokio::spawn(pass_through(stream, host, backend));
            Err(io::Error::other("passed through to another server"))
        })
    }
}

/// Splice `client` to `backend` until either side closes.
async fn pass_through(mut client: TcpStream, host: String, backend: SocketAddr) {
    metrics::counter!("sni_passed_through_total", "host" => host.clone()).increment(1);
    let mut upstream = match TcpStream::connect(backend).await {
        Ok(upstream) => upstream,
        Err(error) => {
            metrics::counter!("sni_backend_unreachable_total", "host" => host.clone()).increment(1);
            tracing::warn!(host, %backend, %error, "SNI route: backend unreachable");
            return;
        }
    };
    let _ = client.set_nodelay(true);
    let _ = upstream.set_nodelay(true);
    match tokio::io::copy_bidirectional(&mut client, &mut upstream).await {
        Ok((sent, received)) => {
            tracing::debug!(host, %backend, sent, received, "SNI route: closed")
        }
        Err(error) => tracing::debug!(host, %backend, %error, "SNI route: ended"),
    }
}

// ── misdirected requests ─────────────────────────────────────────────────────

/// Answer 421 Misdirected Request to a request for a routed host that reached
/// dragoman anyway — over a connection the client opened for another name and
/// reused, as HTTP/2 allows when one certificate covers both. The client then
/// opens a connection for that host, which the acceptor routes.
pub async fn misdirected(State(routes): State<Arc<Routes>>, req: Request, next: Next) -> Response {
    let host = req
        .headers()
        .get(header::HOST)
        .and_then(|v| v.to_str().ok())
        .or_else(|| req.uri().host())
        .map(|h| {
            h.rsplit_once(':').map_or(h, |(host, port)| {
                // Keep an IPv6 literal's colons; drop only a port.
                if port.bytes().all(|b| b.is_ascii_digit()) {
                    host
                } else {
                    h
                }
            })
        });
    if let Some(host) = host.filter(|host| routes.get(host).is_some()) {
        metrics::counter!("http_misdirected_total").increment(1);
        tracing::info!(host, path = %req.uri().path(), "misdirected request for a routed host");
        return (StatusCode::MISDIRECTED_REQUEST, "misdirected request").into_response();
    }
    next.run(req).await
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{Router, body::Body, routing::get};
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tower::ServiceExt;

    fn provider() -> Arc<rustls::crypto::CryptoProvider> {
        Arc::new(rustls::crypto::aws_lc_rs::default_provider())
    }

    fn client_config(roots: rustls::RootCertStore) -> Arc<rustls::ClientConfig> {
        Arc::new(
            rustls::ClientConfig::builder_with_provider(provider())
                .with_safe_default_protocol_versions()
                .unwrap()
                .with_root_certificates(roots)
                .with_no_client_auth(),
        )
    }

    /// The ClientHello rustls sends to open a connection to `name`.
    fn client_hello(name: &str) -> Vec<u8> {
        let server_name = rustls::pki_types::ServerName::try_from(name.to_string()).unwrap();
        let mut conn = rustls::ClientConnection::new(
            client_config(rustls::RootCertStore::empty()),
            server_name,
        )
        .unwrap();
        let mut out = Vec::new();
        conn.write_tls(&mut out).unwrap();
        out
    }

    #[test]
    fn reads_the_server_name_of_a_client_hello() {
        let hello = client_hello("Defter.CommonMeta.org");
        assert_eq!(
            server_name(&hello),
            Hello::Complete(Some("defter.commonmeta.org".into()))
        );
        // Every prefix short of the whole record waits for more.
        for cut in [0, 1, 4, 5, 60, hello.len() - 1] {
            assert_eq!(
                server_name(&hello[..cut]),
                Hello::Incomplete,
                "cut at {cut}"
            );
        }
    }

    #[test]
    fn no_name_for_ip_addresses_or_plain_text() {
        // rustls sends no SNI for an IP address.
        assert_eq!(
            server_name(&client_hello("127.0.0.1")),
            Hello::Complete(None)
        );
        assert_eq!(server_name(b"GET / HTTP/1.1\r\n"), Hello::Complete(None));
        assert_eq!(server_name(b"G"), Hello::Complete(None));
        // A handshake record that is not a ClientHello.
        assert_eq!(
            server_name(&[0x16, 0x03, 0x01, 0x00, 0x01, 0x02]),
            Hello::Complete(None)
        );
        // A record claiming to be larger than any ClientHello.
        assert_eq!(
            server_name(&[0x16, 0x03, 0x01, 0xff, 0xff]),
            Hello::Complete(None)
        );
    }

    #[test]
    fn parses_routes() {
        assert_eq!(
            parse_route("Defter.CommonMeta.org.=127.0.0.1:8443"),
            Ok((
                "defter.commonmeta.org".into(),
                "127.0.0.1:8443".parse().unwrap()
            ))
        );
        assert!(parse_route("defter.commonmeta.org").is_err());
        assert!(parse_route("=127.0.0.1:8443").is_err());
        assert!(parse_route("bad host=127.0.0.1:8443").is_err());
        assert!(parse_route("defter.commonmeta.org=localhost").is_err());
        let routes = Routes::new([parse_route("a.example=127.0.0.1:1").unwrap()]);
        assert_eq!(routes.get("A.Example."), "127.0.0.1:1".parse().ok());
        assert_eq!(routes.get("b.example"), None);
    }

    #[tokio::test]
    async fn requests_for_routed_hosts_are_misdirected() {
        let routes = Arc::new(Routes::new([
            parse_route("defter.example=127.0.0.1:1").unwrap()
        ]));
        let app = Router::new()
            .route("/", get(|| async { "dragoman" }))
            .layer(axum::middleware::from_fn_with_state(routes, misdirected));
        let status = |host: &'static str| {
            let app = app.clone();
            async move {
                let req = axum::http::Request::builder()
                    .uri("/")
                    .header(header::HOST, host)
                    .body(Body::empty())
                    .unwrap();
                app.oneshot(req).await.unwrap().status()
            }
        };
        assert_eq!(
            status("defter.example").await,
            StatusCode::MISDIRECTED_REQUEST
        );
        assert_eq!(
            status("DEFTER.example:443").await,
            StatusCode::MISDIRECTED_REQUEST
        );
        assert_eq!(status("commonmeta.example").await, StatusCode::OK);
    }

    /// A routed name reaches the backend byte for byte; any other name, and a
    /// connection without SNI, is served by dragoman's TLS stack.
    #[tokio::test]
    async fn routes_by_server_name_and_falls_back_to_dragoman() {
        // As main does: rustls is built with more than one provider here.
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
        // The backend: records what it receives and answers with a marker.
        let backend = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let backend_addr = backend.local_addr().unwrap();
        let (seen_tx, mut seen_rx) = tokio::sync::mpsc::unbounded_channel();
        tokio::spawn(async move {
            loop {
                let (mut conn, _) = backend.accept().await.unwrap();
                let seen = seen_tx.clone();
                tokio::spawn(async move {
                    let mut buf = vec![0u8; MAX_HELLO];
                    let n = conn.read(&mut buf).await.unwrap();
                    seen.send(buf[..n].to_vec()).unwrap();
                    conn.write_all(b"backend").await.unwrap();
                });
            }
        });

        // dragoman: TLS for commonmeta.test, with one route.
        let cert = rcgen::generate_simple_self_signed(vec!["commonmeta.test".into()]).unwrap();
        let config = axum_server::tls_rustls::RustlsConfig::from_der(
            vec![cert.cert.der().to_vec()],
            cert.signing_key.serialize_der(),
        )
        .await
        .unwrap();
        let routes = Arc::new(Routes::new([(String::from("backend.test"), backend_addr)]));
        let acceptor =
            SniAcceptor::new(axum_server::tls_rustls::RustlsAcceptor::new(config), routes);
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        listener.set_nonblocking(true).unwrap();
        let addr = listener.local_addr().unwrap();
        let app = Router::new().route("/", get(|| async { "dragoman" }));
        tokio::spawn(async move {
            axum_server::from_tcp(listener)
                .unwrap()
                .acceptor(acceptor)
                .serve(app.into_make_service())
                .await
                .unwrap();
        });

        // A routed name: the backend gets the ClientHello unchanged.
        let hello = client_hello("backend.test");
        let mut conn = TcpStream::connect(addr).await.unwrap();
        conn.write_all(&hello).await.unwrap();
        let mut answer = Vec::new();
        conn.read_to_end(&mut answer).await.unwrap();
        assert_eq!(answer, b"backend");
        assert_eq!(seen_rx.recv().await.unwrap(), hello);

        // dragoman's own name: a TLS session with dragoman.
        let mut roots = rustls::RootCertStore::empty();
        roots.add(cert.cert.der().clone()).unwrap();
        let connector = tokio_rustls::TlsConnector::from(client_config(roots));
        let tcp = TcpStream::connect(addr).await.unwrap();
        let name = rustls::pki_types::ServerName::try_from("commonmeta.test").unwrap();
        let mut tls = connector.connect(name, tcp).await.unwrap();
        tls.write_all(b"GET / HTTP/1.1\r\nHost: commonmeta.test\r\nConnection: close\r\n\r\n")
            .await
            .unwrap();
        let mut response = Vec::new();
        tls.read_to_end(&mut response).await.ok();
        let response = String::from_utf8_lossy(&response);
        assert!(response.starts_with("HTTP/1.1 200"), "{response}");
        assert!(response.ends_with("dragoman"), "{response}");
        assert!(
            seen_rx.try_recv().is_err(),
            "the backend saw no other connection"
        );
    }
}