From db9de1ba3cee03347e04fa312abef5b4e8e07a18 Mon Sep 17 00:00:00 2001
From: Marc 'risson' Schmitt <marc.schmitt@risson.space>
Date: Tue, 7 Apr 2026 17:11:53 +0000
Subject: [PATCH] packages/ak-axum/server: init (#21317)

---
 Cargo.lock                     |   1 +
 packages/ak-axum/Cargo.toml    |   1 +
 packages/ak-axum/src/lib.rs    |   1 +
 packages/ak-axum/src/server.rs | 147 +++++++++++++++++++++++++++++++++
 4 files changed, 150 insertions(+)
 create mode 100644 packages/ak-axum/src/server.rs

diff --git a/Cargo.lock b/Cargo.lock
index dbd32f7558..a0f37735f2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -129,6 +129,7 @@ version = "2026.5.0-rc1"
 dependencies = [
  "authentik-common",
  "axum",
+ "axum-server",
  "durstr",
  "eyre",
  "tokio",
diff --git a/packages/ak-axum/Cargo.toml b/packages/ak-axum/Cargo.toml
index 556f295eea..305ab89798 100644
--- a/packages/ak-axum/Cargo.toml
+++ b/packages/ak-axum/Cargo.toml
@@ -11,6 +11,7 @@ publish.workspace = true
 
 [dependencies]
 ak-common.workspace = true
+axum-server.workspace = true
 axum.workspace = true
 durstr.workspace = true
 eyre.workspace = true
diff --git a/packages/ak-axum/src/lib.rs b/packages/ak-axum/src/lib.rs
index 8ade489e20..f0dccda6d3 100644
--- a/packages/ak-axum/src/lib.rs
+++ b/packages/ak-axum/src/lib.rs
@@ -2,4 +2,5 @@
 
 pub mod error;
 pub mod router;
+pub mod server;
 pub mod tracing;
diff --git a/packages/ak-axum/src/server.rs b/packages/ak-axum/src/server.rs
new file mode 100644
index 0000000000..3869fe2525
--- /dev/null
+++ b/packages/ak-axum/src/server.rs
@@ -0,0 +1,147 @@
+//! Utilities to run an axum server.
+
+use std::{net, os::unix};
+
+use ak_common::arbiter::{Arbiter, Tasks};
+use axum::Router;
+use axum_server::{
+    Handle,
+    accept::DefaultAcceptor,
+    tls_rustls::{RustlsAcceptor, RustlsConfig},
+};
+use eyre::Result;
+use tracing::info;
+
+async fn run_plain(
+    arbiter: Arbiter,
+    name: &str,
+    router: Router,
+    addr: net::SocketAddr,
+    allow_failure: bool,
+) -> Result<()> {
+    info!(addr = addr.to_string(), "starting {name} server");
+
+    let handle = Handle::new();
+    arbiter.add_net_handle(handle.clone()).await;
+
+    let res = axum_server::Server::bind(addr)
+        .acceptor(DefaultAcceptor::new())
+        .handle(handle)
+        .serve(router.into_make_service_with_connect_info::<net::SocketAddr>())
+        .await;
+    if res.is_err() && allow_failure {
+        arbiter.shutdown().await;
+        return Ok(());
+    }
+    res?;
+
+    Ok(())
+}
+
+/// Start a plaintext server.
+///
+/// `name` is only used for observability purposes and should describe which module is starting the
+/// server.
+///
+/// `allow_failure` allows the server to fail silently.
+pub fn start_plain(
+    tasks: &mut Tasks,
+    name: &'static str,
+    router: Router,
+    addr: net::SocketAddr,
+    allow_failure: bool,
+) -> Result<()> {
+    let arbiter = tasks.arbiter();
+    tasks
+        .build_task()
+        .name(&format!("{}::run_plain({name}, {addr})", module_path!()))
+        .spawn(run_plain(arbiter, name, router, addr, allow_failure))?;
+    Ok(())
+}
+
+pub(crate) async fn run_unix(
+    arbiter: Arbiter,
+    name: &str,
+    router: Router,
+    addr: unix::net::SocketAddr,
+    allow_failure: bool,
+) -> Result<()> {
+    info!(addr = ?addr, "starting {name} server");
+
+    let handle = Handle::new();
+    arbiter.add_unix_handle(handle.clone()).await;
+
+    let res = axum_server::Server::bind(addr)
+        .acceptor(DefaultAcceptor::new())
+        .handle(handle)
+        .serve(router.into_make_service())
+        .await;
+    if res.is_err() && allow_failure {
+        arbiter.shutdown().await;
+        return Ok(());
+    }
+    res?;
+
+    Ok(())
+}
+
+/// Start a Unix socket server.
+///
+/// `name` is only used for observability purposes and should describe which module is starting the
+/// server.
+///
+/// `allow_failure` allows the server to fail silently.
+pub fn start_unix(
+    tasks: &mut Tasks,
+    name: &'static str,
+    router: Router,
+    addr: unix::net::SocketAddr,
+    allow_failure: bool,
+) -> Result<()> {
+    let arbiter = tasks.arbiter();
+    tasks
+        .build_task()
+        .name(&format!("{}::run_unix({name}, {addr:?})", module_path!()))
+        .spawn(run_unix(arbiter, name, router, addr, allow_failure))?;
+    Ok(())
+}
+
+async fn run_tls(
+    arbiter: Arbiter,
+    name: &str,
+    router: Router,
+    addr: net::SocketAddr,
+    config: RustlsConfig,
+) -> Result<()> {
+    info!(addr = addr.to_string(), "starting {name} server");
+
+    let handle = Handle::new();
+    arbiter.add_net_handle(handle.clone()).await;
+
+    axum_server::Server::bind(addr)
+        .acceptor(RustlsAcceptor::new(config).acceptor(DefaultAcceptor::new()))
+        .handle(handle)
+        .serve(router.into_make_service_with_connect_info::<net::SocketAddr>())
+        .await?;
+
+    Ok(())
+}
+
+/// Start a TLS server.
+///
+/// `name` is only used for observability purposes and should describe which module is starting the
+/// server.
+pub fn start_tls(
+    tasks: &mut Tasks,
+    name: &'static str,
+    router: Router,
+    addr: net::SocketAddr,
+    config: RustlsConfig,
+) -> Result<()> {
+    let arbiter = tasks.arbiter();
+    tasks
+        .build_task()
+        .name(&format!("{}::run_tls({name}, {addr})", module_path!()))
+        .spawn(run_tls(arbiter, name, router, addr, config))?;
+    Ok(())
+}