diff --git a/Cargo.lock b/Cargo.lock index 80765a9..d29868d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1253,6 +1253,7 @@ dependencies = [ "serde", "thiserror", "tokio", + "tokio-util", "toml", ] diff --git a/wavehunter/Cargo.toml b/wavehunter/Cargo.toml index 914a8df..067e353 100644 --- a/wavehunter/Cargo.toml +++ b/wavehunter/Cargo.toml @@ -15,3 +15,4 @@ futures-core = "0.3.30" thiserror = "1.0.52" log = "0.4.20" env_logger = "0.10.1" +tokio-util = { version = "0.7.10", features = ["io"] } diff --git a/wavehunter/src/main.rs b/wavehunter/src/main.rs index f9cd13f..0b74fa8 100644 --- a/wavehunter/src/main.rs +++ b/wavehunter/src/main.rs @@ -3,7 +3,7 @@ mod error; mod server; use crate::config::{parse_config, parse_args}; -use crate::server::{ServerState, serve_pcap}; +use crate::server::{ServerState, serve_pcap, serve_qmdl}; use crate::error::WavehunterError; use log::debug; @@ -46,6 +46,7 @@ async fn run_server(config: &config::Config, qmdl_bytes_written: Arc>) -> Result { - error!("error parsing message: {:?}", e); - }, + Err(e) => error!("error parsing message: {:?}", e), } } }, @@ -110,6 +111,18 @@ impl Stream for ChannelReader { } } +pub async fn serve_qmdl(State(state): State>) -> Result { + let qmdl_file = AsyncFile::open(&state.qmdl_path).await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("error opening QMDL file: {}", e)))?; + let qmdl_bytes_written = *state.qmdl_bytes_written.read().await; + let limited_qmdl_file = qmdl_file.take(qmdl_bytes_written as u64); + let qmdl_stream = ReaderStream::new(limited_qmdl_file); + + let headers = [(CONTENT_TYPE, "application/octet-stream")]; + let body = Body::from_stream(qmdl_stream); + Ok((headers, body).into_response()) +} + pub struct ServerState { pub qmdl_bytes_written: Arc>, pub qmdl_path: String,