use crate::segment::error::{Result, SegmentError};
use aes_gcm::aead::{Aead, KeyInit};
use aes_gcm::{Aes256Gcm, Nonce};
pub(crate) const NONCE_SIZE: usize = 12;
fn fill_random(dest: &mut [u8]) {
use rand::TryRng;
rand::rngs::SysRng
.try_fill_bytes(dest)
.expect("the operating system random source must be available");
}
pub(crate) const TAG_SIZE: usize = 16;
#[derive(Clone)]
pub struct FieldEncryptionKey {
pub key_id: String,
key: [u8; 32],
}
impl FieldEncryptionKey {
#[must_use]
pub fn new(key_id: impl Into<String>, key_bytes: [u8; 32]) -> Self {
Self {
key_id: key_id.into(),
key: key_bytes,
}
}
#[must_use]
pub fn key_bytes(&self) -> &[u8; 32] {
&self.key
}
}
impl std::fmt::Debug for FieldEncryptionKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FieldEncryptionKey")
.field("key_id", &self.key_id)
.field("key", &"[REDACTED]")
.finish()
}
}
#[derive(Debug, Clone, Default)]
pub struct FieldEncryptionConfig {
pub columns: std::collections::HashMap<String, FieldEncryptionKey>,
}
impl FieldEncryptionConfig {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn encrypt_column(&mut self, column_name: impl Into<String>, key: FieldEncryptionKey) {
self.columns.insert(column_name.into(), key);
}
#[must_use]
pub fn is_encrypted(&self, column_name: &str) -> bool {
self.columns.contains_key(column_name)
}
#[must_use]
pub fn key_for(&self, column_name: &str) -> Option<&FieldEncryptionKey> {
self.columns.get(column_name)
}
}
pub trait FieldKeyProvider: Send + Sync {
fn get_key(&self, key_id: &str) -> Option<[u8; 32]>;
}
#[derive(Debug, Clone, Default)]
pub struct StaticKeyProvider {
keys: std::collections::HashMap<String, [u8; 32]>,
}
impl StaticKeyProvider {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn add_key(&mut self, key_id: impl Into<String>, key_bytes: [u8; 32]) {
self.keys.insert(key_id.into(), key_bytes);
}
}
impl FieldKeyProvider for StaticKeyProvider {
fn get_key(&self, key_id: &str) -> Option<[u8; 32]> {
self.keys.get(key_id).copied()
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct BlockContext<'a> {
pub column: &'a str,
pub segment_created_at: i64,
}
impl BlockContext<'_> {
fn associated_data(self) -> Vec<u8> {
let mut aad = Vec::with_capacity(self.column.len() + 32);
aad.extend_from_slice(b"chronix-field-v1");
aad.extend_from_slice(&(self.column.len() as u32).to_le_bytes());
aad.extend_from_slice(self.column.as_bytes());
aad.extend_from_slice(&self.segment_created_at.to_le_bytes());
aad
}
}
pub(crate) fn encrypt_block(
plaintext: &[u8],
key: &[u8; 32],
context: BlockContext<'_>,
) -> Result<Vec<u8>> {
let cipher = Aes256Gcm::new_from_slice(key).map_err(|e| SegmentError::CorruptFile {
detail: format!("invalid encryption key: {e}"),
})?;
let mut nonce_bytes = [0u8; NONCE_SIZE];
fill_random(&mut nonce_bytes);
let nonce = Nonce::try_from(&nonce_bytes[..]).map_err(|_| SegmentError::CorruptFile {
detail: "nonce must be 12 bytes".into(),
})?;
let aad = context.associated_data();
let ciphertext = cipher
.encrypt(
&nonce,
aes_gcm::aead::Payload {
msg: plaintext,
aad: &aad,
},
)
.map_err(|e| SegmentError::CorruptFile {
detail: format!("field encryption failed: {e}"),
})?;
let mut result = Vec::with_capacity(NONCE_SIZE + ciphertext.len());
result.extend_from_slice(&nonce);
result.extend_from_slice(&ciphertext);
Ok(result)
}
pub(crate) fn decrypt_block(
encrypted: &[u8],
key: &[u8; 32],
context: BlockContext<'_>,
) -> Result<Vec<u8>> {
if encrypted.len() < NONCE_SIZE + TAG_SIZE {
return Err(SegmentError::CorruptFile {
detail: format!(
"encrypted block too short: {} bytes (need at least {})",
encrypted.len(),
NONCE_SIZE + TAG_SIZE,
),
});
}
let cipher = Aes256Gcm::new_from_slice(key).map_err(|e| SegmentError::CorruptFile {
detail: format!("invalid decryption key: {e}"),
})?;
let nonce =
Nonce::try_from(&encrypted[..NONCE_SIZE]).map_err(|_| SegmentError::CorruptFile {
detail: "nonce must be 12 bytes".into(),
})?;
let ciphertext = &encrypted[NONCE_SIZE..];
let aad = context.associated_data();
cipher
.decrypt(
&nonce,
aes_gcm::aead::Payload {
msg: ciphertext,
aad: &aad,
},
)
.map_err(|e| SegmentError::CorruptFile {
detail: format!(
"field decryption failed (wrong key, moved block, or tampered data): {e}"
),
})
}
#[cfg(test)]
mod tests {
use super::*;
fn ctx() -> BlockContext<'static> {
BlockContext {
column: "value",
segment_created_at: 1_700_000_000_000_000_000,
}
}
fn test_key() -> [u8; 32] {
let mut key = [0u8; 32];
for (i, b) in key.iter_mut().enumerate() {
*b = i as u8;
}
key
}
#[test]
fn encrypt_decrypt_roundtrip() {
let key = test_key();
let plaintext = b"Hello, field-level encryption!";
let encrypted = encrypt_block(plaintext, &key, ctx()).unwrap();
assert_ne!(encrypted, plaintext);
assert!(encrypted.len() > plaintext.len()); let decrypted = decrypt_block(&encrypted, &key, ctx()).unwrap();
assert_eq!(decrypted, plaintext);
}
#[test]
fn different_nonces_produce_different_ciphertext() {
let key = test_key();
let plaintext = b"same data";
let enc1 = encrypt_block(plaintext, &key, ctx()).unwrap();
let enc2 = encrypt_block(plaintext, &key, ctx()).unwrap();
assert_ne!(&enc1[..NONCE_SIZE], &enc2[..NONCE_SIZE]);
}
#[test]
fn wrong_key_fails_decryption() {
let key1 = test_key();
let mut key2 = test_key();
key2[0] ^= 0xFF; let encrypted = encrypt_block(b"secret", &key1, ctx()).unwrap();
let result = decrypt_block(&encrypted, &key2, ctx());
assert!(result.is_err());
}
#[test]
fn tampered_ciphertext_fails() {
let key = test_key();
let mut encrypted = encrypt_block(b"important data", &key, ctx()).unwrap();
let last = encrypted.len() - 1;
encrypted[last] ^= 0x01;
let result = decrypt_block(&encrypted, &key, ctx());
assert!(result.is_err());
}
#[test]
fn too_short_encrypted_block_fails() {
let key = test_key();
let result = decrypt_block(&[0u8; 10], &key, ctx());
assert!(result.is_err());
}
#[test]
fn field_encryption_config_basic() {
let mut config = FieldEncryptionConfig::new();
assert!(!config.is_encrypted("temperature"));
let key = FieldEncryptionKey::new("key-1", test_key());
config.encrypt_column("temperature", key);
assert!(config.is_encrypted("temperature"));
assert!(!config.is_encrypted("humidity"));
assert_eq!(config.key_for("temperature").unwrap().key_id, "key-1");
}
#[test]
fn static_key_provider_lookup() {
let mut provider = StaticKeyProvider::new();
let key = test_key();
provider.add_key("key-1", key);
assert_eq!(provider.get_key("key-1"), Some(key));
assert_eq!(provider.get_key("key-2"), None);
}
#[test]
fn field_encryption_key_debug_redacts() {
let key = FieldEncryptionKey::new("key-1", test_key());
let debug = format!("{key:?}");
assert!(debug.contains("REDACTED"));
assert!(!debug.contains(&format!("{:?}", test_key())));
}
#[test]
fn encrypt_empty_block() {
let key = test_key();
let encrypted = encrypt_block(b"", &key, ctx()).unwrap();
assert_eq!(encrypted.len(), NONCE_SIZE + TAG_SIZE); let decrypted = decrypt_block(&encrypted, &key, ctx()).unwrap();
assert!(decrypted.is_empty());
}
#[test]
fn encrypt_large_block() {
let key = test_key();
let plaintext = vec![0xABu8; 1024 * 1024]; let encrypted = encrypt_block(&plaintext, &key, ctx()).unwrap();
let decrypted = decrypt_block(&encrypted, &key, ctx()).unwrap();
assert_eq!(decrypted, plaintext);
}
#[test]
fn a_block_moved_to_another_column_fails_to_authenticate() {
let key = test_key();
let salary = encrypt_block(
b"180000",
&key,
BlockContext {
column: "salary",
segment_created_at: 42,
},
)
.unwrap();
let err = decrypt_block(
&salary,
&key,
BlockContext {
column: "public_note",
segment_created_at: 42,
},
)
.expect_err("a block must not decrypt under another column's name");
assert!(err.to_string().contains("decryption failed"), "{err}");
}
#[test]
fn a_block_moved_to_another_segment_fails_to_authenticate() {
let key = test_key();
let block = encrypt_block(
b"readings",
&key,
BlockContext {
column: "value",
segment_created_at: 1_000,
},
)
.unwrap();
assert!(
decrypt_block(
&block,
&key,
BlockContext {
column: "value",
segment_created_at: 2_000,
},
)
.is_err(),
"a block must not decrypt inside a different segment"
);
let plain = decrypt_block(
&block,
&key,
BlockContext {
column: "value",
segment_created_at: 1_000,
},
)
.expect("the block reads in its own place");
assert_eq!(plain, b"readings");
}
}