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/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..afc10fa2 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,17 @@ 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 +1277,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, 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/aof.rs b/crates/ember-persistence/src/aof.rs index 9a1e3921..857ae2a9 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,6 +437,8 @@ pub enum FsyncPolicy { pub struct AofWriter { writer: BufWriter, path: PathBuf, + #[cfg(feature = "encryption")] + encryption_key: Option, } impl AofWriter { @@ -443,14 +448,7 @@ impl AofWriter { 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 { @@ -458,12 +456,62 @@ impl AofWriter { writer.flush()?; } - Ok(Self { writer, path }) + Ok(Self { + writer, + path, + #[cfg(feature = "encryption")] + encryption_key: None, + }) + } + + /// 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 { + 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_versioned( + &mut writer, + format::AOF_MAGIC, + format::FORMAT_VERSION_ENCRYPTED, + )?; + writer.flush()?; + } + + Ok(Self { + writer, + path, + encryption_key: Some(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 +552,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 +575,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 +596,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 +636,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 +671,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(); @@ -639,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")) @@ -1121,4 +1290,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()); + } + } } diff --git a/crates/ember-persistence/src/encryption.rs b/crates/ember-persistence/src/encryption.rs new file mode 100644 index 00000000..ff6e895c --- /dev/null +++ b/crates/ember-persistence/src/encryption.rs @@ -0,0 +1,264 @@ +//! 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::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; diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index dba0e0a2..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,6 +76,28 @@ 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_impl(data_dir, shard_id, None) +} + +/// 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_impl(data_dir, shard_id, Some(&key)) +} + +/// 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, + #[allow(unused_variables)] encryption_key: Option>, +) -> 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 +106,7 @@ 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) { + match load_snapshot(&snap_path, encryption_key) { Ok(entries) => { for (key, value, ttl_ms) in entries { map.insert(key, (RecoveredValue::from(value), ttl_ms)); @@ -92,7 +122,7 @@ 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) { + match replay_aof(&aof_path, &mut map, encryption_key) { Ok(count) => { if count > 0 { replayed_aof = true; @@ -131,8 +161,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, + #[allow(unused_variables)] encryption_key: Option>, +) -> 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 +209,15 @@ fn apply_incr(map: &mut HashMap, key: String, del fn replay_aof( path: &Path, map: &mut HashMap, + #[allow(unused_variables)] encryption_key: Option>, ) -> 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; diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index 5bad3424..b7995f5d 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -76,6 +76,8 @@ 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 { @@ -83,25 +85,42 @@ impl SnapshotWriter { /// until [`Self::finish`] is called successfully. pub fn create(path: impl Into, shard_id: u16) -> Result { let final_path = path.into(); - let tmp_path = final_path.with_extension("snap.tmp"); + let (tmp_path, writer) = Self::open_tmp(&final_path)?; + let mut writer = BufWriter::new(writer); - 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)? - }; + 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")] + encryption_key: None, + }) + } + + /// Creates a new encrypted snapshot writer. + #[cfg(feature = "encryption")] + pub fn create_encrypted( + path: impl Into, + shard_id: u16, + key: crate::encryption::EncryptionKey, + ) -> Result { + 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 - 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 { @@ -111,10 +130,28 @@ impl SnapshotWriter { hasher: crc32fast::Hasher::new(), count: 0, finished: false, + 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]`. + /// 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 +193,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 +258,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 { @@ -216,6 +270,38 @@ impl SnapshotReader { let file = File::open(path.as_ref())?; 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)?; + + Ok(Self { + reader, + shard_id, + entry_count, + 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)?; @@ -227,6 +313,7 @@ impl SnapshotReader { read_so_far: 0, hasher: crc32fast::Hasher::new(), version, + encryption_key: Some(key), }) } @@ -236,6 +323,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 +444,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 +832,185 @@ 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(); + } + } } 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