From fe92eba0d6c00d5f80ad01d73b377ed3b8af6652 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 09:29:46 -0500 Subject: [PATCH 1/7] feat: add encryption module and v3 format support to ember-persistence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit adds an optional `encryption` feature flag backed by AES-256-GCM. when enabled, provides: - EncryptionKey type with from_file (raw 32B or 64 hex chars) - encrypt_record / decrypt_record primitives - FormatError::EncryptionRequired and DecryptionFailed variants - FORMAT_VERSION_ENCRYPTED (v3) and version-gated header reading the feature is off by default — no crypto deps are linked and no encryption types exist in the API when disabled. --- crates/ember-persistence/Cargo.toml | 7 + crates/ember-persistence/src/encryption.rs | 255 +++++++++++++++++++++ crates/ember-persistence/src/format.rs | 32 ++- crates/ember-persistence/src/lib.rs | 2 + 4 files changed, 293 insertions(+), 3 deletions(-) create mode 100644 crates/ember-persistence/src/encryption.rs diff --git a/crates/ember-persistence/Cargo.toml b/crates/ember-persistence/Cargo.toml index f5b22b34..eca0d0b0 100644 --- a/crates/ember-persistence/Cargo.toml +++ b/crates/ember-persistence/Cargo.toml @@ -9,11 +9,18 @@ keywords.workspace = true categories.workspace = true readme = "README.md" +[features] +encryption = ["aes-gcm", "rand"] + [dependencies] thiserror = { workspace = true } bytes = { workspace = true } crc32fast = { workspace = true } tracing = { workspace = true } +# optional: encryption at rest (AES-256-GCM) +aes-gcm = { version = "0.10", optional = true } +rand = { workspace = true, optional = true } + [dev-dependencies] tempfile = "3" diff --git a/crates/ember-persistence/src/encryption.rs b/crates/ember-persistence/src/encryption.rs new file mode 100644 index 00000000..b80e2aa3 --- /dev/null +++ b/crates/ember-persistence/src/encryption.rs @@ -0,0 +1,255 @@ +//! Encryption at rest using AES-256-GCM. +//! +//! Each record (AOF or snapshot entry) is encrypted independently with a +//! random 12-byte nonce. AES-GCM provides authenticated encryption — a +//! tampered ciphertext is detected immediately rather than producing garbage. +//! +//! This module is only compiled when the `encryption` feature is enabled. + +use std::fmt; +use std::io; +use std::path::Path; + +use aes_gcm::aead::{Aead, KeyInit, OsRng}; +use aes_gcm::{AeadCore, Aes256Gcm, Nonce}; + +use crate::format::FormatError; + +/// AES-256-GCM nonce size in bytes. +pub const NONCE_SIZE: usize = 12; + +/// AES-256-GCM authentication tag size in bytes. +pub const TAG_SIZE: usize = 16; + +/// A 256-bit encryption key for AES-256-GCM. +/// +/// The key is stored inline — no heap allocation. Implements `Clone` +/// but not `Debug` to avoid accidentally logging key material. +#[derive(Clone)] +pub struct EncryptionKey { + bytes: [u8; 32], +} + +impl fmt::Debug for EncryptionKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("EncryptionKey") + .field("bytes", &"[redacted]") + .finish() + } +} + +impl EncryptionKey { + /// Creates a key from raw bytes. + pub fn from_bytes(bytes: [u8; 32]) -> Self { + Self { bytes } + } + + /// Reads an encryption key from a file. + /// + /// The file may contain either: + /// - exactly 32 raw bytes, or + /// - 64 hex characters (with optional trailing whitespace/newline) + pub fn from_file(path: &Path) -> Result { + let data = std::fs::read(path).map_err(|e| { + FormatError::Io(io::Error::new( + e.kind(), + format!("failed to read encryption key file '{}': {e}", path.display()), + )) + })?; + + // try raw 32 bytes first + if data.len() == 32 { + let mut bytes = [0u8; 32]; + bytes.copy_from_slice(&data); + return Ok(Self { bytes }); + } + + // try hex-encoded (64 chars + optional trailing whitespace) + let trimmed = std::str::from_utf8(&data) + .map_err(|_| { + FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + "encryption key file is not valid UTF-8 or raw 32 bytes", + )) + })? + .trim(); + + if trimmed.len() != 64 { + return Err(FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "encryption key file must be 32 raw bytes or 64 hex characters, got {} bytes", + data.len() + ), + ))); + } + + let mut bytes = [0u8; 32]; + for i in 0..32 { + bytes[i] = u8::from_str_radix(&trimmed[i * 2..i * 2 + 2], 16).map_err(|_| { + FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + "encryption key file contains invalid hex characters", + )) + })?; + } + + Ok(Self { bytes }) + } + + /// Returns the raw key bytes. + fn as_bytes(&self) -> &[u8; 32] { + &self.bytes + } +} + +/// Encrypts a plaintext record using AES-256-GCM. +/// +/// Returns `(nonce, ciphertext)` where ciphertext includes the 16-byte +/// auth tag appended by AES-GCM. +pub fn encrypt_record(key: &EncryptionKey, plaintext: &[u8]) -> Result<([u8; NONCE_SIZE], Vec), FormatError> { + let cipher = Aes256Gcm::new(key.as_bytes().into()); + let nonce = Aes256Gcm::generate_nonce(&mut OsRng); + + let ciphertext = cipher.encrypt(&nonce, plaintext).map_err(|e| { + FormatError::Io(io::Error::new( + io::ErrorKind::Other, + format!("encryption failed: {e}"), + )) + })?; + + let mut nonce_bytes = [0u8; NONCE_SIZE]; + nonce_bytes.copy_from_slice(&nonce); + Ok((nonce_bytes, ciphertext)) +} + +/// Decrypts a ciphertext record using AES-256-GCM. +/// +/// The ciphertext must include the 16-byte auth tag (as produced by +/// [`encrypt_record`]). Returns `DecryptionFailed` if the key is wrong +/// or the data has been tampered with. +pub fn decrypt_record(key: &EncryptionKey, nonce: &[u8; NONCE_SIZE], ciphertext: &[u8]) -> Result, FormatError> { + let cipher = Aes256Gcm::new(key.as_bytes().into()); + let nonce = Nonce::from_slice(nonce); + + cipher.decrypt(nonce, ciphertext).map_err(|_| FormatError::DecryptionFailed) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_key() -> EncryptionKey { + EncryptionKey::from_bytes([0x42; 32]) + } + + #[test] + fn round_trip() { + let key = test_key(); + let plaintext = b"hello, encrypted world"; + + let (nonce, ciphertext) = encrypt_record(&key, plaintext).unwrap(); + let decrypted = decrypt_record(&key, &nonce, &ciphertext).unwrap(); + + assert_eq!(decrypted, plaintext); + } + + #[test] + fn wrong_key_fails() { + let key = test_key(); + let wrong_key = EncryptionKey::from_bytes([0xFF; 32]); + let plaintext = b"secret data"; + + let (nonce, ciphertext) = encrypt_record(&key, plaintext).unwrap(); + let err = decrypt_record(&wrong_key, &nonce, &ciphertext).unwrap_err(); + + assert!(matches!(err, FormatError::DecryptionFailed)); + } + + #[test] + fn tampered_ciphertext_fails() { + let key = test_key(); + let plaintext = b"integrity check"; + + let (nonce, mut ciphertext) = encrypt_record(&key, plaintext).unwrap(); + // flip a byte in the ciphertext + ciphertext[0] ^= 0xFF; + + let err = decrypt_record(&key, &nonce, &ciphertext).unwrap_err(); + assert!(matches!(err, FormatError::DecryptionFailed)); + } + + #[test] + fn empty_plaintext() { + let key = test_key(); + let plaintext = b""; + + let (nonce, ciphertext) = encrypt_record(&key, plaintext).unwrap(); + // ciphertext should be exactly the auth tag size + assert_eq!(ciphertext.len(), TAG_SIZE); + + let decrypted = decrypt_record(&key, &nonce, &ciphertext).unwrap(); + assert!(decrypted.is_empty()); + } + + #[test] + fn different_nonces_per_call() { + let key = test_key(); + let plaintext = b"same data"; + + let (nonce1, ct1) = encrypt_record(&key, plaintext).unwrap(); + let (nonce2, ct2) = encrypt_record(&key, plaintext).unwrap(); + + // nonces should differ (probabilistically guaranteed with random nonces) + assert_ne!(nonce1, nonce2); + // ciphertexts should differ too + assert_ne!(ct1, ct2); + } + + #[test] + fn key_from_raw_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("key.bin"); + let raw = [0xAB; 32]; + std::fs::write(&path, raw).unwrap(); + + let key = EncryptionKey::from_file(&path).unwrap(); + assert_eq!(*key.as_bytes(), raw); + } + + #[test] + fn key_from_hex_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("key.hex"); + let hex = "ab".repeat(32); + std::fs::write(&path, format!("{hex}\n")).unwrap(); + + let key = EncryptionKey::from_file(&path).unwrap(); + assert_eq!(*key.as_bytes(), [0xAB; 32]); + } + + #[test] + fn key_from_bad_file_fails() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("key.bad"); + std::fs::write(&path, "too short").unwrap(); + + let err = EncryptionKey::from_file(&path).unwrap_err(); + assert!(matches!(err, FormatError::Io(_))); + } + + #[test] + fn key_from_missing_file_fails() { + let path = std::path::Path::new("/nonexistent/key.bin"); + let err = EncryptionKey::from_file(path).unwrap_err(); + assert!(matches!(err, FormatError::Io(_))); + } + + #[test] + fn debug_redacts_key() { + let key = test_key(); + let debug = format!("{key:?}"); + assert!(debug.contains("redacted")); + assert!(!debug.contains("42")); + } +} diff --git a/crates/ember-persistence/src/format.rs b/crates/ember-persistence/src/format.rs index 9ac43971..7e1cb7de 100644 --- a/crates/ember-persistence/src/format.rs +++ b/crates/ember-persistence/src/format.rs @@ -14,12 +14,17 @@ pub const AOF_MAGIC: &[u8; 4] = b"EAOF"; /// Magic bytes for the snapshot file header. pub const SNAP_MAGIC: &[u8; 4] = b"ESNP"; -/// Current format version for both AOF and snapshot files. +/// Current unencrypted format version. /// /// v1: original format (strings only) -/// v2: type-tagged entries (string, list, sorted set) +/// v2: type-tagged entries (string, list, sorted set, hash, set) pub const FORMAT_VERSION: u8 = 2; +/// Format version for encrypted files. +/// +/// v3: per-record AES-256-GCM encryption (requires `encryption` feature) +pub const FORMAT_VERSION_ENCRYPTED: u8 = 3; + /// Errors that can occur when reading or writing persistence formats. #[derive(Debug, Error)] pub enum FormatError { @@ -38,6 +43,12 @@ pub enum FormatError { #[error("unknown record tag: {0}")] UnknownTag(u8), + #[error("file is encrypted but no encryption key was provided")] + EncryptionRequired, + + #[error("decryption failed (wrong key or tampered data)")] + DecryptionFailed, + #[error("io error: {0}")] Io(#[from] io::Error), } @@ -162,6 +173,21 @@ pub fn write_header(w: &mut impl Write, magic: &[u8; 4]) -> io::Result<()> { write_u8(w, FORMAT_VERSION) } +/// Writes a file header with an explicit version byte. +pub fn write_header_versioned(w: &mut impl Write, magic: &[u8; 4], version: u8) -> io::Result<()> { + w.write_all(magic)?; + write_u8(w, version) +} + +/// The maximum format version this build can read. +/// +/// When the `encryption` feature is compiled in, v3 (encrypted) files +/// are supported. Without the feature, only v1 and v2 are accepted. +#[cfg(feature = "encryption")] +const MAX_READABLE_VERSION: u8 = FORMAT_VERSION_ENCRYPTED; +#[cfg(not(feature = "encryption"))] +const MAX_READABLE_VERSION: u8 = FORMAT_VERSION; + /// Reads and validates a file header. Returns an error if magic doesn't /// match or version is unsupported. Returns the format version. pub fn read_header(r: &mut impl Read, expected_magic: &[u8; 4]) -> Result { @@ -171,7 +197,7 @@ pub fn read_header(r: &mut impl Read, expected_magic: &[u8; 4]) -> Result FORMAT_VERSION { + if version == 0 || version > MAX_READABLE_VERSION { return Err(FormatError::UnsupportedVersion(version)); } Ok(version) diff --git a/crates/ember-persistence/src/lib.rs b/crates/ember-persistence/src/lib.rs index d213e42c..cb6ab899 100644 --- a/crates/ember-persistence/src/lib.rs +++ b/crates/ember-persistence/src/lib.rs @@ -4,6 +4,8 @@ //! and crash recovery. pub mod aof; +#[cfg(feature = "encryption")] +pub mod encryption; pub mod format; pub mod recovery; pub mod snapshot; From aae8bafb34dadbf5edd30ca07188724a9b6f3f6d Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 09:32:47 -0500 Subject: [PATCH 2/7] feat: add encryption support to AOF writer and reader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AofWriter gains open_encrypted() for v3 (AES-256-GCM) files. each record is written as [nonce: 12B][len: 4B][ciphertext]. AofReader gains open_encrypted() that handles both v2 (plaintext) and v3 (encrypted) files — enabling transparent migration. v3 files opened without a key return FormatError::EncryptionRequired. wrong keys are detected immediately via AEAD authentication. --- crates/ember-persistence/src/aof.rs | 338 +++++++++++++++++++++++++++- 1 file changed, 332 insertions(+), 6 deletions(-) diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 9a1e3921..b02aaa53 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -16,7 +16,10 @@ //! ``` //! The CRC32 covers the tag + payload bytes. +use std::fmt; use std::fs::{self, File, OpenOptions}; +#[cfg(feature = "encryption")] +use std::io::Read as _; use std::io::{self, BufReader, BufWriter, Write}; use std::path::{Path, PathBuf}; @@ -434,13 +437,40 @@ pub enum FsyncPolicy { pub struct AofWriter { writer: BufWriter, path: PathBuf, + #[cfg(feature = "encryption")] + encryption_key: Option, } impl AofWriter { /// Opens (or creates) an AOF file. If the file is new, writes the header. /// If the file already exists, appends to it. pub fn open(path: impl Into) -> Result { - let path = path.into(); + Self::open_inner(path.into(), { + #[cfg(feature = "encryption")] + { None } + #[cfg(not(feature = "encryption"))] + { () } + }) + } + + /// Opens (or creates) an encrypted AOF file using AES-256-GCM. + /// + /// New files get a v3 header. Existing v2 files can be appended to — + /// new records will be written unencrypted (use `BGREWRITEAOF` to + /// migrate the full file to v3). + #[cfg(feature = "encryption")] + pub fn open_encrypted( + path: impl Into, + key: crate::encryption::EncryptionKey, + ) -> Result { + Self::open_inner(path.into(), Some(key)) + } + + fn open_inner( + path: PathBuf, + #[cfg(feature = "encryption")] encryption_key: Option, + #[cfg(not(feature = "encryption"))] _: (), + ) -> Result { let exists = path.exists() && fs::metadata(&path).map(|m| m.len() > 0).unwrap_or(false); let mut opts = OpenOptions::new(); @@ -454,16 +484,45 @@ impl AofWriter { let mut writer = BufWriter::new(file); if !exists { + #[cfg(feature = "encryption")] + if encryption_key.is_some() { + format::write_header_versioned( + &mut writer, + format::AOF_MAGIC, + format::FORMAT_VERSION_ENCRYPTED, + )?; + } else { + format::write_header(&mut writer, format::AOF_MAGIC)?; + } + #[cfg(not(feature = "encryption"))] format::write_header(&mut writer, format::AOF_MAGIC)?; writer.flush()?; } - Ok(Self { writer, path }) + Ok(Self { + writer, + path, + #[cfg(feature = "encryption")] + encryption_key, + }) } - /// Appends a record to the AOF. Writes tag+payload+crc32. + /// Appends a record to the AOF. + /// + /// When an encryption key is set, writes: `[nonce: 12B][len: 4B][ciphertext]`. + /// Otherwise writes the v2 format: `[tag+payload][crc32: 4B]`. pub fn write_record(&mut self, record: &AofRecord) -> Result<(), FormatError> { let payload = record.to_bytes()?; + + #[cfg(feature = "encryption")] + if let Some(ref key) = self.encryption_key { + let (nonce, ciphertext) = crate::encryption::encrypt_record(key, &payload)?; + self.writer.write_all(&nonce)?; + format::write_u32(&mut self.writer, ciphertext.len() as u32)?; + self.writer.write_all(&ciphertext)?; + return Ok(()); + } + let checksum = format::crc32(&payload); self.writer.write_all(&payload)?; format::write_u32(&mut self.writer, checksum)?; @@ -504,7 +563,20 @@ impl AofWriter { } let file = opts.open(&self.path)?; let mut writer = BufWriter::new(file); + + #[cfg(feature = "encryption")] + if self.encryption_key.is_some() { + format::write_header_versioned( + &mut writer, + format::AOF_MAGIC, + format::FORMAT_VERSION_ENCRYPTED, + )?; + } else { + format::write_header(&mut writer, format::AOF_MAGIC)?; + } + #[cfg(not(feature = "encryption"))] format::write_header(&mut writer, format::AOF_MAGIC)?; + writer.flush()?; // ensure the fresh header is durable before we start appending writer.get_ref().sync_all()?; @@ -514,9 +586,20 @@ impl AofWriter { } /// Reader for iterating over AOF records. -#[derive(Debug)] pub struct AofReader { reader: BufReader, + /// Format version from the file header. v2 = plaintext, v3 = encrypted. + version: u8, + #[cfg(feature = "encryption")] + encryption_key: Option, +} + +impl fmt::Debug for AofReader { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AofReader") + .field("version", &self.version) + .finish() + } } impl AofReader { @@ -524,8 +607,38 @@ impl AofReader { pub fn open(path: impl AsRef) -> Result { let file = File::open(path.as_ref())?; let mut reader = BufReader::new(file); - let _version = format::read_header(&mut reader, format::AOF_MAGIC)?; - Ok(Self { reader }) + let version = format::read_header(&mut reader, format::AOF_MAGIC)?; + + if version == format::FORMAT_VERSION_ENCRYPTED { + return Err(FormatError::EncryptionRequired); + } + + Ok(Self { + reader, + version, + #[cfg(feature = "encryption")] + encryption_key: None, + }) + } + + /// Opens an AOF file with an encryption key for decrypting v3 records. + /// + /// Also handles v2 (plaintext) files — the key is simply unused, + /// allowing transparent migration. + #[cfg(feature = "encryption")] + pub fn open_encrypted( + path: impl AsRef, + key: crate::encryption::EncryptionKey, + ) -> Result { + let file = File::open(path.as_ref())?; + let mut reader = BufReader::new(file); + let version = format::read_header(&mut reader, format::AOF_MAGIC)?; + + Ok(Self { + reader, + version, + encryption_key: Some(key), + }) } /// Reads the next record from the AOF. @@ -534,6 +647,16 @@ impl AofReader { /// server crashed mid-write), returns `Ok(None)` rather than an error /// — this is the expected recovery behavior. pub fn read_record(&mut self) -> Result, FormatError> { + #[cfg(feature = "encryption")] + if self.version == format::FORMAT_VERSION_ENCRYPTED { + return self.read_encrypted_record(); + } + + self.read_v2_record() + } + + /// Reads a v2 (plaintext) record: tag + payload + crc32. + fn read_v2_record(&mut self) -> Result, FormatError> { // peek for EOF — try reading the tag byte let tag = match format::read_u8(&mut self.reader) { Ok(t) => t, @@ -559,6 +682,51 @@ impl AofReader { } } + /// Reads a v3 (encrypted) record: nonce + len + ciphertext. + #[cfg(feature = "encryption")] + fn read_encrypted_record(&mut self) -> Result, FormatError> { + let key = self + .encryption_key + .as_ref() + .ok_or(FormatError::EncryptionRequired)?; + + // read the 12-byte nonce + let mut nonce = [0u8; crate::encryption::NONCE_SIZE]; + if let Err(e) = self.reader.read_exact(&mut nonce) { + return if e.kind() == io::ErrorKind::UnexpectedEof { + Ok(None) + } else { + Err(FormatError::Io(e)) + }; + } + + // read ciphertext length and ciphertext + let ct_len = match format::read_u32(&mut self.reader) { + Ok(n) => n as usize, + Err(FormatError::UnexpectedEof) => return Ok(None), + Err(e) => return Err(e), + }; + + if ct_len > format::MAX_FIELD_LEN { + return Err(FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + format!("encrypted record length {ct_len} exceeds maximum"), + ))); + } + + let mut ciphertext = vec![0u8; ct_len]; + if let Err(e) = self.reader.read_exact(&mut ciphertext) { + return if e.kind() == io::ErrorKind::UnexpectedEof { + Ok(None) + } else { + Err(FormatError::Io(e)) + }; + } + + let plaintext = crate::encryption::decrypt_record(key, &nonce, &ciphertext)?; + AofRecord::from_bytes(&plaintext).map(Some) + } + /// Reads the remaining payload bytes (after the tag) and the trailing CRC. fn read_payload_for_tag(&mut self, tag: u8) -> Result<(Vec, u32), FormatError> { let mut payload = Vec::new(); @@ -1121,4 +1289,162 @@ mod tests { let decoded = AofRecord::from_bytes(&bytes).unwrap(); assert_eq!(rec, decoded); } + + #[cfg(feature = "encryption")] + mod encrypted { + use super::*; + use crate::encryption::EncryptionKey; + + fn test_key() -> EncryptionKey { + EncryptionKey::from_bytes([0x42; 32]) + } + + #[test] + fn encrypted_writer_reader_round_trip() { + let dir = temp_dir(); + let path = dir.path().join("enc.aof"); + let key = test_key(); + + let records = vec![ + AofRecord::Set { + key: "a".into(), + value: Bytes::from("1"), + expire_ms: -1, + }, + AofRecord::Del { key: "a".into() }, + AofRecord::LPush { + key: "list".into(), + values: vec![Bytes::from("x"), Bytes::from("y")], + }, + AofRecord::ZAdd { + key: "zs".into(), + members: vec![(1.0, "m".into())], + }, + ]; + + { + let mut writer = AofWriter::open_encrypted(&path, key.clone()).unwrap(); + for rec in &records { + writer.write_record(rec).unwrap(); + } + writer.sync().unwrap(); + } + + let mut reader = AofReader::open_encrypted(&path, key).unwrap(); + let mut got = Vec::new(); + while let Some(rec) = reader.read_record().unwrap() { + got.push(rec); + } + assert_eq!(records, got); + } + + #[test] + fn encrypted_aof_wrong_key_fails() { + let dir = temp_dir(); + let path = dir.path().join("enc_bad.aof"); + let key = test_key(); + let wrong_key = EncryptionKey::from_bytes([0xFF; 32]); + + { + let mut writer = AofWriter::open_encrypted(&path, key).unwrap(); + writer + .write_record(&AofRecord::Set { + key: "k".into(), + value: Bytes::from("v"), + expire_ms: -1, + }) + .unwrap(); + writer.sync().unwrap(); + } + + let mut reader = AofReader::open_encrypted(&path, wrong_key).unwrap(); + let err = reader.read_record().unwrap_err(); + assert!(matches!(err, FormatError::DecryptionFailed)); + } + + #[test] + fn v2_file_readable_with_encryption_key() { + let dir = temp_dir(); + let path = dir.path().join("v2.aof"); + let key = test_key(); + + // write a plaintext v2 file + { + let mut writer = AofWriter::open(&path).unwrap(); + writer + .write_record(&AofRecord::Set { + key: "k".into(), + value: Bytes::from("v"), + expire_ms: -1, + }) + .unwrap(); + writer.sync().unwrap(); + } + + // read with encryption key — should work (v2 is plaintext) + let mut reader = AofReader::open_encrypted(&path, key).unwrap(); + let rec = reader.read_record().unwrap().unwrap(); + assert!(matches!(rec, AofRecord::Set { .. })); + } + + #[test] + fn v3_file_without_key_returns_error() { + let dir = temp_dir(); + let path = dir.path().join("v3_nokey.aof"); + let key = test_key(); + + // write an encrypted v3 file + { + let mut writer = AofWriter::open_encrypted(&path, key).unwrap(); + writer + .write_record(&AofRecord::Set { + key: "k".into(), + value: Bytes::from("v"), + expire_ms: -1, + }) + .unwrap(); + writer.sync().unwrap(); + } + + // try to open without a key + let err = AofReader::open(&path).unwrap_err(); + assert!(matches!(err, FormatError::EncryptionRequired)); + } + + #[test] + fn encrypted_truncate_preserves_encryption() { + let dir = temp_dir(); + let path = dir.path().join("enc_trunc.aof"); + let key = test_key(); + + { + let mut writer = AofWriter::open_encrypted(&path, key.clone()).unwrap(); + writer + .write_record(&AofRecord::Set { + key: "old".into(), + value: Bytes::from("data"), + expire_ms: -1, + }) + .unwrap(); + writer.truncate().unwrap(); + + writer + .write_record(&AofRecord::Set { + key: "new".into(), + value: Bytes::from("fresh"), + expire_ms: -1, + }) + .unwrap(); + writer.sync().unwrap(); + } + + let mut reader = AofReader::open_encrypted(&path, key).unwrap(); + let rec = reader.read_record().unwrap().unwrap(); + match rec { + AofRecord::Set { key, .. } => assert_eq!(key, "new"), + other => panic!("expected Set, got {other:?}"), + } + assert!(reader.read_record().unwrap().is_none()); + } + } } From 9449203861a3bf6a270bdd63efdfd63a790b4e6d Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 09:35:38 -0500 Subject: [PATCH 3/7] feat: add encryption support to snapshot writer and reader SnapshotWriter gains create_encrypted() for v3 files. each entry is encrypted independently as [nonce: 12B][len: 4B][ciphertext]. footer CRC covers the encrypted bytes for truncation detection. SnapshotReader gains open_encrypted() that handles both v2 (plaintext) and v3 (encrypted) files transparently. wrong keys and tampered data are detected via AEAD authentication. --- crates/ember-persistence/src/snapshot.rs | 415 ++++++++++++++++++++++- 1 file changed, 413 insertions(+), 2 deletions(-) diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index 5bad3424..535fda5b 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -76,13 +76,38 @@ pub struct SnapshotWriter { /// Set to `true` after a successful `finish()`. Used by the `Drop` /// impl to clean up incomplete temp files. finished: bool, + #[cfg(feature = "encryption")] + encryption_key: Option, } impl SnapshotWriter { /// Creates a new snapshot writer. The file won't appear at `path` /// until [`Self::finish`] is called successfully. pub fn create(path: impl Into, shard_id: u16) -> Result { - let final_path = path.into(); + Self::create_inner(path.into(), shard_id, { + #[cfg(feature = "encryption")] + { None } + #[cfg(not(feature = "encryption"))] + { () } + }) + } + + /// Creates a new encrypted snapshot writer. + #[cfg(feature = "encryption")] + pub fn create_encrypted( + path: impl Into, + shard_id: u16, + key: crate::encryption::EncryptionKey, + ) -> Result { + Self::create_inner(path.into(), shard_id, Some(key)) + } + + fn create_inner( + final_path: PathBuf, + shard_id: u16, + #[cfg(feature = "encryption")] encryption_key: Option, + #[cfg(not(feature = "encryption"))] _: (), + ) -> Result { let tmp_path = final_path.with_extension("snap.tmp"); let file = { @@ -98,7 +123,19 @@ impl SnapshotWriter { let mut writer = BufWriter::new(file); // write header: magic + version + shard_id + placeholder entry count + #[cfg(feature = "encryption")] + if encryption_key.is_some() { + format::write_header_versioned( + &mut writer, + format::SNAP_MAGIC, + format::FORMAT_VERSION_ENCRYPTED, + )?; + } else { + format::write_header(&mut writer, format::SNAP_MAGIC)?; + } + #[cfg(not(feature = "encryption"))] format::write_header(&mut writer, format::SNAP_MAGIC)?; + format::write_u16(&mut writer, shard_id)?; // entry count — we'll seek back and update, or just write it now // and track. since we're streaming, write 0 and update after. @@ -111,10 +148,15 @@ impl SnapshotWriter { hasher: crc32fast::Hasher::new(), count: 0, finished: false, + #[cfg(feature = "encryption")] + encryption_key, }) } /// Writes a single entry to the snapshot. + /// + /// When encrypted, each entry is written as `[nonce: 12B][len: 4B][ciphertext]`. + /// The footer CRC covers the encrypted bytes (nonce + len + ciphertext). pub fn write_entry(&mut self, entry: &SnapEntry) -> Result<(), FormatError> { let mut buf = Vec::new(); format::write_bytes(&mut buf, entry.key.as_bytes())?; @@ -156,6 +198,21 @@ impl SnapshotWriter { } format::write_i64(&mut buf, entry.expire_ms)?; + #[cfg(feature = "encryption")] + if let Some(ref key) = self.encryption_key { + let (nonce, ciphertext) = crate::encryption::encrypt_record(key, &buf)?; + // footer CRC covers the encrypted envelope + self.hasher.update(&nonce); + let ct_len_bytes = (ciphertext.len() as u32).to_le_bytes(); + self.hasher.update(&ct_len_bytes); + self.hasher.update(&ciphertext); + self.writer.write_all(&nonce)?; + format::write_u32(&mut self.writer, ciphertext.len() as u32)?; + self.writer.write_all(&ciphertext)?; + self.count += 1; + return Ok(()); + } + self.hasher.update(&buf); self.writer.write_all(&buf)?; self.count += 1; @@ -206,8 +263,10 @@ pub struct SnapshotReader { pub entry_count: u32, read_so_far: u32, hasher: crc32fast::Hasher, - /// Format version — v1 has no type tags, v2 has type-tagged entries. + /// Format version — v1 has no type tags, v2 has type-tagged entries, v3 is encrypted. version: u8, + #[cfg(feature = "encryption")] + encryption_key: Option, } impl SnapshotReader { @@ -217,6 +276,11 @@ impl SnapshotReader { let mut reader = BufReader::new(file); let version = format::read_header(&mut reader, format::SNAP_MAGIC)?; + + if version == format::FORMAT_VERSION_ENCRYPTED { + return Err(FormatError::EncryptionRequired); + } + let shard_id = format::read_u16(&mut reader)?; let entry_count = format::read_u32(&mut reader)?; @@ -227,6 +291,34 @@ impl SnapshotReader { read_so_far: 0, hasher: crc32fast::Hasher::new(), version, + #[cfg(feature = "encryption")] + encryption_key: None, + }) + } + + /// Opens a snapshot file with an encryption key for decrypting v3 entries. + /// + /// Also handles v1/v2 (plaintext) files — the key is simply unused. + #[cfg(feature = "encryption")] + pub fn open_encrypted( + path: impl AsRef, + key: crate::encryption::EncryptionKey, + ) -> Result { + let file = File::open(path.as_ref())?; + let mut reader = BufReader::new(file); + + let version = format::read_header(&mut reader, format::SNAP_MAGIC)?; + let shard_id = format::read_u16(&mut reader)?; + let entry_count = format::read_u32(&mut reader)?; + + Ok(Self { + reader, + shard_id, + entry_count, + read_so_far: 0, + hasher: crc32fast::Hasher::new(), + version, + encryption_key: Some(key), }) } @@ -236,6 +328,16 @@ impl SnapshotReader { return Ok(None); } + #[cfg(feature = "encryption")] + if self.version == format::FORMAT_VERSION_ENCRYPTED { + return self.read_encrypted_entry(); + } + + self.read_plaintext_entry() + } + + /// Reads a plaintext (v1/v2) entry. + fn read_plaintext_entry(&mut self) -> Result, FormatError> { let mut buf = Vec::new(); let key_bytes = format::read_bytes(&mut self.reader)?; @@ -347,6 +449,132 @@ impl SnapshotReader { })) } + /// Reads an encrypted (v3) entry: nonce + len + ciphertext. + /// Decrypts to get the same bytes as a plaintext entry, then parses. + #[cfg(feature = "encryption")] + fn read_encrypted_entry(&mut self) -> Result, FormatError> { + use std::io::Read as _; + + let key = self + .encryption_key + .as_ref() + .ok_or(FormatError::EncryptionRequired)?; + + let mut nonce = [0u8; crate::encryption::NONCE_SIZE]; + self.reader + .read_exact(&mut nonce) + .map_err(|e| match e.kind() { + io::ErrorKind::UnexpectedEof => FormatError::UnexpectedEof, + _ => FormatError::Io(e), + })?; + + let ct_len = format::read_u32(&mut self.reader)? as usize; + if ct_len > format::MAX_FIELD_LEN { + return Err(FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + format!("encrypted entry length {ct_len} exceeds maximum"), + ))); + } + + let mut ciphertext = vec![0u8; ct_len]; + self.reader + .read_exact(&mut ciphertext) + .map_err(|e| match e.kind() { + io::ErrorKind::UnexpectedEof => FormatError::UnexpectedEof, + _ => FormatError::Io(e), + })?; + + // footer CRC covers the encrypted envelope + self.hasher.update(&nonce); + let ct_len_bytes = (ct_len as u32).to_le_bytes(); + self.hasher.update(&ct_len_bytes); + self.hasher.update(&ciphertext); + + let plaintext = crate::encryption::decrypt_record(key, &nonce, &ciphertext)?; + + // parse the decrypted bytes using the same logic as v2 + let mut cursor = io::Cursor::new(&plaintext); + let key_bytes = format::read_bytes(&mut cursor)?; + let type_tag = format::read_u8(&mut cursor)?; + let value = match type_tag { + TYPE_STRING => { + let v = format::read_bytes(&mut cursor)?; + SnapValue::String(Bytes::from(v)) + } + TYPE_LIST => { + let count = format::read_u32(&mut cursor)?; + let mut deque = VecDeque::with_capacity(count as usize); + for _ in 0..count { + deque.push_back(Bytes::from(format::read_bytes(&mut cursor)?)); + } + SnapValue::List(deque) + } + TYPE_SORTED_SET => { + let count = format::read_u32(&mut cursor)?; + let mut members = Vec::with_capacity(count as usize); + for _ in 0..count { + let score = format::read_f64(&mut cursor)?; + let member_bytes = format::read_bytes(&mut cursor)?; + let member = String::from_utf8(member_bytes).map_err(|_| { + FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + "member is not valid utf-8", + )) + })?; + members.push((score, member)); + } + SnapValue::SortedSet(members) + } + TYPE_HASH => { + let count = format::read_u32(&mut cursor)?; + let mut map = HashMap::with_capacity(count as usize); + for _ in 0..count { + let field_bytes = format::read_bytes(&mut cursor)?; + let field = String::from_utf8(field_bytes).map_err(|_| { + FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + "hash field is not valid utf-8", + )) + })?; + let value_bytes = format::read_bytes(&mut cursor)?; + map.insert(field, Bytes::from(value_bytes)); + } + SnapValue::Hash(map) + } + TYPE_SET => { + let count = format::read_u32(&mut cursor)?; + let mut set = HashSet::with_capacity(count as usize); + for _ in 0..count { + let member_bytes = format::read_bytes(&mut cursor)?; + let member = String::from_utf8(member_bytes).map_err(|_| { + FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + "set member is not valid utf-8", + )) + })?; + set.insert(member); + } + SnapValue::Set(set) + } + _ => return Err(FormatError::UnknownTag(type_tag)), + }; + let expire_ms = format::read_i64(&mut cursor)?; + + let entry_key = String::from_utf8(key_bytes).map_err(|_| { + FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + "key is not valid utf-8", + )) + })?; + + self.read_so_far += 1; + Ok(Some(SnapEntry { + key: entry_key, + value, + expire_ms, + })) + } + /// Verifies the footer CRC32 after all entries have been read. /// Must be called after reading all entries. pub fn verify_footer(self) -> Result<(), FormatError> { @@ -609,4 +837,187 @@ mod tests { let p = snapshot_path(Path::new("/data"), 5); assert_eq!(p, PathBuf::from("/data/shard-5.snap")); } + + #[cfg(feature = "encryption")] + mod encrypted { + use super::*; + use crate::encryption::EncryptionKey; + + fn test_key() -> EncryptionKey { + EncryptionKey::from_bytes([0x42; 32]) + } + + #[test] + fn encrypted_snapshot_round_trip() { + let dir = temp_dir(); + let path = dir.path().join("enc.snap"); + let key = test_key(); + + let entries = vec![ + SnapEntry { + key: "hello".into(), + value: SnapValue::String(Bytes::from("world")), + expire_ms: -1, + }, + SnapEntry { + key: "ttl".into(), + value: SnapValue::String(Bytes::from("expiring")), + expire_ms: 5000, + }, + ]; + + { + let mut writer = + SnapshotWriter::create_encrypted(&path, 7, key.clone()).unwrap(); + for entry in &entries { + writer.write_entry(entry).unwrap(); + } + writer.finish().unwrap(); + } + + let mut reader = SnapshotReader::open_encrypted(&path, key).unwrap(); + assert_eq!(reader.shard_id, 7); + assert_eq!(reader.entry_count, 2); + + let mut got = Vec::new(); + while let Some(entry) = reader.read_entry().unwrap() { + got.push(entry); + } + assert_eq!(entries, got); + reader.verify_footer().unwrap(); + } + + #[test] + fn encrypted_snapshot_wrong_key_fails() { + let dir = temp_dir(); + let path = dir.path().join("enc_bad.snap"); + let key = test_key(); + let wrong_key = EncryptionKey::from_bytes([0xFF; 32]); + + { + let mut writer = SnapshotWriter::create_encrypted(&path, 0, key).unwrap(); + writer + .write_entry(&SnapEntry { + key: "k".into(), + value: SnapValue::String(Bytes::from("v")), + expire_ms: -1, + }) + .unwrap(); + writer.finish().unwrap(); + } + + let mut reader = SnapshotReader::open_encrypted(&path, wrong_key).unwrap(); + let err = reader.read_entry().unwrap_err(); + assert!(matches!(err, FormatError::DecryptionFailed)); + } + + #[test] + fn v2_snapshot_readable_with_encryption_key() { + let dir = temp_dir(); + let path = dir.path().join("v2.snap"); + let key = test_key(); + + { + let mut writer = SnapshotWriter::create(&path, 0).unwrap(); + writer + .write_entry(&SnapEntry { + key: "k".into(), + value: SnapValue::String(Bytes::from("v")), + expire_ms: -1, + }) + .unwrap(); + writer.finish().unwrap(); + } + + let mut reader = SnapshotReader::open_encrypted(&path, key).unwrap(); + let entry = reader.read_entry().unwrap().unwrap(); + assert_eq!(entry.key, "k"); + reader.verify_footer().unwrap(); + } + + #[test] + fn v3_snapshot_without_key_returns_error() { + let dir = temp_dir(); + let path = dir.path().join("v3_nokey.snap"); + let key = test_key(); + + { + let mut writer = SnapshotWriter::create_encrypted(&path, 0, key).unwrap(); + writer + .write_entry(&SnapEntry { + key: "k".into(), + value: SnapValue::String(Bytes::from("v")), + expire_ms: -1, + }) + .unwrap(); + writer.finish().unwrap(); + } + + let result = SnapshotReader::open(&path); + assert!(matches!(result, Err(FormatError::EncryptionRequired))); + } + + #[test] + fn encrypted_snapshot_with_all_types() { + let dir = temp_dir(); + let path = dir.path().join("enc_types.snap"); + let key = test_key(); + + let mut deque = VecDeque::new(); + deque.push_back(Bytes::from("a")); + deque.push_back(Bytes::from("b")); + + let mut hash = HashMap::new(); + hash.insert("f1".into(), Bytes::from("v1")); + + let mut set = HashSet::new(); + set.insert("m1".into()); + set.insert("m2".into()); + + let entries = vec![ + SnapEntry { + key: "str".into(), + value: SnapValue::String(Bytes::from("val")), + expire_ms: -1, + }, + SnapEntry { + key: "list".into(), + value: SnapValue::List(deque), + expire_ms: 1000, + }, + SnapEntry { + key: "zset".into(), + value: SnapValue::SortedSet(vec![(1.0, "a".into()), (2.0, "b".into())]), + expire_ms: -1, + }, + SnapEntry { + key: "hash".into(), + value: SnapValue::Hash(hash), + expire_ms: -1, + }, + SnapEntry { + key: "set".into(), + value: SnapValue::Set(set), + expire_ms: -1, + }, + ]; + + { + let mut writer = + SnapshotWriter::create_encrypted(&path, 0, key.clone()).unwrap(); + for entry in &entries { + writer.write_entry(entry).unwrap(); + } + writer.finish().unwrap(); + } + + let mut reader = SnapshotReader::open_encrypted(&path, key).unwrap(); + let mut got = Vec::new(); + while let Some(entry) = reader.read_entry().unwrap() { + got.push(entry); + } + assert_eq!(entries, got); + reader.verify_footer().unwrap(); + } + } } From 65fa5cead6b9a5a4ec6a3ea5d856871720f3e8d2 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 09:36:37 -0500 Subject: [PATCH 4/7] feat: pass encryption key through recovery recover_shard_encrypted() accepts an optional encryption key and forwards it to snapshot reader and AOF reader. handles both v2 (plaintext) and v3 (encrypted) files transparently. --- crates/ember-persistence/src/recovery.rs | 70 +++++++++++++++++++++++- 1 file changed, 67 insertions(+), 3 deletions(-) diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index dba0e0a2..3e6e39d2 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -68,6 +68,31 @@ pub struct RecoveryResult { /// Returns a list of live entries to restore into the keyspace. /// Entries whose TTL expired during downtime are silently skipped. pub fn recover_shard(data_dir: &Path, shard_id: u16) -> RecoveryResult { + recover_shard_inner(data_dir, shard_id, { + #[cfg(feature = "encryption")] + { None } + #[cfg(not(feature = "encryption"))] + { () } + }) +} + +/// Recovers a shard's state with an encryption key for decrypting +/// v3 persistence files. Also handles plaintext v2 files transparently. +#[cfg(feature = "encryption")] +pub fn recover_shard_encrypted( + data_dir: &Path, + shard_id: u16, + key: crate::encryption::EncryptionKey, +) -> RecoveryResult { + recover_shard_inner(data_dir, shard_id, Some(key)) +} + +fn recover_shard_inner( + data_dir: &Path, + shard_id: u16, + #[cfg(feature = "encryption")] encryption_key: Option, + #[cfg(not(feature = "encryption"))] _: (), +) -> RecoveryResult { // Track remaining TTL in ms (-1 = no expiry, 0+ = remaining ms) let mut map: HashMap = HashMap::new(); let mut loaded_snapshot = false; @@ -76,7 +101,17 @@ pub fn recover_shard(data_dir: &Path, shard_id: u16) -> RecoveryResult { // step 1: load snapshot let snap_path = snapshot::snapshot_path(data_dir, shard_id); if snap_path.exists() { - match load_snapshot(&snap_path) { + let result = { + #[cfg(feature = "encryption")] + { + load_snapshot(&snap_path, encryption_key.as_ref()) + } + #[cfg(not(feature = "encryption"))] + { + load_snapshot(&snap_path) + } + }; + match result { Ok(entries) => { for (key, value, ttl_ms) in entries { map.insert(key, (RecoveredValue::from(value), ttl_ms)); @@ -92,7 +127,17 @@ pub fn recover_shard(data_dir: &Path, shard_id: u16) -> RecoveryResult { // step 2: replay AOF let aof_path = aof::aof_path(data_dir, shard_id); if aof_path.exists() { - match replay_aof(&aof_path, &mut map) { + let result = { + #[cfg(feature = "encryption")] + { + replay_aof(&aof_path, &mut map, encryption_key.as_ref()) + } + #[cfg(not(feature = "encryption"))] + { + replay_aof(&aof_path, &mut map) + } + }; + match result { Ok(count) => { if count > 0 { replayed_aof = true; @@ -131,8 +176,19 @@ pub fn recover_shard(data_dir: &Path, shard_id: u16) -> RecoveryResult { /// Loads entries from a snapshot file. /// Returns (key, value, ttl_ms) where ttl_ms is -1 for no expiry. -fn load_snapshot(path: &Path) -> Result, FormatError> { +fn load_snapshot( + path: &Path, + #[cfg(feature = "encryption")] encryption_key: Option<&crate::encryption::EncryptionKey>, +) -> Result, FormatError> { + #[cfg(feature = "encryption")] + let mut reader = if let Some(key) = encryption_key { + SnapshotReader::open_encrypted(path, key.clone())? + } else { + SnapshotReader::open(path)? + }; + #[cfg(not(feature = "encryption"))] let mut reader = SnapshotReader::open(path)?; + let mut entries = Vec::new(); while let Some(entry) = reader.read_entry()? { @@ -168,7 +224,15 @@ fn apply_incr(map: &mut HashMap, key: String, del fn replay_aof( path: &Path, map: &mut HashMap, + #[cfg(feature = "encryption")] encryption_key: Option<&crate::encryption::EncryptionKey>, ) -> Result { + #[cfg(feature = "encryption")] + let mut reader = if let Some(key) = encryption_key { + AofReader::open_encrypted(path, key.clone())? + } else { + AofReader::open(path)? + }; + #[cfg(not(feature = "encryption"))] let mut reader = AofReader::open(path)?; let mut count = 0; From 1d1ddde2c3799a7a3f4144201806482832c70ef4 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 09:39:07 -0500 Subject: [PATCH 5/7] feat: forward encryption feature through emberkv-core adds `encryption` feature flag to emberkv-core that forwards to ember-persistence. ShardPersistenceConfig gains an optional encryption_key field (cfg-gated). recovery, AOF writer, and snapshot writer all use encrypted variants when a key is present. --- crates/ember-core/Cargo.toml | 3 +++ crates/ember-core/src/shard.rs | 49 +++++++++++++++++++++++++++++++--- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/crates/ember-core/Cargo.toml b/crates/ember-core/Cargo.toml index 6abda40a..352b9e02 100644 --- a/crates/ember-core/Cargo.toml +++ b/crates/ember-core/Cargo.toml @@ -9,6 +9,9 @@ keywords.workspace = true categories.workspace = true readme = "README.md" +[features] +encryption = ["ember-persistence/encryption"] + [lib] name = "ember_core" diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index 78c8dca8..3120c65a 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -41,6 +41,10 @@ pub struct ShardPersistenceConfig { pub append_only: bool, /// When to fsync the AOF file. pub fsync_policy: FsyncPolicy, + /// Optional encryption key for encrypting data at rest. + /// When set, AOF and snapshot files use the v3 encrypted format. + #[cfg(feature = "encryption")] + pub encryption_key: Option, } /// A protocol-agnostic command sent to a shard. @@ -392,6 +396,13 @@ async fn run_shard( // -- recovery -- if let Some(ref pcfg) = persistence { + #[cfg(feature = "encryption")] + let result = if let Some(ref key) = pcfg.encryption_key { + recovery::recover_shard_encrypted(&pcfg.data_dir, shard_id, key.clone()) + } else { + recovery::recover_shard(&pcfg.data_dir, shard_id) + }; + #[cfg(not(feature = "encryption"))] let result = recovery::recover_shard(&pcfg.data_dir, shard_id); let count = result.entries.len(); for entry in result.entries { @@ -425,7 +436,15 @@ async fn run_shard( let mut aof_writer: Option = match &persistence { Some(pcfg) if pcfg.append_only => { let path = ember_persistence::aof::aof_path(&pcfg.data_dir, shard_id); - match AofWriter::open(path) { + #[cfg(feature = "encryption")] + let result = if let Some(ref key) = pcfg.encryption_key { + AofWriter::open_encrypted(path, key.clone()) + } else { + AofWriter::open(path) + }; + #[cfg(not(feature = "encryption"))] + let result = AofWriter::open(path); + match result { Ok(w) => Some(w), Err(e) => { warn!(shard_id, "failed to open AOF writer: {e}"); @@ -956,7 +975,14 @@ fn handle_snapshot( }; let path = snapshot::snapshot_path(&pcfg.data_dir, shard_id); - match write_snapshot(keyspace, &path, shard_id) { + let result = write_snapshot( + keyspace, + &path, + shard_id, + #[cfg(feature = "encryption")] + pcfg.encryption_key.as_ref(), + ); + match result { Ok(count) => { info!(shard_id, entries = count, "snapshot written"); ShardResponse::Ok @@ -981,7 +1007,14 @@ fn handle_rewrite( }; let path = snapshot::snapshot_path(&pcfg.data_dir, shard_id); - match write_snapshot(keyspace, &path, shard_id) { + let result = write_snapshot( + keyspace, + &path, + shard_id, + #[cfg(feature = "encryption")] + pcfg.encryption_key.as_ref(), + ); + match result { Ok(count) => { // truncate AOF after successful snapshot if let Some(ref mut writer) = aof_writer { @@ -1004,7 +1037,15 @@ fn write_snapshot( keyspace: &Keyspace, path: &std::path::Path, shard_id: u16, + #[cfg(feature = "encryption")] encryption_key: Option<&ember_persistence::encryption::EncryptionKey>, ) -> Result { + #[cfg(feature = "encryption")] + let mut writer = if let Some(key) = encryption_key { + SnapshotWriter::create_encrypted(path, shard_id, key.clone())? + } else { + SnapshotWriter::create(path, shard_id)? + }; + #[cfg(not(feature = "encryption"))] let mut writer = SnapshotWriter::create(path, shard_id)?; let mut count = 0u32; @@ -1234,6 +1275,8 @@ mod tests { data_dir: dir.path().to_owned(), append_only: true, fsync_policy: FsyncPolicy::Always, + #[cfg(feature = "encryption")] + encryption_key: None, }; let config = ShardConfig { shard_id: 0, From 8f4bf34868922d378c3bb19445681875b9d0f4ba Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 09:40:20 -0500 Subject: [PATCH 6/7] feat: add --encryption-key-file CLI flag to ember-server the flag is only available when compiled with --features encryption. reads a 32-byte key (raw or hex) from the given file path and passes it through to shard persistence config. validates that persistence is enabled when encryption is configured. logs when encryption is active. --- crates/ember-server/Cargo.toml | 1 + crates/ember-server/src/main.rs | 36 +++++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/crates/ember-server/Cargo.toml b/crates/ember-server/Cargo.toml index b6d1c713..14ee2857 100644 --- a/crates/ember-server/Cargo.toml +++ b/crates/ember-server/Cargo.toml @@ -12,6 +12,7 @@ readme = "README.md" [features] default = ["jemalloc"] jemalloc = ["tikv-jemallocator"] +encryption = ["emberkv-core/encryption", "ember-persistence/encryption"] [dependencies] bytes = { workspace = true } diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index 10351fd2..ace4dab6 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -115,6 +115,15 @@ struct Args { #[arg(long, default_value = "no", env = "EMBER_TLS_AUTH_CLIENTS")] tls_auth_clients: String, + // -- encryption at rest -- + /// path to a 32-byte key file for encrypting AOF and snapshot files. + /// accepts 32 raw bytes or 64 hex characters. when set, new persistence + /// files use AES-256-GCM encryption. existing plaintext files are read + /// normally and migrated on the next BGREWRITEAOF/BGSAVE. + #[cfg(feature = "encryption")] + #[arg(long, env = "EMBER_ENCRYPTION_KEY_FILE")] + encryption_key_file: Option, + // -- cluster options -- /// enable cluster mode with gossip-based discovery and slot routing #[arg(long, env = "EMBER_CLUSTER_ENABLED")] @@ -196,6 +205,26 @@ async fn main() { std::process::exit(1); } + // load encryption key if configured + #[cfg(feature = "encryption")] + let encryption_key = if let Some(ref key_path) = args.encryption_key_file { + match ember_persistence::encryption::EncryptionKey::from_file(key_path) { + Ok(key) => Some(key), + Err(e) => { + eprintln!("failed to load encryption key: {e}"); + std::process::exit(1); + } + } + } else { + None + }; + + #[cfg(feature = "encryption")] + if encryption_key.is_some() && !args.appendonly && args.data_dir.is_none() { + eprintln!("--encryption-key-file requires --data-dir and --appendonly"); + std::process::exit(1); + } + // build persistence config if data-dir is set or appendonly is enabled let persistence = if args.appendonly || args.data_dir.is_some() { let data_dir = args.data_dir.unwrap_or_else(|| { @@ -215,6 +244,8 @@ async fn main() { data_dir, append_only: args.appendonly, fsync_policy, + #[cfg(feature = "encryption")] + encryption_key, }) } else { None @@ -237,6 +268,11 @@ async fn main() { fsync = ?p.fsync_policy, "persistence enabled" ); + + #[cfg(feature = "encryption")] + if p.encryption_key.is_some() { + info!("encryption at rest enabled (AES-256-GCM)"); + } } // install prometheus metrics exporter if --metrics-port is set From 187ffa78d7819344863c2034e0867e788881af18 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 09:58:24 -0500 Subject: [PATCH 7/7] fix: resolve clippy and formatting warnings for encryption feature refactor cfg-gated code to avoid unit_arg/unused_unit clippy lints: - aof: extract open_persistence_file helper, separate open/open_encrypted - snapshot: extract open_tmp helper, separate create/create_encrypted - recovery: use EncryptionKeyRef type alias for cleaner cfg gating - encryption: use io::Error::other instead of deprecated constructor --- Cargo.lock | 141 ++++++++++++++++++++- crates/ember-core/src/shard.rs | 4 +- crates/ember-persistence/src/aof.rs | 69 +++++----- crates/ember-persistence/src/encryption.rs | 29 +++-- crates/ember-persistence/src/recovery.rs | 53 +++----- crates/ember-persistence/src/snapshot.rs | 87 ++++++------- 6 files changed, 256 insertions(+), 127 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b8b07f67..dcffebed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,41 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common", + "generic-array", +] + +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "aes-gcm" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" +dependencies = [ + "aead", + "aes", + "cipher", + "ctr", + "ghash", + "subtle", +] + [[package]] name = "ahash" version = "0.7.8" @@ -311,6 +346,16 @@ dependencies = [ "half", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", +] + [[package]] name = "clap" version = "4.5.57" @@ -400,6 +445,15 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -478,6 +532,26 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "typenum", +] + +[[package]] +name = "ctr" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" +dependencies = [ + "cipher", +] + [[package]] name = "dashmap" version = "6.1.0" @@ -581,8 +655,10 @@ dependencies = [ name = "ember-persistence" version = "0.4.1" dependencies = [ + "aes-gcm", "bytes", "crc32fast", + "rand 0.9.2", "tempfile", "thiserror 2.0.18", "tracing", @@ -699,7 +775,7 @@ checksum = "0ce92ff622d6dadf7349484f42c93271a0d49b7cc4d466a936405bacbe10aa78" dependencies = [ "cfg-if", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -821,6 +897,16 @@ dependencies = [ "slab", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -844,6 +930,16 @@ dependencies = [ "wasip2", ] +[[package]] +name = "ghash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" +dependencies = [ + "opaque-debug", + "polyval", +] + [[package]] name = "h2" version = "0.4.13" @@ -1064,6 +1160,15 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "ipnet" version = "2.11.0" @@ -1301,6 +1406,12 @@ version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + [[package]] name = "openraft" version = "0.9.21" @@ -1441,6 +1552,18 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "polyval" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" +dependencies = [ + "cfg-if", + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "portable-atomic" version = "1.13.1" @@ -2356,6 +2479,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typenum" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" + [[package]] name = "unicode-ident" version = "1.0.22" @@ -2380,6 +2509,16 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common", + "subtle", +] + [[package]] name = "untrusted" version = "0.9.0" diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index 3120c65a..afc10fa2 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -1037,7 +1037,9 @@ fn write_snapshot( keyspace: &Keyspace, path: &std::path::Path, shard_id: u16, - #[cfg(feature = "encryption")] encryption_key: Option<&ember_persistence::encryption::EncryptionKey>, + #[cfg(feature = "encryption")] encryption_key: Option< + &ember_persistence::encryption::EncryptionKey, + >, ) -> Result { #[cfg(feature = "encryption")] let mut writer = if let Some(key) = encryption_key { diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index b02aaa53..857ae2a9 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -445,11 +445,22 @@ impl AofWriter { /// Opens (or creates) an AOF file. If the file is new, writes the header. /// If the file already exists, appends to it. pub fn open(path: impl Into) -> Result { - Self::open_inner(path.into(), { + let path = path.into(); + let exists = path.exists() && fs::metadata(&path).map(|m| m.len() > 0).unwrap_or(false); + + let file = open_persistence_file(&path)?; + let mut writer = BufWriter::new(file); + + if !exists { + format::write_header(&mut writer, format::AOF_MAGIC)?; + writer.flush()?; + } + + Ok(Self { + writer, + path, #[cfg(feature = "encryption")] - { None } - #[cfg(not(feature = "encryption"))] - { () } + encryption_key: None, }) } @@ -463,47 +474,25 @@ impl AofWriter { path: impl Into, key: crate::encryption::EncryptionKey, ) -> Result { - Self::open_inner(path.into(), Some(key)) - } - - fn open_inner( - path: PathBuf, - #[cfg(feature = "encryption")] encryption_key: Option, - #[cfg(not(feature = "encryption"))] _: (), - ) -> Result { + let path = path.into(); let exists = path.exists() && fs::metadata(&path).map(|m| m.len() > 0).unwrap_or(false); - let mut opts = OpenOptions::new(); - opts.create(true).append(true); - #[cfg(unix)] - { - use std::os::unix::fs::OpenOptionsExt; - opts.mode(0o600); - } - let file = opts.open(&path)?; + let file = open_persistence_file(&path)?; let mut writer = BufWriter::new(file); if !exists { - #[cfg(feature = "encryption")] - if encryption_key.is_some() { - format::write_header_versioned( - &mut writer, - format::AOF_MAGIC, - format::FORMAT_VERSION_ENCRYPTED, - )?; - } else { - format::write_header(&mut writer, format::AOF_MAGIC)?; - } - #[cfg(not(feature = "encryption"))] - format::write_header(&mut writer, format::AOF_MAGIC)?; + format::write_header_versioned( + &mut writer, + format::AOF_MAGIC, + format::FORMAT_VERSION_ENCRYPTED, + )?; writer.flush()?; } Ok(Self { writer, path, - #[cfg(feature = "encryption")] - encryption_key, + encryption_key: Some(key), }) } @@ -807,6 +796,18 @@ impl AofReader { } } +/// Opens a persistence file with create+append and restrictive permissions. +fn open_persistence_file(path: &Path) -> Result { + let mut opts = OpenOptions::new(); + opts.create(true).append(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + opts.mode(0o600); + } + Ok(opts.open(path)?) +} + /// Returns the AOF file path for a given shard in a data directory. pub fn aof_path(data_dir: &Path, shard_id: u16) -> PathBuf { data_dir.join(format!("shard-{shard_id}.aof")) diff --git a/crates/ember-persistence/src/encryption.rs b/crates/ember-persistence/src/encryption.rs index b80e2aa3..ff6e895c 100644 --- a/crates/ember-persistence/src/encryption.rs +++ b/crates/ember-persistence/src/encryption.rs @@ -53,7 +53,10 @@ impl EncryptionKey { let data = std::fs::read(path).map_err(|e| { FormatError::Io(io::Error::new( e.kind(), - format!("failed to read encryption key file '{}': {e}", path.display()), + format!( + "failed to read encryption key file '{}': {e}", + path.display() + ), )) })?; @@ -107,16 +110,16 @@ impl EncryptionKey { /// /// Returns `(nonce, ciphertext)` where ciphertext includes the 16-byte /// auth tag appended by AES-GCM. -pub fn encrypt_record(key: &EncryptionKey, plaintext: &[u8]) -> Result<([u8; NONCE_SIZE], Vec), FormatError> { +pub fn encrypt_record( + key: &EncryptionKey, + plaintext: &[u8], +) -> Result<([u8; NONCE_SIZE], Vec), FormatError> { let cipher = Aes256Gcm::new(key.as_bytes().into()); let nonce = Aes256Gcm::generate_nonce(&mut OsRng); - let ciphertext = cipher.encrypt(&nonce, plaintext).map_err(|e| { - FormatError::Io(io::Error::new( - io::ErrorKind::Other, - format!("encryption failed: {e}"), - )) - })?; + let ciphertext = cipher + .encrypt(&nonce, plaintext) + .map_err(|e| FormatError::Io(io::Error::other(format!("encryption failed: {e}"))))?; let mut nonce_bytes = [0u8; NONCE_SIZE]; nonce_bytes.copy_from_slice(&nonce); @@ -128,11 +131,17 @@ pub fn encrypt_record(key: &EncryptionKey, plaintext: &[u8]) -> Result<([u8; NON /// The ciphertext must include the 16-byte auth tag (as produced by /// [`encrypt_record`]). Returns `DecryptionFailed` if the key is wrong /// or the data has been tampered with. -pub fn decrypt_record(key: &EncryptionKey, nonce: &[u8; NONCE_SIZE], ciphertext: &[u8]) -> Result, FormatError> { +pub fn decrypt_record( + key: &EncryptionKey, + nonce: &[u8; NONCE_SIZE], + ciphertext: &[u8], +) -> Result, FormatError> { let cipher = Aes256Gcm::new(key.as_bytes().into()); let nonce = Nonce::from_slice(nonce); - cipher.decrypt(nonce, ciphertext).map_err(|_| FormatError::DecryptionFailed) + cipher + .decrypt(nonce, ciphertext) + .map_err(|_| FormatError::DecryptionFailed) } #[cfg(test)] diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index 3e6e39d2..c0900764 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -18,6 +18,14 @@ use crate::aof::{self, AofReader, AofRecord}; use crate::format::FormatError; use crate::snapshot::{self, SnapValue, SnapshotReader}; +/// Type alias for an optional encryption key reference. When the +/// `encryption` feature is disabled, this is always `Option<&()>` — +/// always `None` — and all encryption branches compile away. +#[cfg(feature = "encryption")] +type EncryptionKeyRef<'a> = &'a crate::encryption::EncryptionKey; +#[cfg(not(feature = "encryption"))] +type EncryptionKeyRef<'a> = &'a (); + /// The value of a recovered entry. #[derive(Debug, Clone)] pub enum RecoveredValue { @@ -68,12 +76,7 @@ pub struct RecoveryResult { /// Returns a list of live entries to restore into the keyspace. /// Entries whose TTL expired during downtime are silently skipped. pub fn recover_shard(data_dir: &Path, shard_id: u16) -> RecoveryResult { - recover_shard_inner(data_dir, shard_id, { - #[cfg(feature = "encryption")] - { None } - #[cfg(not(feature = "encryption"))] - { () } - }) + recover_shard_impl(data_dir, shard_id, None) } /// Recovers a shard's state with an encryption key for decrypting @@ -84,14 +87,16 @@ pub fn recover_shard_encrypted( shard_id: u16, key: crate::encryption::EncryptionKey, ) -> RecoveryResult { - recover_shard_inner(data_dir, shard_id, Some(key)) + recover_shard_impl(data_dir, shard_id, Some(&key)) } -fn recover_shard_inner( +/// Shared implementation. When encryption is not compiled in, the key +/// parameter is always `None` and all encryption branches are dead code +/// that the compiler will remove. +fn recover_shard_impl( data_dir: &Path, shard_id: u16, - #[cfg(feature = "encryption")] encryption_key: Option, - #[cfg(not(feature = "encryption"))] _: (), + #[allow(unused_variables)] encryption_key: Option>, ) -> RecoveryResult { // Track remaining TTL in ms (-1 = no expiry, 0+ = remaining ms) let mut map: HashMap = HashMap::new(); @@ -101,17 +106,7 @@ fn recover_shard_inner( // step 1: load snapshot let snap_path = snapshot::snapshot_path(data_dir, shard_id); if snap_path.exists() { - let result = { - #[cfg(feature = "encryption")] - { - load_snapshot(&snap_path, encryption_key.as_ref()) - } - #[cfg(not(feature = "encryption"))] - { - load_snapshot(&snap_path) - } - }; - match result { + match load_snapshot(&snap_path, encryption_key) { Ok(entries) => { for (key, value, ttl_ms) in entries { map.insert(key, (RecoveredValue::from(value), ttl_ms)); @@ -127,17 +122,7 @@ fn recover_shard_inner( // step 2: replay AOF let aof_path = aof::aof_path(data_dir, shard_id); if aof_path.exists() { - let result = { - #[cfg(feature = "encryption")] - { - replay_aof(&aof_path, &mut map, encryption_key.as_ref()) - } - #[cfg(not(feature = "encryption"))] - { - replay_aof(&aof_path, &mut map) - } - }; - match result { + match replay_aof(&aof_path, &mut map, encryption_key) { Ok(count) => { if count > 0 { replayed_aof = true; @@ -178,7 +163,7 @@ fn recover_shard_inner( /// Returns (key, value, ttl_ms) where ttl_ms is -1 for no expiry. fn load_snapshot( path: &Path, - #[cfg(feature = "encryption")] encryption_key: Option<&crate::encryption::EncryptionKey>, + #[allow(unused_variables)] encryption_key: Option>, ) -> Result, FormatError> { #[cfg(feature = "encryption")] let mut reader = if let Some(key) = encryption_key { @@ -224,7 +209,7 @@ fn apply_incr(map: &mut HashMap, key: String, del fn replay_aof( path: &Path, map: &mut HashMap, - #[cfg(feature = "encryption")] encryption_key: Option<&crate::encryption::EncryptionKey>, + #[allow(unused_variables)] encryption_key: Option>, ) -> Result { #[cfg(feature = "encryption")] let mut reader = if let Some(key) = encryption_key { diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index 535fda5b..b7995f5d 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -84,11 +84,23 @@ impl SnapshotWriter { /// Creates a new snapshot writer. The file won't appear at `path` /// until [`Self::finish`] is called successfully. pub fn create(path: impl Into, shard_id: u16) -> Result { - Self::create_inner(path.into(), shard_id, { + let final_path = path.into(); + let (tmp_path, writer) = Self::open_tmp(&final_path)?; + let mut writer = BufWriter::new(writer); + + format::write_header(&mut writer, format::SNAP_MAGIC)?; + format::write_u16(&mut writer, shard_id)?; + format::write_u32(&mut writer, 0)?; + + Ok(Self { + final_path, + tmp_path, + writer, + hasher: crc32fast::Hasher::new(), + count: 0, + finished: false, #[cfg(feature = "encryption")] - { None } - #[cfg(not(feature = "encryption"))] - { () } + encryption_key: None, }) } @@ -99,46 +111,16 @@ impl SnapshotWriter { shard_id: u16, key: crate::encryption::EncryptionKey, ) -> Result { - Self::create_inner(path.into(), shard_id, Some(key)) - } - - fn create_inner( - final_path: PathBuf, - shard_id: u16, - #[cfg(feature = "encryption")] encryption_key: Option, - #[cfg(not(feature = "encryption"))] _: (), - ) -> Result { - let tmp_path = final_path.with_extension("snap.tmp"); - - let file = { - let mut opts = OpenOptions::new(); - opts.write(true).create(true).truncate(true); - #[cfg(unix)] - { - use std::os::unix::fs::OpenOptionsExt; - opts.mode(0o600); - } - opts.open(&tmp_path)? - }; + let final_path = path.into(); + let (tmp_path, file) = Self::open_tmp(&final_path)?; let mut writer = BufWriter::new(file); - // write header: magic + version + shard_id + placeholder entry count - #[cfg(feature = "encryption")] - if encryption_key.is_some() { - format::write_header_versioned( - &mut writer, - format::SNAP_MAGIC, - format::FORMAT_VERSION_ENCRYPTED, - )?; - } else { - format::write_header(&mut writer, format::SNAP_MAGIC)?; - } - #[cfg(not(feature = "encryption"))] - format::write_header(&mut writer, format::SNAP_MAGIC)?; - + format::write_header_versioned( + &mut writer, + format::SNAP_MAGIC, + format::FORMAT_VERSION_ENCRYPTED, + )?; format::write_u16(&mut writer, shard_id)?; - // entry count — we'll seek back and update, or just write it now - // and track. since we're streaming, write 0 and update after. format::write_u32(&mut writer, 0)?; Ok(Self { @@ -148,11 +130,24 @@ impl SnapshotWriter { hasher: crc32fast::Hasher::new(), count: 0, finished: false, - #[cfg(feature = "encryption")] - encryption_key, + encryption_key: Some(key), }) } + /// Opens the temp file for writing. + fn open_tmp(final_path: &Path) -> Result<(PathBuf, File), FormatError> { + let tmp_path = final_path.with_extension("snap.tmp"); + let mut opts = OpenOptions::new(); + opts.write(true).create(true).truncate(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + opts.mode(0o600); + } + let file = opts.open(&tmp_path)?; + Ok((tmp_path, file)) + } + /// Writes a single entry to the snapshot. /// /// When encrypted, each entry is written as `[nonce: 12B][len: 4B][ciphertext]`. @@ -867,8 +862,7 @@ mod tests { ]; { - let mut writer = - SnapshotWriter::create_encrypted(&path, 7, key.clone()).unwrap(); + let mut writer = SnapshotWriter::create_encrypted(&path, 7, key.clone()).unwrap(); for entry in &entries { writer.write_entry(entry).unwrap(); } @@ -1003,8 +997,7 @@ mod tests { ]; { - let mut writer = - SnapshotWriter::create_encrypted(&path, 0, key.clone()).unwrap(); + let mut writer = SnapshotWriter::create_encrypted(&path, 0, key.clone()).unwrap(); for entry in &entries { writer.write_entry(entry).unwrap(); }