use std::collections::HashMap;
use std::path::Path;
use crate::{DrivenError, Result};
use super::SharedRules;
#[derive(Debug, Clone)]
pub struct RuleSnapshot {
pub id: u64,
pub timestamp: u64,
pub version: u64,
pub rules: HashMap<u32, Vec<u8>>,
pub metadata: HashMap<String, String>,
}
impl RuleSnapshot {
pub fn new(id: u64) -> Self {
Self {
id,
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
version: 0,
rules: HashMap::new(),
metadata: HashMap::new(),
}
}
pub fn capture(shared: &SharedRules, id: u64) -> Self {
let mut snapshot = Self::new(id);
snapshot.version = shared.version();
for rule_id in shared.ids() {
if let Some(rule) = shared.get(rule_id) {
snapshot.rules.insert(rule_id, rule.as_bytes().to_vec());
}
}
snapshot
}
pub fn restore(&self, shared: &SharedRules) {
shared.clear();
for (_id, content) in &self.rules {
shared.insert(content.clone());
}
}
pub fn size(&self) -> usize {
self.rules.values().map(|v| v.len()).sum()
}
pub fn rule_count(&self) -> usize {
self.rules.len()
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut output = Vec::new();
output.extend_from_slice(&self.id.to_le_bytes());
output.extend_from_slice(&self.timestamp.to_le_bytes());
output.extend_from_slice(&self.version.to_le_bytes());
output.extend_from_slice(&(self.rules.len() as u32).to_le_bytes());
for (id, content) in &self.rules {
output.extend_from_slice(&id.to_le_bytes());
output.extend_from_slice(&(content.len() as u32).to_le_bytes());
output.extend_from_slice(content);
}
output.extend_from_slice(&(self.metadata.len() as u32).to_le_bytes());
for (key, value) in &self.metadata {
output.extend_from_slice(&(key.len() as u16).to_le_bytes());
output.extend_from_slice(key.as_bytes());
output.extend_from_slice(&(value.len() as u16).to_le_bytes());
output.extend_from_slice(value.as_bytes());
}
output
}
pub fn from_bytes(data: &[u8]) -> Result<Self> {
if data.len() < 28 {
return Err(DrivenError::InvalidBinary("Snapshot too small".into()));
}
let mut pos = 0;
let id = u64::from_le_bytes([
data[pos],
data[pos + 1],
data[pos + 2],
data[pos + 3],
data[pos + 4],
data[pos + 5],
data[pos + 6],
data[pos + 7],
]);
pos += 8;
let timestamp = u64::from_le_bytes([
data[pos],
data[pos + 1],
data[pos + 2],
data[pos + 3],
data[pos + 4],
data[pos + 5],
data[pos + 6],
data[pos + 7],
]);
pos += 8;
let version = u64::from_le_bytes([
data[pos],
data[pos + 1],
data[pos + 2],
data[pos + 3],
data[pos + 4],
data[pos + 5],
data[pos + 6],
data[pos + 7],
]);
pos += 8;
let rule_count =
u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]) as usize;
pos += 4;
let mut rules = HashMap::with_capacity(rule_count);
for _ in 0..rule_count {
if pos + 8 > data.len() {
return Err(DrivenError::InvalidBinary("Truncated snapshot".into()));
}
let id = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
pos += 4;
let len = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]])
as usize;
pos += 4;
if pos + len > data.len() {
return Err(DrivenError::InvalidBinary("Truncated rule data".into()));
}
rules.insert(id, data[pos..pos + len].to_vec());
pos += len;
}
let mut metadata = HashMap::new();
if pos + 4 <= data.len() {
let meta_count =
u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]])
as usize;
pos += 4;
for _ in 0..meta_count {
if pos + 2 > data.len() {
break;
}
let key_len = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize;
pos += 2;
if pos + key_len > data.len() {
break;
}
let key = String::from_utf8_lossy(&data[pos..pos + key_len]).to_string();
pos += key_len;
if pos + 2 > data.len() {
break;
}
let val_len = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize;
pos += 2;
if pos + val_len > data.len() {
break;
}
let value = String::from_utf8_lossy(&data[pos..pos + val_len]).to_string();
pos += val_len;
metadata.insert(key, value);
}
}
Ok(Self {
id,
timestamp,
version,
rules,
metadata,
})
}
}
#[derive(Debug)]
pub struct SnapshotManager {
directory: std::path::PathBuf,
next_id: u64,
max_snapshots: usize,
}
impl SnapshotManager {
pub fn new(directory: impl AsRef<Path>) -> Result<Self> {
let directory = directory.as_ref().to_path_buf();
std::fs::create_dir_all(&directory)?;
let next_id = Self::find_highest_id(&directory)? + 1;
Ok(Self {
directory,
next_id,
max_snapshots: 10,
})
}
pub fn with_max_snapshots(mut self, max: usize) -> Self {
self.max_snapshots = max;
self
}
pub fn create(&mut self, shared: &SharedRules) -> Result<RuleSnapshot> {
let snapshot = RuleSnapshot::capture(shared, self.next_id);
self.next_id += 1;
self.save(&snapshot)?;
self.cleanup()?;
Ok(snapshot)
}
pub fn save(&self, snapshot: &RuleSnapshot) -> Result<()> {
let path = self.snapshot_path(snapshot.id);
let data = snapshot.to_bytes();
std::fs::write(&path, &data)?;
Ok(())
}
pub fn load(&self, id: u64) -> Result<RuleSnapshot> {
let path = self.snapshot_path(id);
let data = std::fs::read(&path)?;
RuleSnapshot::from_bytes(&data)
}
pub fn list(&self) -> Result<Vec<u64>> {
let mut ids = Vec::new();
for entry in std::fs::read_dir(&self.directory)? {
let entry = entry?;
if let Some(name) = entry.file_name().to_str() {
if let Some(id_str) = name
.strip_prefix("snapshot_")
.and_then(|s| s.strip_suffix(".drv"))
{
if let Ok(id) = id_str.parse() {
ids.push(id);
}
}
}
}
ids.sort();
Ok(ids)
}
pub fn load_latest(&self) -> Result<Option<RuleSnapshot>> {
let ids = self.list()?;
match ids.last() {
Some(&id) => Ok(Some(self.load(id)?)),
None => Ok(None),
}
}
pub fn delete(&self, id: u64) -> Result<()> {
let path = self.snapshot_path(id);
std::fs::remove_file(&path)?;
Ok(())
}
fn cleanup(&self) -> Result<()> {
let mut ids = self.list()?;
while ids.len() > self.max_snapshots {
if let Some(oldest) = ids.first() {
self.delete(*oldest)?;
ids.remove(0);
}
}
Ok(())
}
fn snapshot_path(&self, id: u64) -> std::path::PathBuf {
self.directory.join(format!("snapshot_{:08}.drv", id))
}
fn find_highest_id(directory: &Path) -> Result<u64> {
let mut highest = 0u64;
if directory.exists() {
for entry in std::fs::read_dir(directory)? {
let entry = entry?;
if let Some(name) = entry.file_name().to_str() {
if let Some(id_str) = name
.strip_prefix("snapshot_")
.and_then(|s| s.strip_suffix(".drv"))
{
if let Ok(id) = id_str.parse::<u64>() {
highest = highest.max(id);
}
}
}
}
}
Ok(highest)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_snapshot_roundtrip() {
let mut snapshot = RuleSnapshot::new(1);
snapshot.rules.insert(1, b"rule one".to_vec());
snapshot.rules.insert(2, b"rule two".to_vec());
snapshot
.metadata
.insert("key".to_string(), "value".to_string());
let bytes = snapshot.to_bytes();
let restored = RuleSnapshot::from_bytes(&bytes).unwrap();
assert_eq!(restored.id, 1);
assert_eq!(restored.rules.len(), 2);
assert_eq!(restored.rules.get(&1).unwrap(), b"rule one");
}
#[test]
fn test_capture_restore() {
let shared = SharedRules::new();
shared.insert(b"test rule".to_vec());
let snapshot = RuleSnapshot::capture(&shared, 1);
assert_eq!(snapshot.rule_count(), 1);
let new_shared = SharedRules::new();
snapshot.restore(&new_shared);
assert_eq!(new_shared.len(), 1);
}
}