/ radicle-httpd / src / tracing_extra.rs
tracing_extra.rs
 1  use std::fmt;
 2  use std::net::SocketAddr;
 3  use std::sync::atomic::{AtomicU64, Ordering};
 4  use std::sync::Arc;
 5  
 6  use axum::extract::ConnectInfo;
 7  use axum::http::Request;
 8  use axum::middleware::Next;
 9  use axum::response::IntoResponse;
10  use axum::Extension;
11  use hyper::{Method, StatusCode, Uri, Version};
12  
13  pub use radicle_term::ansi::Paint;
14  
15  #[derive(Clone)]
16  pub struct RequestId(Arc<AtomicU64>);
17  
18  impl RequestId {
19      pub fn new() -> RequestId {
20          RequestId(Arc::new(0.into()))
21      }
22  
23      pub fn next(&mut self) -> u64 {
24          self.0.fetch_add(1, Ordering::SeqCst)
25      }
26  }
27  
28  pub struct TracingInfo {
29      pub connect_info: ConnectInfo<SocketAddr>,
30      pub method: Method,
31      pub version: Version,
32      pub uri: Uri,
33  }
34  
35  pub struct ColoredStatus(pub StatusCode);
36  
37  impl fmt::Display for ColoredStatus {
38      fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
39          match self.0.as_u16() {
40              200..=299 => write!(f, "{}", Paint::green(self.0)),
41              300..=399 => write!(f, "{}", Paint::blue(self.0)),
42              400..=499 => write!(f, "{}", Paint::red(self.0)),
43              _ => write!(f, "{}", Paint::yellow(self.0)),
44          }
45      }
46  }
47  
48  pub async fn tracing_middleware<B>(request: Request<B>, next: Next<B>) -> impl IntoResponse {
49      let connect_info = *request
50          .extensions()
51          .get::<ConnectInfo<std::net::SocketAddr>>()
52          .unwrap();
53  
54      let method = request.method().clone();
55      let version = request.version();
56      let uri = request.uri().clone();
57  
58      let tracing_info = TracingInfo {
59          connect_info,
60          method,
61          version,
62          uri,
63      };
64  
65      let response = next.run(request).await;
66  
67      (Extension(tracing_info), response)
68  }