Add support for compressed QMDL

check/src/main.rs

@@ -1,15 +1,13 @@
 use clap::Parser;
-use futures::TryStreamExt;
 use log::{debug, error, info, warn};
 use pcap_file_tokio::pcapng::{Block, PcapNgReader};
 use rayhunter::{
     analysis::analyzer::{AnalysisRow, AnalyzerConfig, EventType, Harness},
-    diag::DataType,
     gsmtap_parser,
     pcap::GsmtapPcapWriter,
-    qmdl::QmdlReader,
+    qmdl::QmdlMessageReader,
 };
-use std::{collections::HashMap, future, path::PathBuf, pin::pin};
+use std::{collections::HashMap, path::PathBuf};
 use tokio::fs::File;
 use walkdir::WalkDir;
 
@@ -113,26 +111,16 @@ async fn analyze_pcap(pcap_path: &str, show_skipped: bool) {
 async fn analyze_qmdl(qmdl_path: &str, show_skipped: bool) {
     let mut harness = Harness::new_with_config(&AnalyzerConfig::default());
     let qmdl_file = &mut File::open(&qmdl_path).await.expect("failed to open file");
-    let file_size = qmdl_file
-        .metadata()
+    let mut qmdl_reader = QmdlMessageReader::new(qmdl_file)
         .await
-        .expect("failed to get QMDL file metadata")
-        .len();
-    let mut qmdl_reader = QmdlReader::new(qmdl_file, Some(file_size as usize));
-    let mut qmdl_stream = pin!(
-        qmdl_reader
-            .as_stream()
-            .try_filter(|container| future::ready(container.data_type == DataType::UserSpace))
-    );
+        .expect("failed to open QmdlReader");
     let mut report = Report::new(qmdl_path);
-    while let Some(container) = qmdl_stream
-        .try_next()
+    while let Some(maybe_message) = qmdl_reader
+        .get_next_message()
         .await
-        .expect("failed getting QMDL container")
+        .expect("failed to get message")
     {
-        for row in harness.analyze_qmdl_messages(container) {
-            report.process_row(row);
-        }
+        report.process_row(harness.analyze_qmdl_message(maybe_message));
     }
     report.print_summary(show_skipped);
 }
@@ -141,27 +129,28 @@ async fn pcapify(qmdl_path: &PathBuf) {
     let qmdl_file = &mut File::open(&qmdl_path)
         .await
         .expect("failed to open qmdl file");
-    let qmdl_file_size = qmdl_file.metadata().await.unwrap().len();
-    let mut qmdl_reader = QmdlReader::new(qmdl_file, Some(qmdl_file_size as usize));
+    let mut qmdl_reader = QmdlMessageReader::new(qmdl_file)
+        .await
+        .expect("failed to open QmdlReader");
     let mut pcap_path = qmdl_path.clone();
     pcap_path.set_extension("pcapng");
     let pcap_file = &mut File::create(&pcap_path)
         .await
         .expect("failed to open pcap file");
     let mut pcap_writer = GsmtapPcapWriter::new(pcap_file).await.unwrap();
     pcap_writer.write_iface_header().await.unwrap();
-    while let Some(container) = qmdl_reader
-        .get_next_messages_container()
+    while let Some(maybe_message) = qmdl_reader
+        .get_next_message()
         .await
-        .expect("failed to get container")
+        .expect("failed to get message")
     {
-        for msg in container.messages().into_iter().flatten() {
-            if let Ok(Some((timestamp, parsed))) = gsmtap_parser::parse(msg) {
-                pcap_writer
-                    .write_gsmtap_message(parsed, timestamp, None)
-                    .await
-                    .expect("failed to write");
-            }
+        if let Ok(msg) = maybe_message
+            && let Ok(Some((timestamp, parsed))) = gsmtap_parser::parse(msg)
+        {
+            pcap_writer
+                .write_gsmtap_message(parsed, timestamp, None)
+                .await
+                .expect("failed to write");
         }
     }
     info!("wrote pcap to {:?}", &pcap_path);
@@ -197,9 +186,7 @@ async fn main() {
         let name_str = name.to_str().unwrap();
         let path = entry.path();
         let path_str = path.to_str().unwrap();
-        // instead of relying on the QMDL extension, can we check if a file is
-        // QMDL by inspecting the contents?
-        if name_str.ends_with(".qmdl") {
+        if name_str.ends_with(".qmdl") || name_str.ends_with(".qmdl.gz") {
             info!("**** Beginning analysis of {name_str}");
             analyze_qmdl(path_str, args.show_skipped).await;
             if args.pcapify {

daemon/Cargo.toml

@@ -35,7 +35,7 @@ futures-macro = "0.3.30"
 include_dir = "0.7.3"
 chrono = { version = "0.4.31", features = ["serde"] }
 tokio-stream = { version = "0.1.14", default-features = false, features = ["io-util"] }
-futures = { version = "0.3.30", default-features = false }
+futures = { version = "0.3.32", default-features = false, features = ["std"] }
 serde_json = "1.0.114"
 image = { version =  "0.25.1", default-features = false, features = ["png", "gif"] }
 tempfile = "3.10.2"

daemon/src/analysis.rs

@@ -1,16 +1,15 @@
+use std::cmp;
 use std::sync::Arc;
-use std::{cmp, future, pin};
 
 use axum::Json;
 use axum::{
     extract::{Path, State},
     http::StatusCode,
 };
-use futures::TryStreamExt;
 use log::{error, info};
 use rayhunter::analysis::analyzer::{AnalyzerConfig, EventType, Harness};
-use rayhunter::diag::{DataType, MessagesContainer};
-use rayhunter::qmdl::QmdlReader;
+use rayhunter::diag::{DiagParsingError, Message, MessagesContainer};
+use rayhunter::qmdl::QmdlMessageReader;
 use serde::Serialize;
 use tokio::fs::File;
 use tokio::io::{AsyncWriteExt, BufWriter};
@@ -47,7 +46,7 @@ impl AnalysisWriter {
 
     // Runs the analysis harness on the given container, serializing the results
     // to the analysis file, returning the whether any warnings were detected
-    pub async fn analyze(
+    pub async fn analyze_container(
         &mut self,
         container: MessagesContainer,
     ) -> Result<EventType, std::io::Error> {
@@ -62,6 +61,17 @@ impl AnalysisWriter {
         Ok(max_type)
     }
 
+    pub async fn analyze_message(
+        &mut self,
+        maybe_qmdl_msg: Result<Message, DiagParsingError>,
+    ) -> Result<EventType, std::io::Error> {
+        let row = self.harness.analyze_qmdl_message(maybe_qmdl_msg);
+        if !row.is_empty() {
+            self.write(&row).await?;
+        }
+        Ok(row.get_max_event_type())
+    }
+
     async fn write<T: Serialize>(&mut self, value: &T) -> Result<(), std::io::Error> {
         let mut value_str = serde_json::to_string(value).unwrap();
         value_str.push('\n');
@@ -135,7 +145,7 @@ async fn perform_analysis(
     analyzer_config: &AnalyzerConfig,
 ) -> Result<(), String> {
     info!("Opening QMDL and analysis file for {name}...");
-    let (analysis_file, qmdl_file) = {
+    let (analysis_file, mut qmdl_reader) = {
         let mut qmdl_store = qmdl_store_lock.write().await;
         let (entry_index, _) = qmdl_store
             .entry_for_name(name)
@@ -149,33 +159,25 @@ async fn perform_analysis(
             .await
             .map_err(|e| format!("{e:?}"))?
             .ok_or("QMDL file not found")?;
+        let qmdl_reader = QmdlMessageReader::new(qmdl_file)
+            .await
+            .map_err(|e| format!("{e:?}"))?;
 
-        (analysis_file, qmdl_file)
+        (analysis_file, qmdl_reader)
     };
 
     let mut analysis_writer = AnalysisWriter::new(analysis_file, analyzer_config)
         .await
         .map_err(|e| format!("{e:?}"))?;
-    let file_size = qmdl_file
-        .metadata()
-        .await
-        .expect("failed to get QMDL file metadata")
-        .len();
-    let mut qmdl_reader = QmdlReader::new(qmdl_file, Some(file_size as usize));
-    let mut qmdl_stream = pin::pin!(
-        qmdl_reader
-            .as_stream()
-            .try_filter(|container| future::ready(container.data_type == DataType::UserSpace))
-    );
 
     info!("Starting analysis for {name}...");
-    while let Some(container) = qmdl_stream
-        .try_next()
+    while let Some(maybe_message) = qmdl_reader
+        .get_next_message()
         .await
-        .expect("failed getting QMDL container")
+        .expect("failed to get message")
     {
         let _ = analysis_writer
-            .analyze(container)
+            .analyze_message(maybe_message)
             .await
             .map_err(|e| format!("{e:?}"))?;
     }