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/routes.rs

@@ -1,15 +1,20 @@
 use axum::{
     async_trait,
+    body::Body,
     extract::{rejection::JsonRejection, ws::WebSocketUpgrade, FromRequest, Request, State},
     http::StatusCode,
-    response::IntoResponse,
+    middleware::{self, Next},
+    response::{IntoResponse, Response},
     routing::{get, post},
     Router,
 };
 use compositor_pipeline::Pipeline;
+use core::str;
+use http_body_util::BodyExt;
+use log::info;
 use serde_json::{json, Value};
 
-use crate::state::{ApiState, Response};
+use crate::state::{self, ApiState};
 
 use compositor_api::error::ApiError;
 
@@ -54,9 +59,9 @@ pub fn routes(state: ApiState) -> Router {
         .route("/:id/register", post(register_request::handle_shader))
         .route("/:id/unregister", post(unregister_request::handle_shader));
 
-    async fn handle_start(State(state): State<ApiState>) -> Result<Response, ApiError> {
+    async fn handle_start(State(state): State<ApiState>) -> Result<state::Response, ApiError> {
         Pipeline::start(&state.pipeline);
-        Ok(Response::Ok {})
+        Ok(state::Response::Ok {})
     }
 
     Router::new()
@@ -75,6 +80,7 @@ pub fn routes(state: ApiState) -> Router {
                 "instance_id": state.config.instance_id
             }))),
         )
+        .layer(middleware::from_fn(log_request_response_body))
         .with_state(state)
 }
 
@@ -111,3 +117,42 @@ where
         }
     }
 }
+
+async fn log_request_response_body(
+    request: Request,
+    next: Next,
+) -> Result<impl IntoResponse, Response> {
+    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();
+
+    info!("Request body: {:?}", str::from_utf8(&bytes).unwrap());
+
+    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();
+
+    info!("Response body: {:?}", str::from_utf8(&bytes).unwrap());
+
+    Ok(Response::from_parts(parts, Body::from(bytes)))
+}