Ocean logging middleware

DeFiCh · Oct 12, 2024 · 8dd6699 · 8dd6699
1 parent 054a8e4
commit 8dd6699
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 7 deletions.
diff --git a/lib/Cargo.toml b/lib/Cargo.toml
@@ -74,7 +74,7 @@ hex-literal = "0.4"
 bincode = "1.3"
 serde_with = "3.0"
 heck = "0.4"
-tower-http = { version = "0.4.0", features = ["full"] }
+tower-http = { version = "0.4", features = ["full"] }
 tower = "0.4.13"
 hyper = "0.14.20"
 

diff --git a/lib/ain-ocean/src/api/mod.rs b/lib/ain-ocean/src/api/mod.rs
@@ -1,13 +1,14 @@
-use std::{net::SocketAddr, str::FromStr, sync::Arc};
+use std::{net::SocketAddr, str::FromStr, sync::Arc, time::Instant};
 
 use axum::{
-    body::Body,
-    extract::{ConnectInfo, Request},
-    http::{HeaderValue, StatusCode},
-    middleware::{from_fn, Next},
+    body::{to_bytes, Body},
+    extract::{ConnectInfo, OriginalUri, Request},
+    http::{HeaderMap, HeaderValue, StatusCode},
+    middleware::{self, from_fn, Next},
     response::{IntoResponse, Response},
     Json, Router,
 };
+use log::{debug, log_enabled, trace, Level};
 
 mod address;
 mod block;
@@ -32,6 +33,7 @@ mod transactions;
 
 use defichain_rpc::Client;
 use serde::{Deserialize, Serialize};
+use serde_json::json;
 
 use crate::{network::Network, Result, Services};
 
@@ -134,7 +136,8 @@ pub async fn ocean_router(
             format!("/v0/{}", context.network).as_str(),
             main_router.merge(debug_router),
         )
-        .layer(from_fn(cors)))
+        .layer(from_fn(cors))
+        .layer(middleware::from_fn(logging_middleware)))
 }
 
 async fn localhost_only(
@@ -168,3 +171,69 @@ async fn localhost_only(
         Err(StatusCode::FORBIDDEN)
     }
 }
+
+const MAX_BODY_SIZE: usize = 1024 * 1024 * 16; // 16MB limit for body logging
+
+async fn logging_middleware(
+    OriginalUri(original_uri): OriginalUri,
+    req: Request<Body>,
+    next: Next,
+) -> std::result::Result<impl IntoResponse, StatusCode> {
+    let method = req.method().clone();
+    let path = req.uri().path().to_owned();
+    let query = req.uri().query().unwrap_or("").to_owned();
+
+    debug!("Request: {} {}", method, path);
+
+    if log_enabled!(Level::Trace) {
+        let headers = format_headers(req.headers());
+        let request_log = json!({
+            "method": method.as_str(),
+            "path": path,
+            "query": query,
+            "headers": headers,
+            "original_uri": original_uri.to_string(),
+        });
+        if let Ok(json) = serde_json::to_string(&request_log) {
+            trace!("Request: {json}");
+        }
+    }
+
+    let start = Instant::now();
+    let res = next.run(req).await;
+    let latency = start.elapsed();
+
+    debug!("Response: {} {} {:?}", method, path, latency);
+
+    if log_enabled!(Level::Trace) {
+        let (parts, body) = res.into_parts();
+        let bytes = to_bytes(body, MAX_BODY_SIZE).await.unwrap_or_default();
+        let body_str = String::from_utf8_lossy(&bytes);
+
+        let response_log = json!({
+            "status": parts.status.as_u16(),
+            "headers": format_headers(&parts.headers),
+            "body": body_str,
+        });
+        if let Ok(json) = serde_json::to_string(&response_log) {
+            trace!("Response: {json}",);
+        }
+
+        Ok(Response::from_parts(parts, Body::from(bytes)))
+    } else {
+        Ok(res)
+    }
+}
+
+fn format_headers(headers: &HeaderMap) -> serde_json::Value {
+    let mut map = serde_json::Map::new();
+    for (key, value) in headers.iter() {
+        if let Ok(v) = value.to_str() {
+            map.insert(
+                key.as_str().to_owned(),
+                serde_json::Value::String(v.to_owned()),
+            );
+        }
+    }
+    serde_json::Value::Object(map)
+}