improved logging

Cargo.lock

@@ -182,6 +182,28 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "axum-extra"
+version = "0.12.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fef252edff26ddba56bbcdf2ee3307b8129acb86f5749b68990c168a6fcc9c76"
+dependencies = [
+ "axum",
+ "axum-core",
+ "bytes",
+ "futures-core",
+ "futures-util",
+ "headers",
+ "http",
+ "http-body",
+ "http-body-util",
+ "mime",
+ "pin-project-lite",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
 [[package]]
 name = "base64"
 version = "0.22.1"
@@ -602,6 +624,7 @@ dependencies = [
  "anyhow",
  "argon2",
  "axum",
+ "axum-extra",
  "chrono",
  "clap",
  "dashmap",
@@ -813,6 +836,30 @@ dependencies = [
  "hashbrown 0.15.5",
 ]
 
+[[package]]
+name = "headers"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb"
+dependencies = [
+ "base64",
+ "bytes",
+ "headers-core",
+ "http",
+ "httpdate",
+ "mime",
+ "sha1",
+]
+
+[[package]]
+name = "headers-core"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4"
+dependencies = [
+ "http",
+]
+
 [[package]]
 name = "heck"
 version = "0.5.0"
@@ -1196,6 +1243,15 @@ version = "0.4.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
 
+[[package]]
+name = "matchers"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9"
+dependencies = [
+ "regex-automata",
+]
+
 [[package]]
 name = "matchit"
 version = "0.8.4"
@@ -2569,10 +2625,14 @@ version = "0.3.22"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2f30143827ddab0d256fd843b7a66d164e9f271cfa0dde49142c5ca0ca291f1e"
 dependencies = [
+ "matchers",
  "nu-ansi-term",
+ "once_cell",
+ "regex-automata",
  "sharded-slab",
  "smallvec",
  "thread_local",
+ "tracing",
  "tracing-core",
  "tracing-log",
 ]

Cargo.toml

@@ -7,6 +7,7 @@ edition = "2024"
 anyhow = "1.0.99"
 argon2 = "0.5.3"
 axum = { version = "0.8.4", features = ["multipart", "ws"] }
+axum-extra = { version = "0.12.5", features = ["typed-header"] }
 chrono = { version = "0.4.42", features = ["serde"] }
 clap = { version = "4.5.53", features = ["derive"] }
 dashmap = "6.1.0"
@@ -19,6 +20,6 @@ tokio = { version = "1.47.1", features = ["rt-multi-thread", "macros"] }
 tower-http = { version = "0.6.6", features = ["cors", "limit", "trace"] }
 tower_governor = "0.8.0"
 tracing = "0.1.41"
-tracing-subscriber = "0.3.20"
+tracing-subscriber = { version = "0.3.20", features = ["env-filter"] }
 uuid = { version = "1.19.0", features = ["serde", "v7"] }
 validator = "0.20.0"

src/main.rs

@@ -19,6 +19,7 @@ use tower_http::{
 };
 use tracing::Level;
 use tracing::info;
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
 mod auth;
 mod db;
@@ -58,11 +59,24 @@ pub struct Cli {
 async fn main() -> anyhow::Result<()> {
     let cli = Cli::parse();
 
-    tracing_subscriber::fmt()
+    let log_level = if cli.verbose { "debug" } else { "info" };
-        .with_max_level(tracing::Level::INFO)
+
+    tracing_subscriber::registry()
+        .with(
+            tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| {
+                format!(
+                    "{}={},tower_http={},sqlx=warn",
+                    env!("CARGO_CRATE_NAME"),
+                    log_level,
+                    log_level
+                )
+                .into()
+            }),
+        )
+        .with(tracing_subscriber::fmt::layer())
         .init();
 
-    tracing::info!("Connecting to database...");
+    // tracing::info!("Database connection is lazy, it may be unreachable");
     let db_pool = db::init_db(cli.database).await?;
 
     let cors = CorsLayer::new()
@@ -96,7 +110,7 @@ async fn main() -> anyhow::Result<()> {
         prohibit_registration: cli.no_registration,
     });
 
-    let mut app = Router::new()
+    let app = Router::new()
         .route("/version", get(get_version))
         .merge(routes::users::routes())
         .merge(routes::rooms::routes())
@@ -107,16 +121,13 @@ async fn main() -> anyhow::Result<()> {
         .layer(Extension(realtime))
         .layer(Extension(config))
         .layer(GovernorLayer::new(governor_conf))
-        .layer(cors);
+        .layer(cors)
-
+        .layer(
-    if cli.verbose {
-        app = app.layer(
             TraceLayer::new_for_http()
-                .make_span_with(DefaultMakeSpan::new().level(Level::INFO))
+                .make_span_with(DefaultMakeSpan::new().level(Level::DEBUG))
-                .on_response(DefaultOnResponse::new().level(Level::INFO)),
+                .on_request(())
-        )
+                .on_response(DefaultOnResponse::new().level(Level::DEBUG)),
-        // .layer(middleware::from_fn(log_json_body));
+        );
-    }
 
     let port = cli.port;
     let addr = format!("0.0.0.0:{port}");

src/routes/rooms.rs

@@ -296,8 +296,6 @@ async fn send_invite(
             )
         })?;
 
-    tracing::info!("bro");
-
     let room_name = room_name_from_uuid(&db, payload.room_uuid).await?;
 
     Ok((

src/routes/users.rs

@@ -279,7 +279,7 @@ async fn upload_avatar(
     let claims = verify_jwt(headers)?;
 
     let user_id = user_id_from_uuid(&db, claims.sub).await?;
-    tracing::info!("User ID {} is uploading {} bytes)", user_id, body.len());
+    tracing::debug!("User ID {} is uploading {} bytes)", user_id, body.len());
 
     let base_dir = &config.avatar_dir;
     let file_extension = "png"; // TODO: detect MIME type

src/routes/ws.rs

@@ -1,6 +1,6 @@
 use axum::Json;
-use axum::extract::Query;
 use axum::extract::ws::{Message as WsMessage, WebSocket};
+use axum::extract::{ConnectInfo, Query};
 use axum::http::HeaderMap;
 use axum::routing::get;
 use axum::{
@@ -9,7 +9,9 @@ use axum::{
     http::StatusCode,
     response::IntoResponse,
 };
+use axum_extra::{TypedHeader, headers};
 use serde::Deserialize;
+use std::net::SocketAddr;
 use std::time::Duration;
 use tokio::select;
 use uuid::Uuid;
@@ -77,8 +79,10 @@ pub async fn issue_ws_token(
 
 async fn ws_handler(
     ws: WebSocketUpgrade,
+    user_agent: Option<TypedHeader<headers::UserAgent>>,
     Path(room_uuid): Path<Uuid>,
     Query(query): Query<WsAuthQuery>,
+    ConnectInfo(addr): ConnectInfo<SocketAddr>,
     Extension(realtime): Extension<Realtime>,
     Extension(db): Extension<sqlx::PgPool>,
 ) -> Result<impl IntoResponse, axum::http::StatusCode> {
@@ -110,11 +114,19 @@ async fn ws_handler(
     let sender = realtime.sender_for(room_id);
     let receiver = sender.subscribe();
 
-    Ok(ws.on_upgrade(move |socket| handle_socket(socket, receiver)))
+    let user_agent = if let Some(TypedHeader(user_agent)) = user_agent {
+        user_agent.to_string()
+    } else {
+        String::from("Unknown browser")
+    };
+    tracing::debug!("`{user_agent}` at {addr} connected.");
+
+    Ok(ws.on_upgrade(move |socket| handle_socket(socket, addr, receiver)))
 }
 
 async fn handle_socket(
     mut socket: WebSocket,
+    who: SocketAddr,
     mut receiver: tokio::sync::broadcast::Receiver<crate::routes::messages::Message>,
 ) {
     let mut ping_interval = tokio::time::interval(Duration::from_secs(30));
@@ -126,7 +138,7 @@ async fn handle_socket(
                 if let Ok(msg) = msg {
                     if let Ok(json) = serde_json::to_string(&msg) {
                         if socket.send(WsMessage::Text(json.into())).await.is_err() {
-                            tracing::warn!("WebSocket send failed, closing socket");
+                            tracing::error!("Failed to send message to {who}, closing connection");
                             break;
                         }
                     }
@@ -138,33 +150,33 @@ async fn handle_socket(
             // Send Ping
             _ = ping_interval.tick() => {
                 if socket.send(WsMessage::Ping(vec![].into())).await.is_err() {
-                    tracing::error!("Failed to send ping, closing connection");
+                    tracing::error!("Failed to send ping to {who}, closing connection");
                     break;
                 }
 
-                tracing::info!("Ping sent successfully");
+                // tracing::debug!("Ping sent to {who}");
             }
 
             // Get incoming messages from client
             client_msg = socket.recv() => {
-                match client_msg {
+                if let Some(Ok(msg)) = client_msg {
-                    Some(Ok(msg)) => {
                     match msg {
-                            WsMessage::Pong(_) => {
+                        // WsMessage::Pong(_) => {
-                                tracing::info!("Received Pong");
+                        //     tracing::debug!("Received Pong from {who}");
-                            }
+                        // }
-                            WsMessage::Ping(_) => {
+                        // WsMessage::Ping(_) => {
-                                tracing::info!("Received Ping from client");
+                        //     tracing::info!("Received Ping from client");
-                            }
+                        // }
                         // WsMessage::Text(_) => {}
                         WsMessage::Close(_) => {
-                                tracing::info!("Client disconnected");
+                            tracing::debug!("Client disconnected");
                             break;
                         }
                         _ => {}
                     }
-                    },
+                } else {
-                    _ => break,
+                    tracing::debug!("Client {who} abruptly disconnected");
+                    break;
                 }
             }
         }