From 70ee206031649490d6b48bc38a3fbe215c54036c Mon Sep 17 00:00:00 2001
From: Lambda <csgl@xiretza.xyz>
Date: Sun, 15 Sep 2024 17:04:47 +0000
Subject: [PATCH] Extract source address for requests

---
 Cargo.lock       |  2 ++
 Cargo.toml       |  2 +-
 src/cli/serve.rs | 36 +++++++++++++++++++++++++++++++-----
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 65aa7f4e..80ea091b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -158,6 +158,7 @@ dependencies = [
  "serde_path_to_error",
  "serde_urlencoded",
  "sync_wrapper 1.0.1",
+ "tokio",
  "tower 0.5.1",
  "tower-layer",
  "tower-service",
@@ -3069,6 +3070,7 @@ dependencies = [
  "futures-util",
  "pin-project-lite",
  "sync_wrapper 0.1.2",
+ "tokio",
  "tower-layer",
  "tower-service",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index eebcbd2a..6157e966 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -88,7 +88,7 @@ workspace = true
 [dependencies]
 argon2 = "0.5.3"
 async-trait = "0.1.82"
-axum = { version = "0.7.6", default-features = false, features = ["form", "http1", "http2", "json", "matched-path", "tracing"] }
+axum = { version = "0.7.6", default-features = false, features = ["form", "http1", "http2", "json", "matched-path", "tokio", "tracing"] }
 axum-extra = { version = "0.9.4", features = ["typed-header"] }
 axum-server = { version = "0.7.1", features = ["tls-rustls-no-provider"] }
 base64 = "0.22.1"
diff --git a/src/cli/serve.rs b/src/cli/serve.rs
index aa507522..69d5eb2b 100644
--- a/src/cli/serve.rs
+++ b/src/cli/serve.rs
@@ -4,9 +4,13 @@ use std::{
 };
 
 use axum::{
-    extract::{DefaultBodyLimit, FromRequestParts, MatchedPath},
+    extract::{
+        connect_info::IntoMakeServiceWithConnectInfo, ConnectInfo,
+        DefaultBodyLimit, FromRequestParts, MatchedPath,
+    },
+    middleware::AddExtension,
     response::IntoResponse,
-    routing::{any, get, on, IntoMakeService, MethodFilter, Route},
+    routing::{any, get, on, MethodFilter, Route},
     Router,
 };
 use axum_server::{
@@ -21,6 +25,7 @@ use http::{
     Method, StatusCode, Uri,
 };
 use hyper::body::Incoming;
+use proxy_header::ProxyHeader;
 use ruma::api::{
     client::{
         error::{Error as RumaError, ErrorBody, ErrorKind},
@@ -212,9 +217,13 @@ where
         &mut self,
         listen: ListenConfig,
         server: Server<A>,
-        app: IntoMakeService<Router>,
+        app: IntoMakeServiceWithConnectInfo<Router, SocketAddr>,
     ) where
-        A: Accept<TcpStream, Router> + Clone + Send + Sync + 'static,
+        A: Accept<TcpStream, AddExtension<Router, ConnectInfo<SocketAddr>>>
+            + Clone
+            + Send
+            + Sync
+            + 'static,
         A::Stream: AsyncRead + AsyncWrite + Unpin + Send + Sync,
         A::Service: SendService<http::Request<Incoming>> + Send,
         A::Future: Send,
@@ -235,7 +244,7 @@ where
     ) -> Result<(), error::Serve> {
         let app = routes(self.config, &listen.components)
             .layer(self.middlewares.clone())
-            .into_make_service();
+            .into_make_service_with_connect_info::<SocketAddr>();
 
         match listen.transport {
             ListenTransport::Tcp {
@@ -298,11 +307,28 @@ async fn run_server() -> Result<(), error::Serve> {
 
                     let method = request.method();
 
+                    let source_address = request
+                        .extensions()
+                        .get::<ProxyHeader<'_>>()
+                        .map_or_else(
+                            || {
+                                request
+                                    .extensions()
+                                    .get::<ConnectInfo<SocketAddr>>()
+                                    .map(|&ConnectInfo(addr)| addr)
+                            },
+                            |h| h.proxied_address().map(|addr| addr.source),
+                        );
+
                     tracing::info_span!(
                         "http_request",
                         otel.name = format!("{method} {endpoint}"),
                         %method,
                         %endpoint,
+                        source_address = source_address.map_or(
+                            "[unknown]".to_owned(),
+                            |a| a.to_string()
+                        ),
                     )
                 })
                 .on_request(