light_magic/
encrypted.rs

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

light_magic/
encrypted.rs
