armdb 0.7.0

sharded bitcask key-value storage optimized for NVMe
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228

use ring::aead::{self, AES_256_GCM, Aad, LessSafeKey, Nonce, UnboundKey};

use crate::error::{DbError, DbResult};

pub const TAG_LEN: usize = 16;
pub const PAGE_SIZE: usize = 4096;

/// Page-level AES-256-GCM cipher for encrypting data files.
///
/// Each 4096-byte page is encrypted independently with a nonce derived from
/// `(file_id, page_number)`. Uniqueness of `(key, nonce)` holds because the key
/// is a per-shard subkey (`derive`) built from the master key and a per-database
/// salt, so identical `(file_id, page)` values in different shards or databases
/// encrypt under different keys. The 16-byte GCM tag is stored in a `.tags` file.
pub struct PageCipher {
    key: LessSafeKey,
}

impl PageCipher {
    /// Create a cipher from a 32-byte AES-256 key.
    pub fn new(key_bytes: &[u8; 32]) -> DbResult<Self> {
        let unbound = UnboundKey::new(&AES_256_GCM, key_bytes)
            .map_err(|e| DbError::EncryptionError(format!("invalid key: {e}")))?;
        Ok(Self {
            key: LessSafeKey::new(unbound),
        })
    }

    /// Derive a per-shard cipher via `HKDF-SHA256(ikm = master, salt).expand([shard_id])`.
    ///
    /// Separates the AES-GCM key domain per shard, and — through the per-database
    /// `salt` — per database. The `(file_id, page)` nonce is unchanged and stays
    /// unique within each derived key, so no `(key, nonce)` pair is ever reused
    /// across shards or databases sharing one master key.
    pub fn derive(master: &[u8; 32], salt: &[u8; 16], shard_id: u8) -> DbResult<Self> {
        use ring::hkdf::{HKDF_SHA256, Salt};

        let prk = Salt::new(HKDF_SHA256, salt).extract(master);
        // Bind the info bytes to a local so the `Okm` borrow outlives `expand`.
        let shard_info = [shard_id];
        let info: [&[u8]; 1] = [&shard_info];
        // Expanding with `HKDF_SHA256` as the output KeyType yields 32 bytes
        // (SHA-256 output length) — exactly the AES-256 key size.
        let okm = prk
            .expand(&info, HKDF_SHA256)
            .map_err(|_| DbError::EncryptionError("hkdf expand failed".into()))?;
        let mut key = [0u8; 32];
        okm.fill(&mut key)
            .map_err(|_| DbError::EncryptionError("hkdf fill failed".into()))?;
        Self::new(&key)
    }

    /// Read a hex-encoded 32-byte key from an environment variable.
    pub fn key_from_env(var: &str) -> DbResult<[u8; 32]> {
        let hex = std::env::var(var)
            .map_err(|e| DbError::EncryptionError(format!("env var {var}: {e}")))?;
        let hex = hex.trim();
        if hex.len() != 64 {
            return Err(DbError::EncryptionError(format!(
                "env var {var}: expected 64 hex chars, got {}",
                hex.len()
            )));
        }
        let mut key = [0u8; 32];
        for i in 0..32 {
            key[i] = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16).map_err(|e| {
                DbError::EncryptionError(format!("env var {var}: invalid hex: {e}"))
            })?;
        }
        Ok(key)
    }

    /// Encrypt a 4096-byte page in-place, returning the 16-byte GCM tag.
    ///
    /// Nonce = `file_id(4 LE) || page_number(8 LE)`.
    pub fn encrypt_page(
        &self,
        file_id: u32,
        page_number: u64,
        page: &mut [u8],
    ) -> DbResult<[u8; TAG_LEN]> {
        debug_assert_eq!(page.len(), PAGE_SIZE);
        let nonce = make_page_nonce(file_id, page_number);
        let tag = self
            .key
            .seal_in_place_separate_tag(nonce, Aad::empty(), page)
            .map_err(|_| DbError::EncryptionError("encrypt failed".into()))?;
        let mut tag_bytes = [0u8; TAG_LEN];
        tag_bytes.copy_from_slice(tag.as_ref());
        Ok(tag_bytes)
    }

    /// Decrypt a 4096-byte page in-place using its GCM tag.
    ///
    /// Returns `Err(EncryptionError)` if authentication fails (wrong key or corrupted data).
    pub fn decrypt_page(
        &self,
        file_id: u32,
        page_number: u64,
        page: &mut [u8],
        tag: &[u8; TAG_LEN],
    ) -> DbResult<()> {
        debug_assert_eq!(page.len(), PAGE_SIZE);
        // ring::open_in_place expects ciphertext || tag concatenated.
        let mut buf = Vec::with_capacity(PAGE_SIZE + TAG_LEN);
        buf.extend_from_slice(page);
        buf.extend_from_slice(tag);
        let nonce = make_page_nonce(file_id, page_number);
        let plaintext = self
            .key
            .open_in_place(nonce, Aad::empty(), &mut buf)
            .map_err(|_| DbError::EncryptionError("decryption/authentication failed".into()))?;
        page.copy_from_slice(&plaintext[..PAGE_SIZE]);
        Ok(())
    }
}

