light_magic/
encrypted.rs

1use aes_gcm::{
2    aead::{rand_core::RngCore, Aead, KeyInit, OsRng},
3    Aes256Gcm, Key, Nonce,
4};
5use argon2::{self, Argon2};
6use bincode::{self, Options};
7use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
8use serde::{de::DeserializeOwned, Deserialize, Serialize};
9use std::{
10    ffi::{OsStr, OsString},
11    fmt,
12    fs::{self, File},
13    io::{self, Read, Write},
14    ops::{Deref, DerefMut},
15    path::{Path, PathBuf},
16};
17use tracing::{error, info};
18use zeroize::Zeroize;
19
20const SALT_LEN: usize = 16;
21const NONCE_LEN: usize = 12;
22
23#[inline]
24fn bincode_cfg() -> impl bincode::Options {
25    bincode::DefaultOptions::new().with_fixint_encoding()
26}
27
28/// Encrypted envelope (serialized with bincode)
29#[derive(Serialize, Deserialize)]
30pub struct EncryptedData {
31    salt: [u8; SALT_LEN],
32    nonce: [u8; NONCE_LEN],
33    ciphertext: Vec<u8>,
34}
35
36/// Implement on your Database type.
37pub trait EncryptedDataStore: Default + Serialize {
38    fn open<P>(db: P, password: &str) -> io::Result<EncryptedAtomicDatabase<Self>>
39    where
40        P: AsRef<Path>,
41        Self: DeserializeOwned,
42    {
43        let db_path = db.as_ref();
44        if db_path.exists() {
45            EncryptedAtomicDatabase::load(db_path, password)
46        } else {
47            EncryptedAtomicDatabase::create_new(db_path, password)
48        }
49    }
50
51    fn create_from_str<P>(
52        data: &str,
53        path: P,
54        password: &str,
55    ) -> io::Result<EncryptedAtomicDatabase<Self>>
56    where
57        P: AsRef<Path>,
58        Self: DeserializeOwned,
59    {
60        let db_path = path.as_ref();
61        if !db_path.exists() {
62            EncryptedAtomicDatabase::create_from_str(data, path, password)
63        } else {
64            Err(io::Error::new(
65                io::ErrorKind::AlreadyExists,
66                "A file already exists at the provided path!",
67            ))
68        }
69    }
70
71    /// Loads after decrypting `EncryptedData` read from `file`.
72    fn load_encrypted(file: impl Read, key: &Key<Aes256Gcm>) -> io::Result<Self>
73    where
74        Self: DeserializeOwned,
75    {
76        let encrypted: EncryptedData = bincode_cfg().deserialize_from(file).map_err(|e| {
77            io::Error::new(
78                io::ErrorKind::InvalidData,
79                format!("Failed to deserialize encrypted data: {e}"),
80            )
81        })?;
82        Self::decrypt(&encrypted, key)
83    }
84
85    /// Saves current state as encrypted `EncryptedData` to `file`.
86    fn save_encrypted(
87        &self,
88        mut file: impl Write,
89        key: &Key<Aes256Gcm>,
90        salt: [u8; SALT_LEN],
91    ) -> io::Result<()> {
92        let encrypted = self.encrypt(key, salt)?;
93        bincode_cfg()
94            .serialize_into(&mut file, &encrypted)
95            .map_err(|e| {
96                io::Error::new(
97                    io::ErrorKind::Other,
98                    format!("Failed to write encrypted data to file: {e}"),
99                )
100            })
101    }
102
103    /// Encrypts the current data and returns the envelope.
104    fn encrypt(&self, key: &Key<Aes256Gcm>, salt: [u8; SALT_LEN]) -> io::Result<EncryptedData> {
105        // Non-allocating nonce
106        let mut nonce = [0u8; NONCE_LEN];
107        OsRng.fill_bytes(&mut nonce);
108
109        // Serialize plaintext
110        let plaintext = bincode_cfg().serialize(self).map_err(|e| {
111            io::Error::new(
112                io::ErrorKind::InvalidData,
113                format!("Serialization failed: {e}"),
114            )
115        })?;
116
117        let cipher = Aes256Gcm::new(key);
118        let ct = cipher
119            .encrypt(Nonce::from_slice(&nonce), plaintext.as_ref())
120            .map_err(|e| io::Error::new(io::ErrorKind::Other, format!("Encryption failed: {e}")))?;
121
122        Ok(EncryptedData {
123            salt,
124            nonce,
125            ciphertext: ct,
126        })
127    }
128
129    /// Decrypts the envelope into data.
130    fn decrypt(encrypted: &EncryptedData, key: &Key<Aes256Gcm>) -> io::Result<Self>
131    where
132        Self: DeserializeOwned,
133    {
134        let cipher = Aes256Gcm::new(key);
135        let pt = cipher
136            .decrypt(
137                Nonce::from_slice(&encrypted.nonce),
138                encrypted.ciphertext.as_ref(),
139            )
140            .map_err(|e| {
141                io::Error::new(
142                    io::ErrorKind::InvalidData,
143                    format!("Decryption failed: Incorrect password or corrupted data. {e}"),
144                )
145            })?;
146
147        let data = bincode_cfg().deserialize(&pt).map_err(|e| {
148            io::Error::new(
149                io::ErrorKind::InvalidData,
150                format!("Failed to deserialize decrypted data: {e}"),
151            )
152        })?;
153
154        Ok(data)
155    }
156}
157
158/// Derive a 32-byte key from the password and salt using Argon2id
159fn derive_key(password: &str, salt: &[u8]) -> io::Result<Key<Aes256Gcm>> {
160    let mut key = [0u8; 32];
161    Argon2::default()
162        .hash_password_into(password.as_bytes(), salt, &mut key)
163        .map_err(|_| io::Error::new(io::ErrorKind::Other, "Key derivation failed"))?;
164
165    let out = *Key::<Aes256Gcm>::from_slice(&key);
166    key.zeroize(); // wipe stack buffer
167    Ok(out)
168}
169
170/// Synchronized Wrapper, that automatically saves changes when path and tmp are defined
171pub struct EncryptedAtomicDatabase<T: EncryptedDataStore> {
172    path: PathBuf,
173    tmp: PathBuf,
174    data: RwLock<T>,
175    key: RwLock<Key<Aes256Gcm>>,
176    salt: RwLock<[u8; SALT_LEN]>,
177}
178
179impl<T: EncryptedDataStore + DeserializeOwned> EncryptedAtomicDatabase<T> {
180    /// Load the database with the provided password.
181    pub fn load<P: AsRef<Path>>(path: P, password: &str) -> io::Result<Self> {
182        let new_path = path.as_ref().to_path_buf();
183        let tmp = Self::tmp_path(&new_path)?;
184
185        // Read the whole envelope once; don't reopen
186        let mut file = File::open(&new_path)?;
187        let encrypted: EncryptedData = bincode_cfg().deserialize_from(&mut file).map_err(|e| {
188            io::Error::new(
189                io::ErrorKind::InvalidData,
190                format!("Failed to deserialize encrypted data: {e}"),
191            )
192        })?;
193        let key = derive_key(password, &encrypted.salt)?;
194        let data = T::decrypt(&encrypted, &key)?;
195
196        Ok(Self {
197            path: new_path,
198            tmp,
199            data: RwLock::new(data),
200            key: RwLock::new(key),
201            salt: RwLock::new(encrypted.salt),
202        })
203    }
204
205    /// Load from an in-memory string and persist if successful.
206    pub fn create_from_str<P: AsRef<Path>>(
207        data: &str,
208        path: P,
209        password: &str,
210    ) -> io::Result<Self> {
211        let new_path = path.as_ref().to_path_buf();
212        let tmp = Self::tmp_path(&new_path)?;
213
214        let encrypted: EncryptedData = bincode_cfg().deserialize(data.as_bytes()).map_err(|e| {
215            io::Error::new(
216                io::ErrorKind::InvalidData,
217                format!("Failed to deserialize encrypted data: {e}"),
218            )
219        })?;
220        let key = derive_key(password, &encrypted.salt)?;
221        let data = T::decrypt(&encrypted, &key)?;
222
223        atomic_write_encrypted(&tmp, &new_path, &data, &key, encrypted.salt)?;
224
225        Ok(Self {
226            path: new_path,
227            tmp,
228            data: RwLock::new(data),
229            key: RwLock::new(key),
230            salt: RwLock::new(encrypted.salt),
231        })
232    }
233
234    /// Create a new database and save it with the provided password.
235    pub fn create_new<P: AsRef<Path>>(path: P, password: &str) -> io::Result<Self> {
236        let new_path = path.as_ref().to_path_buf();
237        let tmp = Self::tmp_path(&new_path)?;
238
239        // Generate salt
240        let mut salt_bytes = [0u8; SALT_LEN];
241        OsRng.fill_bytes(&mut salt_bytes);
242        let key = derive_key(password, &salt_bytes)?;
243
244        let data = Default::default();
245        atomic_write_encrypted(&tmp, &new_path, &data, &key, salt_bytes)?;
246
247        Ok(Self {
248            path: new_path,
249            tmp,
250            data: RwLock::new(data),
251            key: RwLock::new(key),
252            salt: RwLock::new(salt_bytes),
253        })
254    }
255
256    /// Lock the database for reading.
257    pub fn read(&self) -> EncryptedAtomicDatabaseRead<'_, T> {
258        EncryptedAtomicDatabaseRead {
259            data: self.data.read(),
260        }
261    }
262
263    /// Lock the database for writing. Saves changes atomically on drop.
264    pub fn write(&self) -> EncryptedAtomicDatabaseWrite<'_, T> {
265        let key = *self.key.read();
266        let salt = *self.salt.read();
267        EncryptedAtomicDatabaseWrite {
268            path: self.path.as_ref(),
269            tmp: self.tmp.as_ref(),
270            data: self.data.write(),
271            key,
272            salt,
273        }
274    }
275
276    /// Change the password (re-encrypt with a new salt+key).
277    pub fn change_password(&self, new_password: &str) -> io::Result<()> {
278        let data_guard = self.data.read();
279
280        let mut new_salt = [0u8; SALT_LEN];
281        OsRng.fill_bytes(&mut new_salt);
282        let new_key = derive_key(new_password, &new_salt)?;
283
284        atomic_write_encrypted(&self.tmp, &self.path, &*data_guard, &new_key, new_salt)?;
285
286        {
287            let mut key_lock = self.key.write();
288            *key_lock = new_key;
289        }
290        {
291            let mut salt_lock = self.salt.write();
292            *salt_lock = new_salt;
293        }
294
295        Ok(())
296    }
297
298    fn tmp_path(path: &Path) -> io::Result<PathBuf> {
299        let mut tmp_name = OsString::from(".");
300        tmp_name.push(path.file_name().unwrap_or(OsStr::new("db")));
301        tmp_name.push("~");
302        let tmp = path.with_file_name(tmp_name);
303        if tmp.exists() {
304            error!(
305                "Found orphaned database temporary file '{tmp:?}'. The server has recently crashed or is already running. Delete this before continuing!"
306            );
307            return Err(io::Error::new(
308                io::ErrorKind::AlreadyExists,
309                "Orphaned temporary file exists",
310            ));
311        }
312        Ok(tmp)
313    }
314}
315
316/// Atomic write routine with encryption
317fn atomic_write_encrypted<T: EncryptedDataStore>(
318    tmp: &Path,
319    path: &Path,
320    data: &T,
321    key: &Key<Aes256Gcm>,
322    salt: [u8; SALT_LEN],
323) -> io::Result<()> {
324    {
325        let tmpfile = File::create(tmp)?;
326        data.save_encrypted(tmpfile, key, salt)?;
327    }
328    fs::rename(tmp, path)?;
329    Ok(())
330}
331
332impl<T: EncryptedDataStore> fmt::Debug for EncryptedAtomicDatabase<T> {
333    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
334        f.debug_struct("EncryptedAtomicDatabase")
335            .field("file", &self.path)
336            .finish()
337    }
338}
339
340impl<T: EncryptedDataStore> Drop for EncryptedAtomicDatabase<T> {
341    fn drop(&mut self) {
342        info!("Saving database");
343        let data_guard = self.data.read();
344        let key = self.key.read();
345        let salt = self.salt.read();
346        if let Err(e) = atomic_write_encrypted(&self.tmp, &self.path, &*data_guard, &key, *salt) {
347            error!("Failed to save database: {}", e);
348        }
349    }
350}
351
352pub struct EncryptedAtomicDatabaseRead<'a, T: EncryptedDataStore> {
353    data: RwLockReadGuard<'a, T>,
354}
355
356impl<'a, T: EncryptedDataStore> Deref for EncryptedAtomicDatabaseRead<'a, T> {
357    type Target = T;
358    fn deref(&self) -> &Self::Target {
359        &self.data
360    }
361}
362
363pub struct EncryptedAtomicDatabaseWrite<'a, T: EncryptedDataStore> {
364    tmp: &'a Path,
365    path: &'a Path,
366    data: RwLockWriteGuard<'a, T>,
367    key: Key<Aes256Gcm>,
368    salt: [u8; SALT_LEN],
369}
370
371impl<'a, T: EncryptedDataStore> Deref for EncryptedAtomicDatabaseWrite<'a, T> {
372    type Target = T;
373    fn deref(&self) -> &Self::Target {
374        &self.data
375    }
376}
377
378impl<'a, T: EncryptedDataStore> DerefMut for EncryptedAtomicDatabaseWrite<'a, T> {
379    fn deref_mut(&mut self) -> &mut Self::Target {
380        &mut self.data
381    }
382}
383
384impl<'a, T: EncryptedDataStore> Drop for EncryptedAtomicDatabaseWrite<'a, T> {
385    fn drop(&mut self) {
386        info!("Saving database");
387        if let Err(e) =
388            atomic_write_encrypted(self.tmp, self.path, &*self.data, &self.key, self.salt)
389        {
390            error!("Failed to save database: {}", e);
391        }
392    }
393}