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

Add support for compressed QMDL

Browse Source 
This reworks the QmdlWriter to output gzipped QMDL files by default,
and allows QmdlReader to operate on either compressed or uncompressed
QMDLs.

QmdlReader has been significantly rewritten to expose a single AsyncRead
interface to both compressed and uncompressed QMDL sources.
This commit is contained in:
Will Greenberg
2026-03-30 15:56:03 -07:00
parent 0580a8af33
commit 0a93e93838
11 changed files with 314 additions and 167 deletions
Download Patch File Download Diff File Expand all files Collapse all files
daemon/Cargo.toml
+1 -1
View File
@@ -33,7 +33,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
+3 -10
View File
@@ -10,7 +10,6 @@ use futures::TryStreamExt;
use log::{error, info};
use rayhunter::analysis::analyzer::{AnalyzerConfig, EventType, Harness};
use rayhunter::diag::{DataType, MessagesContainer};
use rayhunter::qmdl::QmdlReader;
use serde::Serialize;
use tokio::fs::File;
use tokio::io::{AsyncWriteExt, BufWriter};
@@ -135,7 +134,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, qmdl_reader) = {
let mut qmdl_store = qmdl_store_lock.write().await;
let (entry_index, _) = qmdl_store
.entry_for_name(name)
@@ -144,23 +143,17 @@ async fn perform_analysis(
.clear_and_open_entry_analysis(entry_index)
.await
.map_err(|e| format!("{e:?}"))?;
let qmdl_file = qmdl_store
let qmdl_reader = qmdl_store
.open_entry_qmdl(entry_index)
.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()
daemon/src/diag.rs
+19 -9
View File
@@ -143,7 +143,7 @@ impl DiagTask {
DiskSpaceCheck::Failed => {}
}
let (qmdl_file, analysis_file) = match qmdl_store.new_entry().await {
let (qmdl_gz_file, analysis_file) = match qmdl_store.new_entry().await {
Ok(files) => files,
Err(e) => {
let msg = format!("failed creating QMDL file entry: {e}");
@@ -152,7 +152,7 @@ impl DiagTask {
}
};
self.stop_current_recording().await;
let qmdl_writer = QmdlWriter::new(qmdl_file);
let qmdl_writer = QmdlWriter::new(qmdl_gz_file);
let analysis_writer = match AnalysisWriter::new(analysis_file, &self.analyzer_config).await
{
Ok(writer) => Box::new(writer),
@@ -237,13 +237,23 @@ impl DiagTask {
let mut state = DiagState::Stopped;
std::mem::swap(&mut self.state, &mut state);
if let DiagState::Recording {
analysis_writer, ..
qmdl_writer,
analysis_writer,
..
} = state
{
analysis_writer
.close()
.await
.expect("failed to close analysis writer");
match (qmdl_writer.close().await, analysis_writer.close().await) {
(Ok(()), Ok(())) => {},
(qmdl_result, analysis_result) => {
if let Err(err) = qmdl_result {
error!("failed to close QmdlWriter: {:?}", err);
}
if let Err(err) = analysis_result {
error!("failed to close AnalysisWriter: {:?}", err);
}
panic!();
}
}
}
}
@@ -315,13 +325,13 @@ impl DiagTask {
}
debug!(
"total QMDL bytes written: {}, updating manifest...",
qmdl_writer.total_written
qmdl_writer.total_uncompressed_bytes
);
let index = qmdl_store
.current_entry
.expect("DiagDevice had qmdl_writer, but QmdlStore didn't have current entry???");
if let Err(e) = qmdl_store
.update_entry_qmdl_size(index, qmdl_writer.total_written)
.update_entry_qmdl_size(index, qmdl_writer.total_uncompressed_bytes)
.await
{
let reason = format!("failed to update manifest (disk full?): {e}");
daemon/src/pcap.rs
+4 -9
View File
@@ -45,23 +45,20 @@ pub async fn get_pcap(
StatusCode::NOT_FOUND,
format!("couldn't find manifest entry with name {qmdl_name}"),
))?;
if entry.qmdl_size_bytes == 0 {
if entry.uncompressed_qmdl_size_bytes == 0 {
return Err((
StatusCode::SERVICE_UNAVAILABLE,
"QMDL file is empty, try again in a bit!".to_string(),
));
}
let qmdl_size_bytes = entry.qmdl_size_bytes;
let qmdl_file = qmdl_store
let qmdl_reader = qmdl_store
.open_entry_qmdl(entry_index)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("{e:?}")))?;
// the QMDL reader should stop at the last successfully written data chunk
// (entry.size_bytes)
let (reader, writer) = duplex(1024);
tokio::spawn(async move {
if let Err(e) = generate_pcap_data(writer, qmdl_file, qmdl_size_bytes).await {
if let Err(e) = generate_pcap_data(writer, qmdl_reader).await {
error!("failed to generate PCAP: {e:?}");
}
});
@@ -73,8 +70,7 @@ pub async fn get_pcap(
pub async fn generate_pcap_data<R, W>(
writer: W,
qmdl_file: R,
qmdl_size_bytes: usize,
mut reader: QmdlReader<R>,
) -> Result<(), Error>
where
W: AsyncWrite + Unpin + Send,
@@ -83,7 +79,6 @@ where
let mut pcap_writer = GsmtapPcapWriter::new(writer).await?;
pcap_writer.write_iface_header().await?;
let mut reader = QmdlReader::new(qmdl_file, Some(qmdl_size_bytes));
while let Some(container) = reader.get_next_messages_container().await? {
if container.data_type != DataType::UserSpace {
continue;
daemon/src/qmdl_store.rs
+28 -14
View File
@@ -4,6 +4,7 @@ use std::path::{Path, PathBuf};
use chrono::{DateTime, Local};
use log::{info, warn};
use rayhunter::qmdl::QmdlReader;
use rayhunter::util::RuntimeMetadata;
use serde::{Deserialize, Serialize};
use thiserror::Error;
@@ -58,7 +59,7 @@ pub struct ManifestEntry {
#[cfg_attr(feature = "apidocs", schema(value_type = String))]
pub last_message_time: Option<DateTime<Local>>,
/// The size of the QMDL file in bytes
pub qmdl_size_bytes: usize,
pub uncompressed_qmdl_size_bytes: usize,
/// The rayhunter daemon version which generated the file
pub rayhunter_version: Option<String>,
/// The OS which created the file
@@ -67,6 +68,8 @@ pub struct ManifestEntry {
pub arch: Option<String>,
#[serde(default)]
pub stop_reason: Option<String>,
#[serde(default)]
pub compressed: bool,
}
impl ManifestEntry {
@@ -77,17 +80,22 @@ impl ManifestEntry {
name: format!("{}", now.timestamp()),
start_time: now,
last_message_time: None,
qmdl_size_bytes: 0,
uncompressed_qmdl_size_bytes: 0,
rayhunter_version: Some(metadata.rayhunter_version),
system_os: Some(metadata.system_os),
arch: Some(metadata.arch),
stop_reason: None,
compressed: true,
}
}
pub fn get_qmdl_filepath<P: AsRef<Path>>(&self, path: P) -> PathBuf {
let mut filepath = path.as_ref().join(&self.name);
filepath.set_extension("qmdl");
if self.compressed {
filepath.set_extension("qmdl.gz");
} else {
filepath.set_extension("qmdl");
}
filepath
}
@@ -153,8 +161,9 @@ impl RecordingStore {
}
// Does a best-effort attempt to recover the manifest from a directory of
// QMDL files. We expect these files to be named like "<timestamp>.qmdl",
// and skip any files which don't match that pattern.
// QMDL files. We expect these files to be named like "<timestamp>.qmdl"
// or "<timestamp>.qmdl.gz", and skip any files which don't match that
// pattern.
pub async fn recover<P>(path: P) -> Result<Self, RecordingStoreError>
where
P: AsRef<Path>,
@@ -174,11 +183,14 @@ impl RecordingStore {
continue;
};
if !filename.ends_with(".qmdl") {
let (stem, compressed) = if filename.ends_with(".qmdl") {
(filename.trim_end_matches(".qmdl"), false)
} else if filename.ends_with(".qmdl.gz") {
(filename.trim_end_matches(".qmdl.gz"), true)
} else {
continue;
}
};
let stem = filename.trim_end_matches(".qmdl");
let Ok(start_timestamp) = stem.parse::<i64>() else {
warn!("QMDL file has invalid name {os_filename:?}, skipping");
continue;
@@ -205,9 +217,10 @@ impl RecordingStore {
info!("successfully recovered QMDL entry {os_filename:?}!");
manifest_entries.push(ManifestEntry {
name: stem.to_string(),
compressed,
start_time: start_time.into(),
last_message_time: Some(last_message_time.into()),
qmdl_size_bytes: metadata.size() as usize,
uncompressed_qmdl_size_bytes: metadata.size() as usize,
rayhunter_version: None,
system_os: None,
arch: None,
@@ -265,11 +278,12 @@ impl RecordingStore {
}
// Returns the corresponding QMDL file for a given entry
pub async fn open_entry_qmdl(&self, entry_index: usize) -> Result<File, RecordingStoreError> {
pub async fn open_entry_qmdl(&self, entry_index: usize) -> Result<QmdlReader<File>, RecordingStoreError> {
let entry = &self.manifest.entries[entry_index];
File::open(entry.get_qmdl_filepath(&self.path))
let file = File::open(entry.get_qmdl_filepath(&self.path))
.await
.map_err(RecordingStoreError::ReadFileError)
.map_err(RecordingStoreError::ReadFileError)?;
Ok(QmdlReader::new(file, entry.compressed, Some(entry.uncompressed_qmdl_size_bytes)))
}
// Returns the corresponding QMDL file for a given entry
@@ -314,7 +328,7 @@ impl RecordingStore {
entry_index: usize,
size_bytes: usize,
) -> Result<(), RecordingStoreError> {
self.manifest.entries[entry_index].qmdl_size_bytes = size_bytes;
self.manifest.entries[entry_index].uncompressed_qmdl_size_bytes = size_bytes;
self.manifest.entries[entry_index].last_message_time =
Some(rayhunter::clock::get_adjusted_now());
self.write_manifest().await
@@ -490,7 +504,7 @@ mod tests {
.entry_for_name(&store.manifest.entries[entry_index].name)
.unwrap();
assert!(entry.last_message_time.is_some());
assert_eq!(store.manifest.entries[entry_index].qmdl_size_bytes, 1000);
assert_eq!(store.manifest.entries[entry_index].uncompressed_qmdl_size_bytes, 1000);
assert_eq!(
RecordingStore::read_manifest(dir.path()).await.unwrap(),
store.manifest
daemon/src/server.rs
+20 -30
View File
@@ -12,9 +12,10 @@ use axum::response::{IntoResponse, Response};
use chrono::{DateTime, Local};
use log::{error, warn};
use serde::{Deserialize, Serialize};
use tokio::io::copy;
use std::sync::Arc;
use tokio::fs::write;
use tokio::io::{AsyncReadExt, copy, duplex};
use tokio::io::duplex;
use tokio::sync::RwLock;
use tokio::sync::mpsc::Sender;
use tokio_util::compat::FuturesAsyncWriteCompatExt;
@@ -64,7 +65,7 @@ pub async fn get_qmdl(
StatusCode::NOT_FOUND,
format!("couldn't find qmdl file with name {qmdl_idx}"),
))?;
let qmdl_file = qmdl_store
let qmdl_reader = qmdl_store
.open_entry_qmdl(entry_index)
.await
.map_err(|err| {
@@ -73,14 +74,12 @@ pub async fn get_qmdl(
format!("error opening QMDL file: {err}"),
)
})?;
let limited_qmdl_file = qmdl_file.take(entry.qmdl_size_bytes as u64);
let qmdl_stream = ReaderStream::new(limited_qmdl_file);
let headers = [
(CONTENT_TYPE, "application/octet-stream"),
(CONTENT_LENGTH, &entry.qmdl_size_bytes.to_string()),
(CONTENT_LENGTH, &entry.uncompressed_qmdl_size_bytes.to_string()),
];
let body = Body::from_stream(qmdl_stream);
let body = Body::from_stream(qmdl_reader.as_stream());
Ok((headers, body).into_response())
}
@@ -308,21 +307,21 @@ 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, qmdl_size_bytes) = {
let (entry_index, compressed) = {
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,
format!("couldn't find entry with name {qmdl_idx}"),
))?;
if entry.qmdl_size_bytes == 0 {
if entry.uncompressed_qmdl_size_bytes == 0 {
return Err((
StatusCode::SERVICE_UNAVAILABLE,
"QMDL file is empty, try again in a bit!".to_string(),
));
}
(entry_index, entry.qmdl_size_bytes)
(entry_index, entry.compressed)
};
let qmdl_store_lock = state.qmdl_store_lock.clone();
@@ -335,22 +334,18 @@ pub async fn get_zip(
// Add QMDL file
{
let extension = if compressed { "qmdl.gz" } else { "qmdl" };
let entry =
ZipEntryBuilder::new(format!("{qmdl_idx}.qmdl").into(), Compression::Stored);
ZipEntryBuilder::new(format!("{qmdl_idx}.{extension}").into(), Compression::Stored);
// FuturesAsyncWriteCompatExt::compat_write because async-zip's entrystream does
// not impl tokio's AsyncWrite, but only future's AsyncWrite. This can be removed
// once https://github.com/Majored/rs-async-zip/pull/160 is released.
let mut entry_writer = zip.write_entry_stream(entry).await?.compat_write();
let mut qmdl_file = {
let qmdl_store = qmdl_store_lock.read().await;
qmdl_store
.open_entry_qmdl(entry_index)
.await?
.take(qmdl_size_bytes as u64)
};
copy(&mut qmdl_file, &mut entry_writer).await?;
let qmdl_store = qmdl_store_lock.read().await;
let mut qmdl_reader = qmdl_store
.open_entry_qmdl(entry_index)
.await?;
copy(&mut qmdl_reader, &mut entry_writer).await?;
entry_writer.into_inner().close().await?;
}
@@ -360,17 +355,12 @@ pub async fn get_zip(
ZipEntryBuilder::new(format!("{qmdl_idx}.pcapng").into(), Compression::Stored);
let mut entry_writer = zip.write_entry_stream(entry).await?.compat_write();
let qmdl_file_for_pcap = {
let qmdl_store = qmdl_store_lock.read().await;
qmdl_store
.open_entry_qmdl(entry_index)
.await?
.take(qmdl_size_bytes as u64)
};
let qmdl_store = qmdl_store_lock.read().await;
let qmdl_reader = qmdl_store
.open_entry_qmdl(entry_index)
.await?;
if let Err(e) =
generate_pcap_data(&mut entry_writer, qmdl_file_for_pcap, qmdl_size_bytes).await
{
if let Err(e) = generate_pcap_data(&mut entry_writer, qmdl_reader).await {
// if we fail to generate the PCAP file, we should still continue and give the
// user the QMDL.
error!("Failed to generate PCAP: {e:?}");