use anyhow::{Context, Result};
use ring::aead::{self, Aad, LessSafeKey, Nonce, UnboundKey};
use zstd::bulk::Compressor;
use crate::topology::CpuGeneration;
pub const CHUNK_2_MB: usize = 2 * 1024 * 1024;
pub const CHUNK_4_MB: usize = 4 * 1024 * 1024;
pub const CHUNK_8_MB: usize = 8 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChunkPolicy {
Fixed2M,
Fixed4M,
Fixed8M,
Adaptive8To4,
}
#[derive(Debug, Clone, Copy)]
pub struct RuntimeMetrics {
pub put_p99_ms: f64,
pub median_queue_depth: f64,
}
#[derive(Debug, Clone)]
pub struct ChunkController {
policy: ChunkPolicy,
current_chunk_bytes: usize,
downshift_streak: u32,
upshift_streak: u32,
}
#[derive(Debug, Clone)]
pub struct EncryptedChunk {
pub nonce_id: u64,
pub crc32c: u32,
pub plain_len: usize,
pub cipher_len: usize,
pub bytes: Vec<u8>,
}
#[derive(Debug, Clone)]
pub struct PipelineOutput {
pub chunks: Vec<EncryptedChunk>,
pub total_ciphertext_bytes: usize,
}
impl ChunkPolicy {
pub fn from_str(value: &str) -> Result<Self> {
match value {
"fixed_2m" => Ok(Self::Fixed2M),
"fixed_4m" => Ok(Self::Fixed4M),
"fixed_8m" => Ok(Self::Fixed8M),
"adaptive_8_to_4" => Ok(Self::Adaptive8To4),
_ => anyhow::bail!("unsupported chunk policy: {value}"),
}
}
}
impl ChunkController {
pub fn new(cpu_generation: CpuGeneration, policy: ChunkPolicy) -> Self {
let current_chunk_bytes = match policy {
ChunkPolicy::Fixed2M => CHUNK_2_MB,
ChunkPolicy::Fixed4M => CHUNK_4_MB,
ChunkPolicy::Fixed8M => CHUNK_8_MB,
ChunkPolicy::Adaptive8To4 => match cpu_generation {
CpuGeneration::Zen2 => CHUNK_2_MB,
CpuGeneration::Zen3 | CpuGeneration::Unknown => CHUNK_8_MB,
},
};
Self {
policy,
current_chunk_bytes,
downshift_streak: 0,
upshift_streak: 0,
}
}
pub fn current_chunk_bytes(&self) -> usize {
self.current_chunk_bytes
}
pub fn policy(&self) -> ChunkPolicy {
self.policy
}
pub fn observe(&mut self, metrics: RuntimeMetrics) {
if self.policy != ChunkPolicy::Adaptive8To4 {
return;
}
let should_downshift = metrics.put_p99_ms > 200.0 && metrics.median_queue_depth > 64.0;
let should_upshift = metrics.put_p99_ms < 120.0 && metrics.median_queue_depth < 32.0;
if self.current_chunk_bytes == CHUNK_8_MB {
if should_downshift {
self.downshift_streak += 1;
} else {
self.downshift_streak = 0;
}
self.upshift_streak = 0;
if self.downshift_streak >= 3 {
self.current_chunk_bytes = CHUNK_4_MB;
self.downshift_streak = 0;
}
return;
}
if self.current_chunk_bytes == CHUNK_4_MB {
if should_upshift {
self.upshift_streak += 1;
} else {
self.upshift_streak = 0;
}
self.downshift_streak = 0;
if self.upshift_streak >= 30 {
self.current_chunk_bytes = CHUNK_8_MB;
self.upshift_streak = 0;
}
}
}
}
pub fn process_payload_crc_zstd_aes(
payload: &[u8],
chunk_bytes: usize,
key_bytes: &[u8; 32],
nonce_seed: u64,
) -> Result<PipelineOutput> {
if chunk_bytes == 0 {
anyhow::bail!("chunk_bytes must be > 0");
}
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 mut compressor = Compressor::new(1).context("failed initializing zstd compressor")?;
let mut compressed_buffer = vec![0_u8; zstd::zstd_safe::compress_bound(chunk_bytes)];
let mut output = PipelineOutput {
chunks: Vec::new(),
total_ciphertext_bytes: 0,
};
for (chunk_idx, chunk) in payload.chunks(chunk_bytes).enumerate() {
let crc = crc32c::crc32c(chunk);
let compressed_len = compressor
.compress_to_buffer(chunk, &mut compressed_buffer)
.with_context(|| format!("zstd compression failed at chunk {chunk_idx}"))?;
let nonce_id = nonce_seed + chunk_idx as u64 + 1;
let nonce = nonce_for_chunk(nonce_id);
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 {chunk_idx}"))?;
let mut bytes = Vec::with_capacity(compressed_len + tag.as_ref().len());
bytes.extend_from_slice(&compressed_buffer[..compressed_len]);
bytes.extend_from_slice(tag.as_ref());
output.total_ciphertext_bytes += bytes.len();
output.chunks.push(EncryptedChunk {
nonce_id,
crc32c: crc,
plain_len: chunk.len(),
cipher_len: bytes.len(),
bytes,
});
}
Ok(output)
}
pub fn restore_chunk_zstd_aes(
encrypted_chunk: &[u8],
expected_crc32c: u32,
expected_plain_len: usize,
key_bytes: &[u8; 32],
nonce_id: u64,
) -> Result<Vec<u8>> {
if encrypted_chunk.len() < aead::AES_256_GCM.tag_len() {
anyhow::bail!(
"encrypted chunk is too small: {} bytes",
encrypted_chunk.len()
);
}
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 nonce = nonce_for_chunk(nonce_id);
let mut sealed = encrypted_chunk.to_vec();
let decrypted = key
.open_in_place(nonce, Aad::from(expected_crc32c.to_le_bytes()), &mut sealed)
.map_err(|_| anyhow::anyhow!("AES-256-GCM decrypt failed"))?;
let restored = zstd::bulk::decompress(decrypted, expected_plain_len)
.context("zstd decompression failed during restore")?;
let crc = crc32c::crc32c(&restored);
if crc != expected_crc32c {
anyhow::bail!(
"crc32c mismatch after decrypt/decompress: expected {}, got {}",
expected_crc32c,
crc
);
}
Ok(restored)
}
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 adaptive_policy_downshifts_after_three_windows() {
let mut controller = ChunkController::new(CpuGeneration::Zen3, ChunkPolicy::Adaptive8To4);
for _ in 0..2 {
controller.observe(RuntimeMetrics {
put_p99_ms: 250.0,
median_queue_depth: 100.0,
});
assert_eq!(controller.current_chunk_bytes(), CHUNK_8_MB);
}
controller.observe(RuntimeMetrics {
put_p99_ms: 250.0,
median_queue_depth: 100.0,
});
assert_eq!(controller.current_chunk_bytes(), CHUNK_4_MB);
}
#[test]
fn adaptive_policy_upshifts_after_thirty_windows() {
let mut controller = ChunkController::new(CpuGeneration::Zen3, ChunkPolicy::Adaptive8To4);
for _ in 0..3 {
controller.observe(RuntimeMetrics {
put_p99_ms: 300.0,
median_queue_depth: 120.0,
});
}
assert_eq!(controller.current_chunk_bytes(), CHUNK_4_MB);
for _ in 0..30 {
controller.observe(RuntimeMetrics {
put_p99_ms: 80.0,
median_queue_depth: 20.0,
});
}
assert_eq!(controller.current_chunk_bytes(), CHUNK_8_MB);
}
#[test]
fn process_payload_emits_multiple_chunks() {
let payload = vec![7_u8; CHUNK_2_MB + 1024];
let output = process_payload_crc_zstd_aes(&payload, CHUNK_2_MB, &[0x42_u8; 32], 100)
.expect("pipeline processing should succeed");
assert_eq!(output.chunks.len(), 2);
assert!(output.total_ciphertext_bytes > 0);
}
#[test]
fn restore_chunk_round_trip_works() {
let payload = vec![9_u8; CHUNK_2_MB + 321];
let output = process_payload_crc_zstd_aes(&payload, CHUNK_2_MB, &[0xAA_u8; 32], 42)
.expect("pipeline processing should succeed");
let mut restored = Vec::new();
for chunk in output.chunks {
let plain = restore_chunk_zstd_aes(
&chunk.bytes,
chunk.crc32c,
chunk.plain_len,
&[0xAA_u8; 32],
chunk.nonce_id,
)
.expect("restore should succeed");
restored.extend_from_slice(&plain);
}
assert_eq!(restored, payload);
}
}