Merge pull request #28 from flashbots/peg/logging

Add logging - trying to give similar logging behaviour / options as `cvm-reverse-proxy`

Author: ameba23

Cargo.lock

@@ -5,8 +5,8 @@ edition = "2024"
 license = "MIT"
 
 [dependencies]
-tokio = { version = "1.48.0", features = ["full"]}
-tokio-rustls = { version = "0.26.4", default-features = false, features = ["ring"]}
+tokio = { version = "1.48.0", features = ["full"] }
+tokio-rustls = { version = "0.26.4", default-features = false, features = ["ring"] }
 sha2 = "0.10.9"
 x509-parser = "0.18.0"
 thiserror = "2.0.17"
@@ -27,6 +27,8 @@ bytes = "1.10.1"
 http = "1.3.1"
 serde_json = "1.0.145"
 serde = "1.0.228"
+tracing = "0.1.41"
+tracing-subscriber = { version = "0.3.20", features = ["env-filter", "json"] }
 parity-scale-codec = "3.7.5"
 
 [dev-dependencies]

Cargo.toml

@@ -11,6 +11,7 @@ use parity_scale_codec::{Decode, Encode};
 use thiserror::Error;
 use tokio::sync::{mpsc, oneshot};
 use tokio_rustls::rustls::server::{VerifierBuilderError, WebPkiClientVerifier};
+use tracing::{error, warn};
 
 #[cfg(test)]
 mod test_helpers;
@@ -153,7 +154,7 @@ impl ProxyServer {
             )
             .await
             {
-                eprintln!("Failed to handle connection: {err}");
+                warn!("Failed to handle connection: {err}");
             }
         });
 
