use std::time::Instant;
use anyhow::{Context, Result};
use ring::aead::{self, Aad, LessSafeKey, Nonce, UnboundKey};
use zstd::bulk::Compressor;
use crate::alloc_counter;
#[derive(Debug, Clone)]
pub struct ModuleAConfig {
pub payload_bytes: usize,
pub chunk_bytes: usize,
}
#[derive(Debug, Clone)]
pub struct ModuleAStats {
pub payload_bytes: usize,
pub chunk_bytes: usize,
pub chunks: usize,
pub elapsed_ms: f64,
pub throughput_mb_s: f64,
pub total_output_bytes: usize,
pub hot_loop_allocations: u64,
pub crc_fold: u32,
}
impl ModuleAStats {
pub fn to_json(&self) -> String {
format!(
"{{\"module\":\"A\",\"payload_bytes\":{},\"chunk_bytes\":{},\"chunks\":{},\"elapsed_ms\":{:.3},\"throughput_mb_s\":{:.3},\"total_output_bytes\":{},\"hot_loop_allocations\":{},\"crc_fold\":{}}}",
self.payload_bytes,
self.chunk_bytes,
self.chunks,
self.elapsed_ms,
self.throughput_mb_s,
self.total_output_bytes,
self.hot_loop_allocations,
self.crc_fold
)
}
}
pub fn run(config: ModuleAConfig) -> Result<ModuleAStats> {
if config.chunk_bytes == 0 {
anyhow::bail!("chunk_bytes must be > 0");
}
if config.payload_bytes == 0 {
anyhow::bail!("payload_bytes must be > 0");
}
if config.payload_bytes % config.chunk_bytes != 0 {
anyhow::bail!("payload_bytes must be divisible by chunk_bytes");
}
let chunks = config.payload_bytes / config.chunk_bytes;
let mut payload = vec![0_u8; config.payload_bytes];
fill_semi_compressible(&mut payload);
let chunk_output_bound = zstd::zstd_safe::compress_bound(config.chunk_bytes)
.saturating_add(aead::AES_256_GCM.tag_len());
let sink_capacity = chunk_output_bound
.checked_mul(chunks)
.context("sink capacity overflow")?;
let mut compressed_buffer = vec![0_u8; zstd::zstd_safe::compress_bound(config.chunk_bytes)];
let mut blackhole = vec![0_u8; sink_capacity];
let mut compressor = Compressor::new(1).context("failed to initialize zstd compressor")?;
let key_bytes = [0x42_u8; 32];
let unbound = UnboundKey::new(&aead::AES_256_GCM, &key_bytes)
.map_err(|_| anyhow::anyhow!("failed to initialize AES-256-GCM key"))?;
let key = LessSafeKey::new(unbound);
let warm_chunk = &payload[..config.chunk_bytes];
let warm_compressed = compressor
.compress_to_buffer(warm_chunk, &mut compressed_buffer[..])
.context("warm-up zstd compression failed")?;
let warm_nonce = nonce_for_chunk(0);
let warm_tag = key
.seal_in_place_separate_tag(
warm_nonce,
Aad::empty(),
&mut compressed_buffer[..warm_compressed],
)
.map_err(|_| anyhow::anyhow!("warm-up AES-256-GCM encryption failed"))?;
let warm_total = warm_compressed + warm_tag.as_ref().len();
blackhole[..warm_compressed].copy_from_slice(&compressed_buffer[..warm_compressed]);
blackhole[warm_compressed..warm_total].copy_from_slice(warm_tag.as_ref());
let alloc_before = alloc_counter::snapshot();
let start = Instant::now();
let mut sink_offset = 0usize;
let mut crc_fold = 0_u32;
for idx in 0..chunks {
let start_idx = idx * config.chunk_bytes;
let end_idx = start_idx + config.chunk_bytes;
let chunk = &payload[start_idx..end_idx];
let crc = crc32c::crc32c(chunk);
crc_fold ^= crc;
let compressed_len = compressor
.compress_to_buffer(chunk, &mut compressed_buffer[..])
.with_context(|| format!("zstd compression failed at chunk {}", idx))?;
let nonce = nonce_for_chunk(idx as u64 + 1);
let tag = key
.seal_in_place_separate_tag(
nonce,
Aad::from(crc.to_le_bytes()),
&mut compressed_buffer[..compressed_len],
)
.map_err(|_| anyhow::anyhow!("AES-256-GCM encryption failed at chunk {}", idx))?;
let encrypted_len = compressed_len + tag.as_ref().len();
let end_offset = sink_offset + encrypted_len;
if end_offset > blackhole.len() {
anyhow::bail!("blackhole buffer overflow");
}
blackhole[sink_offset..sink_offset + compressed_len]
.copy_from_slice(&compressed_buffer[..compressed_len]);
blackhole[sink_offset + compressed_len..end_offset].copy_from_slice(tag.as_ref());
sink_offset = end_offset;
}
let elapsed = start.elapsed();
let hot_loop_allocations = alloc_counter::allocations_since(alloc_before);
let elapsed_ms = elapsed.as_secs_f64() * 1_000.0;
let throughput_mb_s = (config.payload_bytes as f64 / (1024.0 * 1024.0)) / elapsed.as_secs_f64();
Ok(ModuleAStats {
payload_bytes: config.payload_bytes,
chunk_bytes: config.chunk_bytes,
chunks,
elapsed_ms,
throughput_mb_s,
total_output_bytes: sink_offset,
hot_loop_allocations,
crc_fold,
})
}
fn fill_semi_compressible(buffer: &mut [u8]) {
let pattern = b"TRACER_BULLET_PATTERN_";
for (idx, byte) in buffer.iter_mut().enumerate() {
if idx % 64 < pattern.len() {
*byte = pattern[idx % pattern.len()];
} else {
let mixed = (idx as u64)
.wrapping_mul(6364136223846793005_u64)
.wrapping_add(1);
*byte = (mixed >> 32) as u8;
}
}
}
fn nonce_for_chunk(chunk: u64) -> Nonce {
let mut nonce = [0_u8; 12];
nonce[4..].copy_from_slice(&chunk.to_be_bytes());
Nonce::assume_unique_for_key(nonce)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rejects_non_divisible_payload() {
let err = run(ModuleAConfig {
payload_bytes: 17,
chunk_bytes: 8,
})
.expect_err("expected configuration validation failure");
assert!(err.to_string().contains("divisible"));
}
#[test]
fn processes_small_payload() {
let stats = run(ModuleAConfig {
payload_bytes: 8 * 1024 * 1024,
chunk_bytes: 1024 * 1024,
})
.expect("module-a small payload run should succeed");
assert_eq!(stats.chunks, 8);
assert!(stats.total_output_bytes > 0);
assert!(stats.elapsed_ms >= 0.0);
assert!(stats.throughput_mb_s > 0.0);
}
#[test]
fn rejects_zero_chunk() {
let err = run(ModuleAConfig {
payload_bytes: 1024,
chunk_bytes: 0,
})
.expect_err("expected zero chunk rejection");
assert!(err.to_string().contains("chunk_bytes"));
}
#[test]
fn json_output_contains_module_label() {
let stats = run(ModuleAConfig {
payload_bytes: 1024 * 1024,
chunk_bytes: 256 * 1024,
})
.expect("module-a run should succeed");
let json = stats.to_json();
assert!(json.contains("\"module\":\"A\""));
assert!(json.contains("\"chunks\":4"));
}
}