From 1e35286eec71b391c51d68ac241c8e680eec3b54 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 28 Jan 2026 21:38:32 +0000 Subject: [PATCH] multi session --- Cargo.lock | 182 ++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 1 + src/main.rs | 195 +++++++++++++++++++++++++++++++++++++--------------- 3 files changed, 323 insertions(+), 55 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 038e531..d2bb106 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anstream" version = "0.6.21" @@ -52,6 +61,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + [[package]] name = "bitflags" version = "2.10.0" @@ -67,6 +82,12 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bumpalo" +version = "3.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" + [[package]] name = "byteorder" version = "1.5.0" @@ -95,6 +116,19 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chrono" +version = "0.4.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fac4744fb15ae8337dc853fee7fb3f4e48c0fbaa23d0afe49c447b4fab126118" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "wasm-bindgen", + "windows-link", +] + [[package]] name = "clap" version = "4.5.55" @@ -348,6 +382,30 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.1.1" @@ -462,6 +520,16 @@ version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" +[[package]] +name = "js-sys" +version = "0.3.85" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c942ebf8e95485ca0d52d97da7c5a2c387d0e7f0ba4c35e93bfcaee045955b3" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -538,6 +606,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -735,6 +812,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + [[package]] name = "schannel" version = "0.1.28" @@ -1150,10 +1233,56 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasm-bindgen" +version = "0.2.108" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64024a30ec1e37399cf85a7ffefebdb72205ca1c972291c51512360d90bd8566" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.108" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "008b239d9c740232e71bd39e8ef6429d27097518b6b30bdf9086833bd5b6d608" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.108" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5256bae2d58f54820e6490f9839c49780dff84c65aeab9e772f15d5f0e913a55" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.108" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f01b580c9ac74c8d8f0c0e4afb04eeef2acf145458e52c03845ee9cd23e3d12" +dependencies = [ + "unicode-ident", +] + [[package]] name = "websocket-debug" version = "0.1.0" dependencies = [ + "chrono", "clap", "futures-util", "tokio", @@ -1162,12 +1291,65 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.60.2" diff --git a/Cargo.toml b/Cargo.toml index d72bdfb..68b2051 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,3 +10,4 @@ tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["fmt"] } clap = { version = "4", features = ["derive"] } futures-util = "0.3" +chrono = "0.4" diff --git a/src/main.rs b/src/main.rs index df9a820..997bd43 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,16 +1,31 @@ use std::fs; +use std::path::PathBuf; +use chrono::Local; use clap::Parser; use futures_util::StreamExt; -use tokio_tungstenite::{connect_async, tungstenite::Message}; +use tokio::task::JoinSet; +use tokio_tungstenite::{ + connect_async_with_config, + tungstenite::{ + client::IntoClientRequest, + http::header::{AUTHORIZATION, HeaderValue}, + Message, + }, +}; use tracing::{error, info, warn}; #[derive(Parser)] #[command(name = "websocket-debug")] #[command(about = "A WebSocket debugging tool that logs and saves messages")] struct Args { - /// WebSocket URL to connect to (e.g., ws://localhost:8080 or wss://example.com/ws) - url: String, + /// WebSocket URLs to connect to (e.g., ws://localhost:8080 or wss://example.com/ws) + #[arg(required = true)] + urls: Vec, + + /// Bearer token for Authorization header + #[arg(long)] + bearer_token: Option, } #[tokio::main] @@ -22,67 +37,137 @@ async fn main() -> Result<(), Box> { let args = Args::parse(); - info!("Connecting to {}", args.url); + // Create session directory + let timestamp = Local::now().format("%Y%m%d-%H%M%S"); + let session_dir = PathBuf::from(format!("session-{}", timestamp)); + fs::create_dir_all(&session_dir)?; + info!("Created session directory: {}", session_dir.display()); - let (ws_stream, response) = connect_async(&args.url).await?; - info!("Connected successfully"); - info!("Response status: {}", response.status()); + // Connect to all URLs simultaneously + let mut connect_futures = Vec::new(); + for (idx, url) in args.urls.iter().enumerate() { + let letter = (b'A' + idx as u8) as char; + let bearer_token = args.bearer_token.clone(); + let url = url.clone(); - let (_, mut read) = ws_stream.split(); + connect_futures.push(async move { + let mut request = url + .as_str() + .into_client_request() + .map_err(|e| format!("[{}] Invalid URL {}: {}", letter, url, e))?; - let mut seq_num: u64 = 0; - - while let Some(message_result) = read.next().await { - match message_result { - Ok(message) => { - match message { - Message::Text(text) => { - let preview: String = text.chars().take(50).collect(); - let truncated = if text.len() > 50 { "..." } else { "" }; - info!("[{}] Text: {}{}", seq_num, preview, truncated); - - let filename = format!("{}.txt", seq_num); - if let Err(e) = fs::write(&filename, &text) { - error!("Failed to write {}: {}", filename, e); - } - } - Message::Binary(data) => { - info!("[{}] Binary: {} bytes", seq_num, data.len()); - - let filename = format!("{}.bin", seq_num); - if let Err(e) = fs::write(&filename, &data) { - error!("Failed to write {}: {}", filename, e); - } - } - Message::Ping(data) => { - info!("[{}] Ping: {} bytes", seq_num, data.len()); - continue; // Don't increment seq_num for control frames - } - Message::Pong(data) => { - info!("[{}] Pong: {} bytes", seq_num, data.len()); - continue; // Don't increment seq_num for control frames - } - Message::Close(frame) => { - if let Some(cf) = frame { - info!("Connection closed: {} - {}", cf.code, cf.reason); - } else { - info!("Connection closed"); - } - break; - } - Message::Frame(_) => { - continue; // Raw frames, skip - } - } - seq_num += 1; + if let Some(ref token) = bearer_token { + let auth_value = HeaderValue::from_str(&format!("Bearer {}", token)) + .map_err(|e| format!("Invalid bearer token: {}", e))?; + request.headers_mut().insert(AUTHORIZATION, auth_value); } + + info!("[{}] Connecting to {}", letter, url); + + let (ws_stream, response) = connect_async_with_config(request, None, false) + .await + .map_err(|e| format!("[{}] Failed to connect to {}: {}", letter, url, e))?; + + info!("[{}] Connected successfully (status: {})", letter, response.status()); + + Ok::<_, String>((letter, url, ws_stream)) + }); + } + + // Wait for all connections to establish + let results = futures_util::future::join_all(connect_futures).await; + + let mut connections = Vec::new(); + for result in results { + match result { + Ok(conn) => connections.push(conn), Err(e) => { - warn!("Error receiving message: {}", e); + error!("{}", e); + return Err(e.into()); } } } - info!("WebSocket connection ended. Received {} messages.", seq_num); + info!("All {} connections established", connections.len()); + + // Spawn tasks for each connection + let mut join_set: JoinSet<(char, String)> = JoinSet::new(); + + for (letter, url, ws_stream) in connections { + let session_dir = session_dir.clone(); + + join_set.spawn(async move { + let (_, mut read) = ws_stream.split(); + let mut seq_num: u64 = 0; + + while let Some(message_result) = read.next().await { + match message_result { + Ok(message) => { + match message { + Message::Text(text) => { + let preview: String = text.chars().take(50).collect(); + let truncated = if text.len() > 50 { "..." } else { "" }; + info!("[{}:{}] Text: {}{}", letter, seq_num, preview, truncated); + + let filename = session_dir.join(format!("{}{}.txt", letter, seq_num)); + if let Err(e) = fs::write(&filename, &text) { + error!("[{}] Failed to write {:?}: {}", letter, filename, e); + } + } + Message::Binary(data) => { + info!("[{}:{}] Binary: {} bytes", letter, seq_num, data.len()); + + let filename = session_dir.join(format!("{}{}.bin", letter, seq_num)); + if let Err(e) = fs::write(&filename, &data) { + error!("[{}] Failed to write {:?}: {}", letter, filename, e); + } + } + Message::Ping(data) => { + info!("[{}] Ping: {} bytes", letter, data.len()); + continue; + } + Message::Pong(data) => { + info!("[{}] Pong: {} bytes", letter, data.len()); + continue; + } + Message::Close(frame) => { + if let Some(cf) = frame { + info!("[{}] Connection closed: {} - {}", letter, cf.code, cf.reason); + } else { + info!("[{}] Connection closed", letter); + } + break; + } + Message::Frame(_) => { + continue; + } + } + seq_num += 1; + } + Err(e) => { + warn!("[{}] Error receiving message: {}", letter, e); + } + } + } + + info!("[{}] Session ended. Received {} messages.", letter, seq_num); + (letter, url) + }); + } + + // Wait for all connections to finish + while let Some(result) = join_set.join_next().await { + match result { + Ok((letter, url)) => { + info!("[{}] Disconnected from {}", letter, url); + } + Err(e) => { + error!("Task error: {}", e); + } + } + } + + info!("All connections closed. Session saved to: {}", session_dir.display()); Ok(()) }