multi session

This commit is contained in:
Richard
2026-01-28 21:38:32 +00:00
parent 6ec7b42b22
commit 1e35286eec
3 changed files with 323 additions and 55 deletions

182
Cargo.lock generated
View File

@@ -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"

View File

@@ -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"

View File

@@ -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<String>,
/// Bearer token for Authorization header
#[arg(long)]
bearer_token: Option<String>,
}
#[tokio::main]
@@ -22,67 +37,137 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
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(())
}