async_sni/
async_sni.rs

1//! Run with `cargo run --all-features --example async_sni` command.
2//!
3//! Resolves the certificate from the server name in the ClientHello, before
4//! the handshake completes. The built-in `RustlsAcceptor` picks one
5//! `ServerConfig` up front; this example defers that choice by implementing
6//! `Accept` with `tokio_rustls::LazyConfigAcceptor`.
7//!
8//! Try it with:
9//!
10//! ```text
11//! curl -k --resolve localhost:3000:127.0.0.1 https://localhost:3000/
12//! curl -k --resolve other.test:3000:127.0.0.1 https://other.test:3000/
13//! ```
14//!
15//! The server prints which certificate it resolved for each connection.
16
17use axum::{routing::get, Router};
18use axum_server::{accept::Accept, tls_rustls::RustlsConfig};
19use rustls::server::Acceptor;
20use std::{
21    collections::HashMap, future::Future, io, net::SocketAddr, pin::Pin, sync::Arc, time::Duration,
22};
23use tokio::{
24    io::{AsyncRead, AsyncWrite},
25    time::{timeout_at, Instant},
26};
27use tokio_rustls::{server::TlsStream, LazyConfigAcceptor};
28
29#[tokio::main]
30async fn main() {
31    let app = Router::new().route("/", get(|| async { "Hello, world!" }));
32
33    let acceptor = AsyncSniAcceptor::new(CertificateStore::example());
34    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
35
36    println!("listening on {addr}");
37    println!("try: curl -k --resolve localhost:3000:127.0.0.1 https://localhost:3000/");
38    println!("try: curl -k --resolve other.test:3000:127.0.0.1 https://other.test:3000/");
39
40    axum_server::bind(addr)
41        .acceptor(acceptor)
42        .serve(app.into_make_service())
43        .await
44        .unwrap();
45}
46
47/// An acceptor that resolves the certificate after reading the ClientHello.
48#[derive(Clone)]
49struct AsyncSniAcceptor {
50    store: CertificateStore,
51    handshake_timeout: Duration,
52}
53
54impl AsyncSniAcceptor {
55    fn new(store: CertificateStore) -> Self {
56        Self {
57            store,
58            handshake_timeout: Duration::from_secs(10),
59        }
60    }
61}
62
63impl<I, S> Accept<I, S> for AsyncSniAcceptor
64where
65    I: AsyncRead + AsyncWrite + Unpin + Send + 'static,
66    S: Send + 'static,
67{
68    type Stream = TlsStream<I>;
69    type Service = S;
70    type Future = Pin<Box<dyn Future<Output = io::Result<(Self::Stream, Self::Service)>> + Send>>;
71
72    fn accept(&self, stream: I, service: S) -> Self::Future {
73        let store = self.store.clone();
74        let handshake_timeout = self.handshake_timeout;
75
76        Box::pin(async move {
77            // One deadline covers reading the ClientHello and the handshake.
78            let deadline = Instant::now() + handshake_timeout;
79
80            let start = timeout_at(
81                deadline,
82                LazyConfigAcceptor::new(Acceptor::default(), stream),
83            )
84            .await
85            .map_err(|_| timeout_error())??;
86
87            let server_name = start.client_hello().server_name().map(str::to_owned);
88            let config = store.resolve(server_name.as_deref()).await?;
89
90            let stream = timeout_at(deadline, start.into_stream(config.get_inner()))
91                .await
92                .map_err(|_| timeout_error())??;
93
94            Ok((stream, service))
95        })
96    }
97}
98
99fn timeout_error() -> io::Error {
100    io::Error::new(io::ErrorKind::TimedOut, "tls handshake timed out")
101}
102
103/// Resolves certificates asynchronously. A real store could query a database,
104/// a secret manager, or an ACME directory here.
105#[derive(Clone)]
106struct CertificateStore {
107    certificates: Arc<HashMap<&'static str, (&'static str, &'static str)>>,
108}
109
110impl CertificateStore {
111    fn example() -> Self {
112        Self {
113            certificates: Arc::new(HashMap::from([
114                (
115                    "localhost",
116                    (
117                        "examples/self-signed-certs/sni/localhost-cert.pem",
118                        "examples/self-signed-certs/sni/localhost-key.pem",
119                    ),
120                ),
121                (
122                    "other.test",
123                    (
124                        "examples/self-signed-certs/sni/other-test-cert.pem",
125                        "examples/self-signed-certs/sni/other-test-key.pem",
126                    ),
127                ),
128            ])),
129        }
130    }
131
132    async fn resolve(&self, server_name: Option<&str>) -> io::Result<RustlsConfig> {
133        let server_name = server_name.ok_or_else(|| {
134            io::Error::new(
135                io::ErrorKind::InvalidData,
136                "client did not send a server name",
137            )
138        })?;
139
140        let (cert, key) = self.certificates.get(server_name).ok_or_else(|| {
141            io::Error::new(
142                io::ErrorKind::NotFound,
143                format!("no certificate registered for {server_name}"),
144            )
145        })?;
146
147        println!("resolving certificate for {server_name}");
148
149        // Async I/O happens here, before the handshake continues. `from_pem_file`
150        // also sets the `h2` and `http/1.1` ALPN protocols. A `ServerConfig`
151        // built by hand must set them too.
152        RustlsConfig::from_pem_file(cert, key).await
153    }
154}
async_sni/async_sni.rs

async_sni/
async_sni.rs