fn make_page_nonce(file_id: u32, page_number: u64) -> Nonce {
    let mut nonce_bytes = [0u8; aead::NONCE_LEN];
    nonce_bytes[..4].copy_from_slice(&file_id.to_le_bytes());
    nonce_bytes[4..12].copy_from_slice(&page_number.to_le_bytes());
    Nonce::assume_unique_for_key(nonce_bytes)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encrypt_decrypt_round_trip() {
        let key = [0xABu8; 32];
        let cipher = PageCipher::new(&key).unwrap();

        let original = [42u8; PAGE_SIZE];
        let mut page = original;
        let tag = cipher.encrypt_page(1, 0, &mut page).unwrap();

        // Encrypted data should differ from original
        assert_ne!(page, original);

        cipher.decrypt_page(1, 0, &mut page, &tag).unwrap();
        assert_eq!(page, original);
    }

    #[test]
    fn test_wrong_key_fails() {
        let cipher1 = PageCipher::new(&[0xAA; 32]).unwrap();
        let cipher2 = PageCipher::new(&[0xBB; 32]).unwrap();

        let mut page = [1u8; PAGE_SIZE];
        let tag = cipher1.encrypt_page(1, 0, &mut page).unwrap();

        assert!(cipher2.decrypt_page(1, 0, &mut page, &tag).is_err());
    }

    #[test]
    fn test_wrong_nonce_fails() {
        let cipher = PageCipher::new(&[0xCC; 32]).unwrap();

        let mut page = [7u8; PAGE_SIZE];
        let tag = cipher.encrypt_page(1, 0, &mut page).unwrap();

        // Different page_number
        assert!(cipher.decrypt_page(1, 1, &mut page, &tag).is_err());
    }

    #[test]
    fn test_tampered_data_fails() {
        let cipher = PageCipher::new(&[0xDD; 32]).unwrap();

        let mut page = [9u8; PAGE_SIZE];
        let tag = cipher.encrypt_page(1, 0, &mut page).unwrap();

        page[0] ^= 0xFF; // Corrupt one byte
        assert!(cipher.decrypt_page(1, 0, &mut page, &tag).is_err());
    }

    #[test]
    fn derive_distinct_shard_ids_distinct_ciphertext() {
        let master = [0x11u8; 32];
        let salt = [0x22u8; 16];
        let c0 = PageCipher::derive(&master, &salt, 0).unwrap();
        let c1 = PageCipher::derive(&master, &salt, 1).unwrap();

        let plain = [7u8; PAGE_SIZE];
        let mut p0 = plain;
        let mut p1 = plain;
        let tag0 = c0.encrypt_page(1, 0, &mut p0).unwrap();
        let _tag1 = c1.encrypt_page(1, 0, &mut p1).unwrap();

        // Одинаковый (file_id, page) под разными под-ключами шардов → разный шифртекст.
        assert_ne!(p0, p1);
        // Ключ шарда 1 не расшифровывает страницу шарда 0.
        let mut cross = p0;
        assert!(c1.decrypt_page(1, 0, &mut cross, &tag0).is_err());
    }

    #[test]
    fn derive_distinct_salts_distinct_ciphertext() {
        let master = [0x33u8; 32];
        let c_a = PageCipher::derive(&master, &[0xA0; 16], 0).unwrap();
        let c_b = PageCipher::derive(&master, &[0xB0; 16], 0).unwrap();

        let plain = [9u8; PAGE_SIZE];
        let mut pa = plain;
        let mut pb = plain;
        c_a.encrypt_page(1, 0, &mut pa).unwrap();
        c_b.encrypt_page(1, 0, &mut pb).unwrap();

        // Один мастер-ключ + shard_id, но разные соли (разные коллекции/БД) → разный шифртекст.
        assert_ne!(pa, pb);
    }

    #[test]
    fn derive_is_deterministic() {
        let master = [0x44u8; 32];
        let salt = [0x55u8; 16];
        let c1 = PageCipher::derive(&master, &salt, 3).unwrap();
        let c2 = PageCipher::derive(&master, &salt, 3).unwrap();

        let plain = [1u8; PAGE_SIZE];
        let mut p = plain;
        let tag = c1.encrypt_page(2, 5, &mut p).unwrap();
        // Тот же (salt, shard_id) → тот же ключ → c2 расшифровывает то, что запечатал c1.
        c2.decrypt_page(2, 5, &mut p, &tag).unwrap();
        assert_eq!(p, plain);
    }
}