packages/ak-axum/accept/proxy_protocol: init (goauthentik#21319)

Author: rissson

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

@@ -1 +1,2 @@
+pub mod proxy_protocol;
 pub mod tls;

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

@@ -0,0 +1,85 @@
+use std::{io, time::Duration};
+
+use ak_common::tokio::proxy_protocol::{ProxyProtocolStream, header::Header};
+use axum::{Extension, middleware::AddExtension};
+use axum_server::accept::{Accept, DefaultAcceptor};
+use futures::future::BoxFuture;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tower::Layer as _;
+use tracing::instrument;
+
+#[derive(Clone, Debug)]
+pub struct ProxyProtocolState {
+    pub header: Option<Header<'static>>,
+}
+
+#[derive(Clone)]
+pub(crate) struct ProxyProtocolAcceptor<A = DefaultAcceptor> {
+    inner: A,
+    parsing_timeout: Duration,
+}
+
+impl ProxyProtocolAcceptor {
+    pub(crate) fn new() -> Self {
+        let inner = DefaultAcceptor::new();
+
+        #[cfg(not(test))]
+        let parsing_timeout = Duration::from_secs(10);
+
+        // Don't force tests to wait too long
+        #[cfg(test)]
+        let parsing_timeout = Duration::from_secs(1);
+
+        Self {
+            inner,
+            parsing_timeout,
+        }
+    }
+}
+
+impl Default for ProxyProtocolAcceptor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<A> ProxyProtocolAcceptor<A> {
+    pub(crate) fn acceptor<Acceptor>(self, acceptor: Acceptor) -> ProxyProtocolAcceptor<Acceptor> {
+        ProxyProtocolAcceptor {
+            inner: acceptor,
+            parsing_timeout: self.parsing_timeout,
+        }
+    }
+}
+
+impl<A, I, S> Accept<I, S> for ProxyProtocolAcceptor<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, io::Result<(Self::Stream, Self::Service)>>;
+    type Service = AddExtension<A::Service, ProxyProtocolState>;
+    type Stream = ProxyProtocolStream<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 stream = ProxyProtocolStream::new(stream).await?;
+
+            let proxy_protocol_state = ProxyProtocolState {
+                header: stream.header().cloned(),
+            };
+
+            let service = Extension(proxy_protocol_state).layer(service);
+
+            Ok((stream, service))
+        })
+    }
+}

packages/ak-axum/src/server.rs

@@ -12,7 +12,7 @@ use axum_server::{
 use eyre::Result;
 use tracing::info;
 
-use crate::accept::tls::TlsAcceptor;
+use crate::accept::{proxy_protocol::ProxyProtocolAcceptor, tls::TlsAcceptor};
 
 async fn run_plain(
     arbiter: Arbiter,
@@ -27,7 +27,7 @@ async fn run_plain(
     arbiter.add_net_handle(handle.clone()).await;
 
     let res = axum_server::Server::bind(addr)
-        .acceptor(DefaultAcceptor::new())
+        .acceptor(ProxyProtocolAcceptor::new().acceptor(DefaultAcceptor::new()))
         .handle(handle)
         .serve(router.into_make_service_with_connect_info::<net::SocketAddr>())
         .await;
@@ -121,9 +121,9 @@ async fn run_tls(
     arbiter.add_net_handle(handle.clone()).await;
 
     axum_server::Server::bind(addr)
-        .acceptor(TlsAcceptor::new(
+        .acceptor(ProxyProtocolAcceptor::new().acceptor(TlsAcceptor::new(
             RustlsAcceptor::new(config).acceptor(DefaultAcceptor::new()),
-        ))
+        )))
         .handle(handle)
         .serve(router.into_make_service_with_connect_info::<net::SocketAddr>())
         .await?;