packages/ak-axum/accept/tls: init (goauthentik#21318)

Author: rissson

Cargo.lock

@@ -33,6 +33,7 @@ console-subscriber = "= 0.5.0"
 dotenvy = "= 0.15.7"
 durstr = "= 0.5.1"
 eyre = "= 0.6.12"
+futures = "= 0.3.32"
 glob = "= 0.3.3"
 json-subscriber = "= 0.2.8"
 nix = { version = "= 0.31.2", features = ["signal"] }
@@ -72,9 +73,10 @@ serde_with = { version = "= 3.18.0", default-features = false, features = [
   "base64",
 ] }
 tempfile = "= 3.27.0"
-time = { version = "= 0.3.47", features = ["macros"] }
 thiserror = "= 2.0.18"
+time = { version = "= 0.3.47", features = ["macros"] }
 tokio = { version = "= 1.51.0", features = ["full", "tracing"] }
+tokio-rustls = "= 0.26.4"
 tokio-util = { version = "= 0.7.18", features = ["full"] }
 tower = "= 0.5.3"
 tower-http = { version = "= 0.6.8", features = ["timeout"] }

Cargo.toml

@@ -15,6 +15,8 @@ axum-server.workspace = true
 axum.workspace = true
 durstr.workspace = true
 eyre.workspace = true
+futures.workspace = true
+tokio-rustls.workspace = true
 tokio.workspace = true
 tower-http.workspace = true
 tower.workspace = true

packages/ak-axum/Cargo.toml

@@ -0,0 +1 @@
+pub mod tls;

packages/ak-axum/src/accept/mod.rs

@@ -0,0 +1,54 @@
+use axum::{Extension, middleware::AddExtension};
+use axum_server::{accept::Accept, tls_rustls::RustlsAcceptor};
+use futures::future::BoxFuture;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_rustls::{rustls::pki_types::CertificateDer, server::TlsStream};
+use tower::Layer as _;
+use tracing::instrument;
+
+#[derive(Clone, Debug)]
+pub struct TlsState {
+    pub peer_certificates: Option<Vec<CertificateDer<'static>>>,
+}
+
+#[derive(Clone)]
+pub(crate) struct TlsAcceptor<A> {
+    inner: RustlsAcceptor<A>,
+}
+
+impl<A> TlsAcceptor<A> {
+    pub(crate) fn new(inner: RustlsAcceptor<A>) -> Self {
+        Self { inner }
+    }
+}
+
+impl<A, I, S> Accept<I, S> for TlsAcceptor<A>
+where
+    A: Accept<I, S> + Clone + Send + 'static,
+    A::Stream: AsyncRead + AsyncWrite + Unpin + Send,
+    A::Service: Send,
+    A::Future: Send,
+    I: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    S: Send + 'static,
+{
+    type Future = BoxFuture<'static, std::io::Result<(Self::Stream, Self::Service)>>;
+    type Service = AddExtension<A::Service, TlsState>;
+    type Stream = TlsStream<A::Stream>;
+
+    #[instrument(skip_all)]
+    fn accept(&self, stream: I, service: S) -> Self::Future {
+        let acceptor = self.inner.clone();
+
+        Box::pin(async move {
+            let (stream, service) = acceptor.accept(stream, service).await?;
+            let server_conn = stream.get_ref().1;
+            let tls_state = TlsState {
+                peer_certificates: server_conn.peer_certificates().map(|c| c.to_owned()),
+            };
+
+            let service = Extension(tls_state).layer(service);
+
+            Ok((stream, service))
+        })
+    }
+}

packages/ak-axum/src/accept/tls.rs

@@ -1,5 +1,6 @@
 //! Utilities for working with [`axum`].
 
+pub mod accept;
 pub mod error;
 pub mod router;
 pub mod server;

packages/ak-axum/src/lib.rs

@@ -12,6 +12,8 @@ use axum_server::{
 use eyre::Result;
 use tracing::info;
 
+use crate::accept::tls::TlsAcceptor;
+
 async fn run_plain(
     arbiter: Arbiter,
     name: &str,
@@ -119,7 +121,9 @@ async fn run_tls(
     arbiter.add_net_handle(handle.clone()).await;
 
     axum_server::Server::bind(addr)
-        .acceptor(RustlsAcceptor::new(config).acceptor(DefaultAcceptor::new()))
+        .acceptor(TlsAcceptor::new(
+            RustlsAcceptor::new(config).acceptor(DefaultAcceptor::new()),
+        ))
         .handle(handle)
         .serve(router.into_make_service_with_connect_info::<net::SocketAddr>())
         .await?;