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#[derive(Serialize, Deserialize)]
30pub struct EncryptedData {
31 salt: [u8; SALT_LEN],
32 nonce: [u8; NONCE_LEN],
33 ciphertext: Vec<u8>,
34}
35
36pub trait EncryptedDataStore: Default + Serialize {
40 fn open<P>(db: P, password: &str) -> io::Result<EncryptedAtomicDatabase<Self>>
43 where
44 P: AsRef<Path>,
45 Self: DeserializeOwned,
46 {
47 let db_path = db.as_ref();
48 if db_path.exists() {
49 EncryptedAtomicDatabase::load(db_path, password)
50 } else {
51 EncryptedAtomicDatabase::create_new(db_path, password)
52 }
53 }
54
55 fn create_from_str<P>(
59 data: &str,
60 path: P,
61 password: &str,
62 ) -> io::Result<EncryptedAtomicDatabase<Self>>
63 where
64 P: AsRef<Path>,
65 Self: DeserializeOwned,
66 {
67 let db_path = path.as_ref();
68 if !db_path.exists() {
69 EncryptedAtomicDatabase::create_from_str(data, path, password)
70 } else {
71 Err(io::Error::new(
72 io::ErrorKind::AlreadyExists,
73 "A file already exists at the provided path!",
74 ))
75 }
76 }
77
78 fn load_encrypted(file: impl Read, key: &Key<Aes256Gcm>) -> io::Result<Self>
80 where
81 Self: DeserializeOwned,
82 {
83 let encrypted: EncryptedData = bincode_cfg().deserialize_from(file).map_err(|e| {
84 io::Error::new(
85 io::ErrorKind::InvalidData,
86 format!("Failed to deserialize encrypted data: {e}"),
87 )
88 })?;
89 Self::decrypt(&encrypted, key)
90 }
91
92 fn save_encrypted(
94 &self,
95 mut file: impl Write,
96 key: &Key<Aes256Gcm>,
97 salt: [u8; SALT_LEN],
98 ) -> io::Result<()> {
99 let encrypted = self.encrypt(key, salt)?;
100 bincode_cfg()
101 .serialize_into(&mut file, &encrypted)
102 .map_err(|e| {
103 io::Error::new(
104 io::ErrorKind::Other,
105 format!("Failed to write encrypted data to file: {e}"),
106 )
107 })
108 }
109
110 fn encrypt(&self, key: &Key<Aes256Gcm>, salt: [u8; SALT_LEN]) -> io::Result<EncryptedData> {
112 let mut nonce = [0u8; NONCE_LEN];
114 OsRng.fill_bytes(&mut nonce);
115
116 let plaintext = bincode_cfg().serialize(self).map_err(|e| {
118 io::Error::new(
119 io::ErrorKind::InvalidData,
120 format!("Serialization failed: {e}"),
121 )
122 })?;
123
124 let cipher = Aes256Gcm::new(key);
125 let ct = cipher
126 .encrypt(Nonce::from_slice(&nonce), plaintext.as_ref())
127 .map_err(|e| io::Error::new(io::ErrorKind::Other, format!("Encryption failed: {e}")))?;
128
129 Ok(EncryptedData {
130 salt,
131 nonce,
132 ciphertext: ct,
133 })
134 }
135
136 fn decrypt(encrypted: &EncryptedData, key: &Key<Aes256Gcm>) -> io::Result<Self>
138 where
139 Self: DeserializeOwned,
140 {
141 let cipher = Aes256Gcm::new(key);
142 let pt = cipher
143 .decrypt(
144 Nonce::from_slice(&encrypted.nonce),
145 encrypted.ciphertext.as_ref(),
146 )
147 .map_err(|e| {
148 io::Error::new(
149 io::ErrorKind::InvalidData,
150 format!("Decryption failed: Incorrect password or corrupted data. {e}"),
151 )
152 })?;
153
154 let data = bincode_cfg().deserialize(&pt).map_err(|e| {
155 io::Error::new(
156 io::ErrorKind::InvalidData,
157 format!("Failed to deserialize decrypted data: {e}"),
158 )
159 })?;
160
161 Ok(data)
162 }
163}
164
165fn derive_key(password: &str, salt: &[u8]) -> io::Result<Key<Aes256Gcm>> {
167 let mut key = [0u8; 32];
168 Argon2::default()
169 .hash_password_into(password.as_bytes(), salt, &mut key)
170 .map_err(|_| io::Error::new(io::ErrorKind::Other, "Key derivation failed"))?;
171
172 let out = *Key::<Aes256Gcm>::from_slice(&key);
173 key.zeroize(); Ok(out)
175}
176
177pub struct EncryptedAtomicDatabase<T: EncryptedDataStore> {
179 path: PathBuf,
180 tmp: PathBuf,
181 data: RwLock<T>,
182 key: RwLock<Key<Aes256Gcm>>,
183 salt: RwLock<[u8; SALT_LEN]>,
184}
185
186impl<T: EncryptedDataStore + DeserializeOwned> EncryptedAtomicDatabase<T> {
187 pub fn load<P: AsRef<Path>>(path: P, password: &str) -> io::Result<Self> {
189 let new_path = path.as_ref().to_path_buf();
190 let tmp = Self::tmp_path(&new_path)?;
191
192 let mut file = File::open(&new_path)?;
194 let encrypted: EncryptedData = bincode_cfg().deserialize_from(&mut file).map_err(|e| {
195 io::Error::new(
196 io::ErrorKind::InvalidData,
197 format!("Failed to deserialize encrypted data: {e}"),
198 )
199 })?;
200 let key = derive_key(password, &encrypted.salt)?;
201 let data = T::decrypt(&encrypted, &key)?;
202
203 Ok(Self {
204 path: new_path,
205 tmp,
206 data: RwLock::new(data),
207 key: RwLock::new(key),
208 salt: RwLock::new(encrypted.salt),
209 })
210 }
211
212 pub fn create_from_str<P: AsRef<Path>>(
215 data: &str,
216 path: P,
217 password: &str,
218 ) -> io::Result<Self> {
219 let new_path = path.as_ref().to_path_buf();
220 let tmp = Self::tmp_path(&new_path)?;
221
222 let encrypted: EncryptedData = bincode_cfg().deserialize(data.as_bytes()).map_err(|e| {
223 io::Error::new(
224 io::ErrorKind::InvalidData,
225 format!("Failed to deserialize encrypted data: {e}"),
226 )
227 })?;
228 let key = derive_key(password, &encrypted.salt)?;
229 let data = T::decrypt(&encrypted, &key)?;
230
231 atomic_write_encrypted(&tmp, &new_path, &data, &key, encrypted.salt)?;
232
233 Ok(Self {
234 path: new_path,
235 tmp,
236 data: RwLock::new(data),
237 key: RwLock::new(key),
238 salt: RwLock::new(encrypted.salt),
239 })
240 }
241
242 pub fn create_new<P: AsRef<Path>>(path: P, password: &str) -> io::Result<Self> {
244 let new_path = path.as_ref().to_path_buf();
245 let tmp = Self::tmp_path(&new_path)?;
246
247 let mut salt_bytes = [0u8; SALT_LEN];
249 OsRng.fill_bytes(&mut salt_bytes);
250 let key = derive_key(password, &salt_bytes)?;
251
252 let data = Default::default();
253 atomic_write_encrypted(&tmp, &new_path, &data, &key, salt_bytes)?;
254
255 Ok(Self {
256 path: new_path,
257 tmp,
258 data: RwLock::new(data),
259 key: RwLock::new(key),
260 salt: RwLock::new(salt_bytes),
261 })
262 }
263
264 pub fn read(&self) -> EncryptedAtomicDatabaseRead<'_, T> {
266 EncryptedAtomicDatabaseRead {
267 data: self.data.read(),
268 }
269 }
270
271 pub fn write(&self) -> EncryptedAtomicDatabaseWrite<'_, T> {
273 let key = *self.key.read();
274 let salt = *self.salt.read();
275 EncryptedAtomicDatabaseWrite {
276 path: self.path.as_ref(),
277 tmp: self.tmp.as_ref(),
278 data: self.data.write(),
279 key,
280 salt,
281 }
282 }
283
284 pub fn change_password(&self, new_password: &str) -> io::Result<()> {
286 let data_guard = self.data.read();
287
288 let mut new_salt = [0u8; SALT_LEN];
289 OsRng.fill_bytes(&mut new_salt);
290 let new_key = derive_key(new_password, &new_salt)?;
291
292 atomic_write_encrypted(&self.tmp, &self.path, &*data_guard, &new_key, new_salt)?;
293
294 {
295 let mut key_lock = self.key.write();
296 *key_lock = new_key;
297 }
298 {
299 let mut salt_lock = self.salt.write();
300 *salt_lock = new_salt;
301 }
302
303 Ok(())
304 }
305
306 fn tmp_path(path: &Path) -> io::Result<PathBuf> {
307 let mut tmp_name = OsString::from(".");
308 tmp_name.push(path.file_name().unwrap_or(OsStr::new("db")));
309 tmp_name.push("~");
310 let tmp = path.with_file_name(tmp_name);
311 if tmp.exists() {
312 error!(
313 "Found orphaned database temporary file '{tmp:?}'. The server has recently crashed or is already running. Delete this before continuing!"
314 );
315 return Err(io::Error::new(
316 io::ErrorKind::AlreadyExists,
317 "Orphaned temporary file exists",
318 ));
319 }
320 Ok(tmp)
321 }
322}
323
324fn atomic_write_encrypted<T: EncryptedDataStore>(
326 tmp: &Path,
327 path: &Path,
328 data: &T,
329 key: &Key<Aes256Gcm>,
330 salt: [u8; SALT_LEN],
331) -> io::Result<()> {
332 {
333 let tmpfile = File::create(tmp)?;
334 data.save_encrypted(tmpfile, key, salt)?;
335 }
336 fs::rename(tmp, path)?;
337 Ok(())
338}
339
340impl<T: EncryptedDataStore> fmt::Debug for EncryptedAtomicDatabase<T> {
341 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
342 f.debug_struct("EncryptedAtomicDatabase")
343 .field("file", &self.path)
344 .finish()
345 }
346}
347
348impl<T: EncryptedDataStore> Drop for EncryptedAtomicDatabase<T> {
349 fn drop(&mut self) {
350 info!("Saving database");
351 let data_guard = self.data.read();
352 let key = self.key.read();
353 let salt = self.salt.read();
354 if let Err(e) = atomic_write_encrypted(&self.tmp, &self.path, &*data_guard, &key, *salt) {
355 error!("Failed to save database: {}", e);
356 }
357 }
358}
359
360pub struct EncryptedAtomicDatabaseRead<'a, T: EncryptedDataStore> {
361 data: RwLockReadGuard<'a, T>,
362}
363
364impl<'a, T: EncryptedDataStore> Deref for EncryptedAtomicDatabaseRead<'a, T> {
365 type Target = T;
366 fn deref(&self) -> &Self::Target {
367 &self.data
368 }
369}
370
371pub struct EncryptedAtomicDatabaseWrite<'a, T: EncryptedDataStore> {
372 tmp: &'a Path,
373 path: &'a Path,
374 data: RwLockWriteGuard<'a, T>,
375 key: Key<Aes256Gcm>,
376 salt: [u8; SALT_LEN],
377}
378
379impl<'a, T: EncryptedDataStore> Deref for EncryptedAtomicDatabaseWrite<'a, T> {
380 type Target = T;
381 fn deref(&self) -> &Self::Target {
382 &self.data
383 }
384}
385
386impl<'a, T: EncryptedDataStore> DerefMut for EncryptedAtomicDatabaseWrite<'a, T> {
387 fn deref_mut(&mut self) -> &mut Self::Target {
388 &mut self.data
389 }
390}
391
392impl<'a, T: EncryptedDataStore> Drop for EncryptedAtomicDatabaseWrite<'a, T> {
393 fn drop(&mut self) {
394 info!("Saving database");
395 if let Err(e) =
396 atomic_write_encrypted(self.tmp, self.path, &*self.data, &self.key, self.salt)
397 {
398 error!("Failed to save database: {}", e);
399 }
400 }
401}