@@ -244,7 +245,7 @@ impl ProxyServer {
                     Err(e) => {
                         // This error is highly unlikely - that the measurement values fail to
                         // encode to JSON or fit in an HTTP header
-                        eprintln!("Failed to encode measurement values: {e}");
+                        error!("Failed to encode measurement values: {e}");
                     }
                 }
                 headers.insert(
@@ -260,7 +261,7 @@ impl ProxyServer {
                         Ok::<Response<BoxBody<bytes::Bytes, hyper::Error>>, hyper::Error>(res)
                     }
                     Err(e) => {
-                        eprintln!("Failed to handle a request from a proxy-client: {e}");
+                        warn!("Failed to handle a request from a proxy-client: {e}");
                         let mut resp = Response::new(full(format!("Request failed: {e}")));
                         *resp.status_mut() = hyper::StatusCode::BAD_GATEWAY;
                         Ok(resp)
@@ -291,15 +292,15 @@ impl ProxyServer {
         // Drive the connection
         tokio::spawn(async move {
             if let Err(e) = conn.await {
-                eprintln!("Client connection error: {e}");
+                warn!("Client connection error: {e}");
             }
         });
 
         // Forward the request from the proxy-client to the target server
         match sender.send_request(req).await {
             Ok(resp) => Ok(resp.map(|b| b.boxed())),
             Err(e) => {
-                eprintln!("send_request error: {e}");
+                warn!("send_request error: {e}");
                 let mut resp = Response::new(full(format!("Request failed: {e}")));
                 *resp.status_mut() = hyper::StatusCode::BAD_GATEWAY;
                 Ok(resp)
@@ -428,7 +429,7 @@ impl ProxyClient {
                                 Err(e) => {
                                     // This error is highly unlikely - that the measurement values fail to
                                     // encode to JSON or fit in an HTTP header
-                                    eprintln!("Failed to encode measurement values: {e}");
+                                    error!("Failed to encode measurement values: {e}");
                                 }
                             }
                             headers.insert(
@@ -440,7 +441,7 @@ impl ProxyClient {
                         (Ok(resp.map(|b| b.boxed())), false)
                     }
                     Err(e) => {
-                        eprintln!("Failed to send request to proxy-server: {e}");
+                        warn!("Failed to send request to proxy-server: {e}");
                         let mut resp = Response::new(full(format!("Request failed: {e}")));
                         *resp.status_mut() = hyper::StatusCode::BAD_GATEWAY;
 
@@ -450,7 +451,7 @@ impl ProxyClient {
 
                 // Send the response back to the source client
                 if response_tx.send(response).is_err() {
-                    eprintln!("Failed to forward response to source client, probably they dropped the connection");
+                    warn!("Failed to forward response to source client, probably they dropped the connection");
                 }
 
                 // If the connection to the proxy server failed, reconnect
@@ -488,7 +489,7 @@ impl ProxyClient {
 
         tokio::spawn(async move {
             if let Err(err) = Self::handle_connection(inbound, requests_tx).await {
-                eprintln!("Failed to handle connection from source client: {err}");
+                warn!("Failed to handle connection from source client: {err}");
             }
         });
 
@@ -510,7 +511,7 @@ impl ProxyClient {
                         Ok::<Response<BoxBody<bytes::Bytes, hyper::Error>>, hyper::Error>(res)
                     }
                     Err(e) => {
-                        eprintln!("send_request error: {e}");
+                        warn!("send_request error: {e}");
                         let mut resp = Response::new(full(format!("Request failed: {e}")));
                         *resp.status_mut() = hyper::StatusCode::BAD_GATEWAY;
                         Ok(resp)
@@ -551,7 +552,7 @@ impl ProxyClient {
                     return output;
                 }
                 Err(e) => {
-                    eprintln!("Reconnect failed: {e}. Retrying in {:#?}...", delay);
+                    warn!("Reconnect failed: {e}. Retrying in {:#?}...", delay);
                     tokio::time::sleep(delay).await;
 
                     // increase delay for next time (exponential), but clamp to max_delay
@@ -634,7 +635,7 @@ impl ProxyClient {
         // Drive the connection
         tokio::spawn(async move {
             if let Err(e) = conn.await {
-                eprintln!("Client connection error: {e}");
+                warn!("Client connection error: {e}");
             }
         });
 

src/lib.rs

@@ -2,6 +2,7 @@ use anyhow::{anyhow, ensure};
 use clap::{Parser, Subcommand};
 use std::{fs::File, net::SocketAddr, path::PathBuf};
 use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};
+use tracing::level_filters::LevelFilter;
 
 use attested_tls_proxy::{
     attestation::{measurements::get_measurements_from_file, AttestationType, AttestationVerifier},
@@ -13,15 +14,13 @@ use attested_tls_proxy::{
 struct Cli {
     #[clap(subcommand)]
     command: CliCommand,
-    // TODO missing:
-    // Name:  "log-json",
-    // Value: false,
-    // Usage: "log in JSON format",
-    //
-    // Name:  "log-debug",
-    // Value: true,
-    // Usage: "log debug messages",
-    //
+    /// Log debug messages
+    #[arg(long, global = true)]
+    log_debug: bool,
+    /// Log in JSON format
+    #[arg(long, global = true)]
+    log_json: bool,
+    // TODO still missing
     // Name:    "log-dcap-quote",
     // EnvVars: []string{"LOG_DCAP_QUOTE"},
     // Value:   false,
@@ -106,6 +105,24 @@ enum CliCommand {
 async fn main() -> anyhow::Result<()> {
     let cli = Cli::parse();
 
+    let level_filter = if cli.log_debug {
+        LevelFilter::DEBUG
+    } else {
+        LevelFilter::WARN
+    };
+
+    let env_filter = tracing_subscriber::EnvFilter::builder()
+        .with_default_directive(level_filter.into())
+        .from_env_lossy();
+
+    let subscriber = tracing_subscriber::fmt::Subscriber::builder().with_env_filter(env_filter);
+
+    if cli.log_json {
+        subscriber.json().init();
+    } else {
+        subscriber.pretty().init();
+    }
+
     match cli.command {
         CliCommand::Client {
             listen_addr,
@@ -170,7 +187,7 @@ async fn main() -> anyhow::Result<()> {
 
             loop {
                 if let Err(err) = client.accept().await {
-                    eprintln!("Failed to handle connection: {err}");
+                    tracing::error!("Failed to handle connection: {err}");
                 }
             }
         }
@@ -211,7 +228,7 @@ async fn main() -> anyhow::Result<()> {
 
             loop {
                 if let Err(err) = server.accept().await {
-                    eprintln!("Failed to handle connection: {err}");
+                    tracing::error!("Failed to handle connection: {err}");
                 }
             }
         }