Mirror/rayhunter
1
0
Fork 0
mirror of https://github.com/EFForg/rayhunter.git synced 2026-07-21 23:18:11 -07:00
Code Issues Packages Projects Releases Wiki Activity

Refactor and simplify QmdlReader

Browse Source 
In the past, QmdlReader was written to share a trait with DiagDevice, so
it had to pretend to be reading MessagesContainers. This needlessly
complicated both its code as well as that of consumers'. Instead,
QmdlReader now returns a stream of diag Messages.

QmdlReader also automatically detects if it's reading a compressed QMDL
stream or not.

Additionally, QmdlReader no longer can be bounded by a filesize limit,
and instead relies on HDLC message framing to detect file truncation.
This works for both compressed and uncompressed QMDL files.
This commit is contained in:
Will Greenberg
2026-05-07 16:46:59 -07:00
parent af4a9aeb95
commit 0cd70ad73c
9 changed files with 368 additions and 270 deletions
Download Patch File Download Diff File Expand all files Collapse all files
daemon/src/analysis.rs
+19 -14
View File
@@ -1,15 +1,14 @@
use std::sync::Arc;
use std::{cmp, future, pin};
use std::cmp;
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::diag::{DiagParsingError, Message, MessagesContainer};
use serde::Serialize;
use tokio::fs::File;
use tokio::io::{AsyncWriteExt, BufWriter};
@@ -46,7 +45,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> {
@@ -61,6 +60,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');
@@ -134,7 +144,7 @@ async fn perform_analysis(
analyzer_config: &AnalyzerConfig,
) -> Result<(), String> {
info!("Opening QMDL and analysis file for {name}...");
let (analysis_file, qmdl_reader) = {
let (analysis_file, mut qmdl_reader) = {
let mut qmdl_store = qmdl_store_lock.write().await;
let (entry_index, _) = qmdl_store
.entry_for_name(name)
@@ -154,20 +164,15 @@ async fn perform_analysis(
let mut analysis_writer = AnalysisWriter::new(analysis_file, analyzer_config)
.await
.map_err(|e| format!("{e:?}"))?;
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:?}"))?;
}
daemon/src/diag.rs
+1 -1
View File
@@ -342,7 +342,7 @@ impl DiagTask {
debug!("done!");
let container_bytes: usize = container.messages.iter().map(|m| m.data.len()).sum();
self.bytes_since_space_check += container_bytes;
let max_type = match analysis_writer.analyze(container).await {
let max_type = match analysis_writer.analyze_container(container).await {
Ok(t) => t,
Err(e) => {
warn!("failed to analyze container: {e}");
daemon/src/pcap.rs
+11 -18
View File
@@ -7,12 +7,11 @@ use axum::http::StatusCode;
use axum::http::header::CONTENT_TYPE;
use axum::response::{IntoResponse, Response};
use log::error;
use rayhunter::diag::DataType;
use rayhunter::gsmtap_parser;
use rayhunter::pcap::GsmtapPcapWriter;
use rayhunter::qmdl::QmdlReader;
use std::sync::Arc;
use tokio::io::{AsyncRead, AsyncWrite, duplex};
use tokio::io::{AsyncRead, AsyncSeek, AsyncWrite, duplex};
use tokio_util::io::ReaderStream;
// Streams a pcap file chunk-by-chunk to the client by reading the QMDL data
@@ -71,28 +70,22 @@ pub async fn get_pcap(
pub async fn generate_pcap_data<R, W>(writer: W, mut reader: QmdlReader<R>) -> Result<(), Error>
where
W: AsyncWrite + Unpin + Send,
R: AsyncRead + Unpin,
R: AsyncRead + AsyncSeek + Unpin,
{
let mut pcap_writer = GsmtapPcapWriter::new(writer).await?;
pcap_writer.write_iface_header().await?;
while let Some(container) = reader.get_next_messages_container().await? {
if container.data_type != DataType::UserSpace {
continue;
}
for maybe_msg in container.into_messages() {
match maybe_msg {
Ok(msg) => {
let maybe_gsmtap_msg = gsmtap_parser::parse(msg)?;
if let Some((timestamp, gsmtap_msg)) = maybe_gsmtap_msg {
pcap_writer
.write_gsmtap_message(gsmtap_msg, timestamp)
.await?;
}
while let Some(maybe_msg) = reader.get_next_message().await? {
match maybe_msg {
Ok(msg) => {
let maybe_gsmtap_msg = gsmtap_parser::parse(msg)?;
if let Some((timestamp, gsmtap_msg)) = maybe_gsmtap_msg {
pcap_writer
.write_gsmtap_message(gsmtap_msg, timestamp)
.await?;
}
Err(e) => error!("error parsing message: {e:?}"),
}
Err(e) => error!("error parsing message: {e:?}"),
}
}
daemon/src/qmdl_store.rs
+3 -5
View File
@@ -288,11 +288,9 @@ impl RecordingStore {
let file = File::open(entry.get_qmdl_filepath(&self.path))
.await
.map_err(RecordingStoreError::ReadFileError)?;
Ok(QmdlReader::new(
file,
entry.compressed,
Some(entry.uncompressed_qmdl_size_bytes),
))
QmdlReader::new(file)
.await
.map_err(RecordingStoreError::ReadFileError)
}
// Returns the corresponding QMDL file for a given entry
daemon/src/server.rs
+89 -23
View File
@@ -82,7 +82,7 @@ pub async fn get_qmdl(
&entry.uncompressed_qmdl_size_bytes.to_string(),
),
];
let body = Body::from_stream(qmdl_reader.as_stream());
let body = Body::from_stream(qmdl_reader.as_qmdl_stream());
Ok((headers, body).into_response())
}
@@ -310,7 +310,7 @@ pub async fn get_zip(
Path(entry_name): Path<String>,
) -> Result<Response, (StatusCode, String)> {
let qmdl_idx = entry_name.trim_end_matches(".zip").to_owned();
let (entry_index, compressed) = {
let (entry_index, _) = {
let qmdl_store = state.qmdl_store_lock.read().await;
let (entry_index, entry) = qmdl_store.entry_for_name(&qmdl_idx).ok_or((
StatusCode::NOT_FOUND,
@@ -338,7 +338,7 @@ pub async fn get_zip(
// Add QMDL file
{
let entry = ZipEntryBuilder::new(
format!("{qmdl_idx}.qmdl").into(),
format!("{qmdl_idx}.qmdl.gz").into(),
Compression::Stored,
);
// FuturesAsyncWriteCompatExt::compat_write because async-zip's
@@ -425,9 +425,13 @@ pub async fn debug_set_display_state(
#[cfg(test)]
mod tests {
use std::io::Cursor;
use super::*;
use async_zip::base::read::mem::ZipFileReader;
use axum::extract::{Path, State};
use futures::AsyncReadExt;
use rayhunter::{diag::{DataType, HdlcEncapsulatedMessage, Message, MessagesContainer}, qmdl::{QmdlReader, QmdlWriter}};
use tempfile::TempDir;
async fn create_test_qmdl_store() -> (TempDir, Arc<RwLock<crate::qmdl_store::RecordingStore>>) {
@@ -441,24 +445,23 @@ mod tests {
async fn create_test_entry_with_data(
store_lock: &Arc<RwLock<crate::qmdl_store::RecordingStore>>,
test_data: &[u8],
test_data: &MessagesContainer,
) -> String {
let entry_name = {
let mut store = store_lock.write().await;
let (mut qmdl_file, _analysis_file) = store.new_entry().await.unwrap();
let (qmdl_gz_file, _analysis_file) = store.new_entry().await.unwrap();
if !test_data.is_empty() {
use tokio::io::AsyncWriteExt;
qmdl_file.write_all(test_data).await.unwrap();
qmdl_file.flush().await.unwrap();
}
let mut writer = QmdlWriter::new(qmdl_gz_file);
writer.write_container(test_data).await.unwrap();
let test_data_len = writer.total_uncompressed_bytes;
writer.close().await.unwrap();
let current_entry = store.current_entry.unwrap();
let entry = &store.manifest.entries[current_entry];
let entry_name = entry.name.clone();
store
.update_entry_qmdl_size(current_entry, test_data.len())
.update_entry_qmdl_size(current_entry, test_data_len)
.await
.unwrap();
entry_name
@@ -492,17 +495,69 @@ mod tests {
})
}
// valid HDLC encapsulated diag message generated from
// rayhunter::diag::test::get_test_message
fn create_test_container() -> MessagesContainer {
MessagesContainer {
data_type: DataType::UserSpace,
num_messages: 1,
messages: vec![
HdlcEncapsulatedMessage {
len: 39,
data: vec![
16,
0,
32,
0,
32,
0,
192,
176,
26,
165,
245,
135,
118,
35,
2,
1,
20,
14,
48,
0,
160,
0,
2,
8,
0,
0,
217,
15,
5,
0,
0,
0,
0,
1,
0,
10,
13,
196,
126,
],
},
],
}
}
#[tokio::test]
async fn test_get_zip_success() {
let (_temp_dir, store_lock) = create_test_qmdl_store().await;
let test_qmdl_data = vec![0x7E, 0x00, 0x00, 0x00, 0x10, 0x00, 0x7E];
let test_qmdl_data = create_test_container();
let entry_name = create_test_entry_with_data(&store_lock, &test_qmdl_data).await;
let state = create_test_server_state(store_lock);
let result = get_zip(State(state), Path(entry_name.clone())).await;
assert!(result.is_ok());
let response = result.unwrap();
let response = get_zip(State(state), Path(entry_name.clone())).await.unwrap();
let headers = response.headers();
assert_eq!(headers.get("content-type").unwrap(), "application/zip");
@@ -511,14 +566,11 @@ mod tests {
let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();
let zip_reader = ZipFileReader::new(body_bytes.to_vec()).await.unwrap();
let filenames = zip_reader
.file()
.entries()
let zip_reader_file = zip_reader.file();
let filenames: Vec<String> = zip_reader_file.entries()
.iter()
.map(|entry| entry.filename().as_str().unwrap().to_owned())
.collect::<Vec<String>>();
.map(|entry| entry.filename().as_str().unwrap().to_string())
.collect();
assert_eq!(
filenames,
vec![
@@ -526,5 +578,19 @@ mod tests {
format!("{entry_name}.pcapng"),
]
);
let mut qmdl_body = Vec::with_capacity(128);
zip_reader.reader_without_entry(0)
.await
.unwrap()
.read_to_end(&mut qmdl_body)
.await
.unwrap();
let mut qmdl_reader = QmdlReader::new(Cursor::new(qmdl_body)).await.unwrap();
let expected_message = Message::from_hdlc(&test_qmdl_data.messages[0].data).unwrap();
assert_eq!(
qmdl_reader.get_next_message().await.unwrap(),
Some(Ok(expected_message)),
);
}
}