Add logger to axum router

Cargo.toml

@@ -95,6 +95,7 @@ tokio = { workspace = true }
 axum = { version = "0.7.4", features = ["ws"] }
 futures-util = { workspace = true }
 wgpu = { workspace = true }
+http-body-util = "0.1.2"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 shared_memory = { workspace = true }

src/lib.rs

@@ -1,5 +1,6 @@
 pub mod config;
 pub mod logger;
+pub mod middleware;
 pub mod routes;
 pub mod server;
 pub mod state;

src/main.rs

@@ -2,6 +2,7 @@ use log::info;
 
 mod config;
 mod logger;
+mod middleware;
 mod routes;
 mod server;
 mod state;

src/middleware.rs

@@ -0,0 +1,77 @@
+use axum::{
+    body::Body,
+    extract::Request,
+    http::StatusCode,
+    middleware::Next,
+    response::{IntoResponse, Response},
+};
+use core::str;
+use http_body_util::BodyExt;
+use serde_json::Value;
+use tracing::{enabled, trace, Level};
+
+pub async fn body_logger_middleware(
+    request: Request,
+    next: Next,
+) -> Result<impl IntoResponse, Response> {
+    if !enabled!(target: "live_compositor::log_request_body", Level::TRACE) {
+        return Ok(next.run(request).await);
+    }
+    let request = buffer_request_body(request).await?;
+    let response = next.run(request).await;
+    let response = buffer_response_body(response).await?;
+
+    Ok(response)
+}
+
+async fn buffer_request_body(request: Request) -> Result<Request, Response> {
+    let (parts, body) = request.into_parts();
+
+    let bytes = body
+        .collect()
+        .await
+        .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response())?
+        .to_bytes();
+
+    match serde_json::from_slice::<Value>(&bytes) {
+        Ok(body_json) => {
+            trace!(target: "live_compositor::log_request_body", method = ?parts.method, path = ?parts.uri, "Request body: {}", body_json);
+        }
+        Err(_) => match str::from_utf8(&bytes) {
+            Ok(body_str) => {
+                trace!(target: "live_compositor::log_request_body", method = ?parts.method, path = ?parts.uri, "Request body: {}", body_str);
+            }
+            Err(_) => {
+                trace!(target: "live_compositor::log_request_body", method = ?parts.method, path = ?parts.uri, "Request body: {:?}", bytes);
+            }
+        },
+    }
+
+    Ok(Request::from_parts(parts, Body::from(bytes)))
+}
+
+async fn buffer_response_body(response: Response) -> Result<Response, Response> {
+    let (parts, body) = response.into_parts();
+
+    let bytes = body
+        .collect()
+        .await
+        .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response())?
+        .to_bytes();
+
+    match serde_json::from_slice::<Value>(&bytes) {
+        Ok(body_json) => {
+            trace!(target: "live_compositor::log_request_body", status=?parts.status, "Response body: {}", body_json);
+        }
+        Err(_) => match str::from_utf8(&bytes) {
+            Ok(body_str) => {
+                trace!(target: "live_compositor::log_request_body", status=?parts.status, "Response body: {}", body_str);
+            }
+            Err(_) => {
+                trace!(target: "live_compositor::log_request_body", status=?parts.status, "Response body: {:?}", bytes);
+            }
+        },
+    }
+
+    Ok(Response::from_parts(parts, Body::from(bytes)))
+}

src/routes.rs

@@ -2,6 +2,7 @@ use axum::{
     async_trait,
     extract::{rejection::JsonRejection, ws::WebSocketUpgrade, FromRequest, Request, State},
     http::StatusCode,
+    middleware,
     response::IntoResponse,
     routing::{get, post},
     Router,
@@ -17,6 +18,7 @@ use self::{
     update_output::handle_keyframe_request, update_output::handle_output_update,
     ws::handle_ws_upgrade,
 };
+use crate::middleware::body_logger_middleware;
 
 mod register_request;
 mod unregister_request;
@@ -75,6 +77,7 @@ pub fn routes(state: ApiState) -> Router {
                 "instance_id": state.config.instance_id
             }))),
         )
+        .layer(middleware::from_fn(body_logger_middleware))
         .with_state(state)
 }