rustls_acme/
axum.rs

1use crate::futures_rustls::rustls::ServerConfig;
2use crate::{AcmeAccept, AcmeAcceptor};
3use futures::prelude::*;
4use futures_rustls::Accept;
5use std::io;
6use std::pin::Pin;
7use std::sync::Arc;
8use std::task::{Context, Poll};
9use tokio::io::{AsyncRead, AsyncWrite};
10use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
11
12#[derive(Clone)]
13pub struct AxumAcceptor {
14    acme_acceptor: AcmeAcceptor,
15    config: Arc<ServerConfig>,
16}
17
18impl AxumAcceptor {
19    pub fn new(acme_acceptor: AcmeAcceptor, config: Arc<ServerConfig>) -> Self {
20        Self { acme_acceptor, config }
21    }
22}
23
24impl<I: AsyncRead + AsyncWrite + Unpin + Send + 'static, S: Send + 'static> axum_server::accept::Accept<I, S> for AxumAcceptor {
25    type Stream = Compat<futures_rustls::server::TlsStream<Compat<I>>>;
26    type Service = S;
27    type Future = AxumAccept<I, S>;
28
29    fn accept(&self, stream: I, service: S) -> Self::Future {
30        let acme_accept = self.acme_acceptor.accept(stream.compat());
31        Self::Future {
32            config: self.config.clone(),
33            acme_accept,
34            tls_accept: None,
35            service: Some(service),
36        }
37    }
38}
39
40pub struct AxumAccept<I: AsyncRead + AsyncWrite + Unpin + Send + 'static, S: Send + 'static> {
41    config: Arc<ServerConfig>,
42    acme_accept: AcmeAccept<Compat<I>>,
43    tls_accept: Option<Accept<Compat<I>>>,
44    service: Option<S>,
45}
46
47impl<I: AsyncRead + AsyncWrite + Unpin + Send + 'static, S: Send + 'static> Unpin for AxumAccept<I, S> {}
48
49impl<I: AsyncRead + AsyncWrite + Unpin + Send + 'static, S: Send + 'static> Future for AxumAccept<I, S> {
50    type Output = io::Result<(Compat<futures_rustls::server::TlsStream<Compat<I>>>, S)>;
51
52    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
53        loop {
54            if let Some(tls_accept) = &mut self.tls_accept {
55                return match Pin::new(&mut *tls_accept).poll(cx) {
56                    Poll::Ready(Ok(tls)) => Poll::Ready(Ok((tls.compat(), self.service.take().unwrap()))),
57                    Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
58                    Poll::Pending => Poll::Pending,
59                };
60            }
61            return match Pin::new(&mut self.acme_accept).poll(cx) {
62                Poll::Ready(Ok(Some(start_handshake))) => {
63                    let config = self.config.clone();
64                    self.tls_accept = Some(start_handshake.into_stream(config));
65                    continue;
66                }
67                Poll::Ready(Ok(None)) => Poll::Ready(Err(io::Error::other("TLS-ALPN-01 validation request"))),
68                Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
69                Poll::Pending => Poll::Pending,
70            };
71        }
72    }
73}