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

rissson · rissson · commit 107b0d1b549b · 2026-04-07T15:55:56.000+02:00
Squashed commit of the following: commit 65f33d6 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 18:35:12 2026 +0200 packages/ak-axum/accept/tls: init Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 116e601 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 18:28:43 2026 +0200 packages/ak-axum/server: init Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 67ed0dc Merge: b4f06d6 095d38c Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 17:17:40 2026 +0200 Merge branch 'rust-config' into rust-axum Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 095d38c Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 17:11:52 2026 +0200 fixup Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 7eb07d5 Merge: cb1f86b 507eb2f Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 17:11:08 2026 +0200 Merge branch 'rust-arbiter' into rust-config commit 507eb2f Merge: 6968d59 0079984 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 17:10:37 2026 +0200 Merge branch 'rust-lib-rename' into rust-arbiter Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 0079984 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 17:08:07 2026 +0200 packages/ak-common: rename from ak-lib Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit b4f06d6 Merge: c08e2a3 cb1f86b Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 17:05:36 2026 +0200 wip Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit c08e2a3 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 17:05:01 2026 +0200 wip Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 75622e9 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 16:57:45 2026 +0200 packages/ak-axum: init Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit cb1f86b Merge: c8fa1b8 a3fea5d Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 16:53:19 2026 +0200 Merge branch 'fix-rustfmt-config' into rust-config commit 6968d59 Merge: 341c9cc a3fea5d Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 16:52:55 2026 +0200 Merge branch 'fix-rustfmt-config' into rust-arbiter commit a3fea5d Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Wed Apr 1 16:32:16 2026 +0200 root: fix rustfmt config Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit c8fa1b8 Merge: 48c833c 341c9cc Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Tue Mar 31 13:56:39 2026 +0200 Merge branch 'rust-arbiter' into rust-config commit 341c9cc Merge: a1cf0a7 55e555c Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Tue Mar 31 13:55:28 2026 +0200 Merge branch 'main' into rust-arbiter Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 48c833c Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 20:31:04 2026 +0200 lint Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 681117d Merge: 4c54511 a1cf0a7 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 20:03:49 2026 +0200 Merge branch 'rust-arbiter' into rust-config Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit a1cf0a7 Merge: 1ee6f11 524b788 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 20:03:01 2026 +0200 Merge branch 'rust-lib' into rust-arbiter Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 524b788 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:57:51 2026 +0200 fixup Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit dc65ab1 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 16:33:49 2026 +0200 packages/ak-lib: init Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 4c54511 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:52:55 2026 +0200 move into lib crate Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit d7141df Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:48:34 2026 +0200 fixup Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 2bab7ed Merge: 2bc79f1 1ee6f11 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:43:24 2026 +0200 Merge branch 'rust-arbiter' into rust-config commit 1ee6f11 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:43:21 2026 +0200 fixup Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 2bc79f1 Merge: e7d3704 27ff039 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:42:58 2026 +0200 Merge branch 'rust-arbiter' into rust-config Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 27ff039 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:40:37 2026 +0200 rename to ak-lib Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit e7d3704 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:32:59 2026 +0200 packages/ak-config: init Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 1e5cb4b Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 19:32:18 2026 +0200 sort out package versions Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 64b9391 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 17:49:16 2026 +0200 lint Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 57edeec Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 17:48:51 2026 +0200 add tests Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit 5294f8a Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 16:38:11 2026 +0200 fixup Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> commit f1257d2 Author: Marc 'risson' Schmitt <marc.schmitt@risson.space> Date: Mon Mar 30 16:33:49 2026 +0200 packages/ak-arbiter: init Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space> Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -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"] }
@@ -76,6 +77,7 @@ 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"] }
diff --git a/packages/ak-axum/Cargo.toml b/packages/ak-axum/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
diff --git a/packages/ak-axum/src/accept/mod.rs b/packages/ak-axum/src/accept/mod.rs
@@ -0,0 +1 @@
+pub mod tls;
diff --git a/packages/ak-axum/src/accept/tls.rs b/packages/ak-axum/src/accept/tls.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))
+        })
+    }
+}
diff --git a/packages/ak-axum/src/lib.rs b/packages/ak-axum/src/lib.rs
@@ -1,5 +1,6 @@
 //! Utilities for working with [`axum`].
 
+pub mod accept;
 pub mod error;
 pub mod router;
 pub mod server;
diff --git a/packages/ak-axum/src/server.rs b/packages/ak-axum/src/server.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?;
