1use aes_gcm::{
2 aead::{rand_core::RngCore, Aead, KeyInit, OsRng},
3 Aes256Gcm, Key, Nonce,
4};
5use argon2::{self, Argon2};
6use bincode::{
7 self,
8 serde::{decode_from_slice, decode_from_std_read, encode_into_std_write, encode_to_vec},
9};
10use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
11use serde::{de::DeserializeOwned, Deserialize, Serialize};
12use std::{
13 ffi::{OsStr, OsString},
14 fmt,
15 fs::{self, File},
16 io::{self, Read, Write},
17 ops::{Deref, DerefMut},
18 path::{Path, PathBuf},
19};
20use tracing::{error, info};
21use zeroize::Zeroize;
22
23const SALT_LEN: usize = 16;
24const NONCE_LEN: usize = 12;
25
26#[inline]
27pub fn bincode_cfg() -> impl bincode::config::Config {
28 bincode::config::standard().with_fixed_int_encoding()
29}
30
31#[derive(Serialize, Deserialize)]
33pub struct EncryptedData {
34 salt: [u8; SALT_LEN],
35 nonce: [u8; NONCE_LEN],
36 ciphertext: Vec<u8>,
37}
38
39pub trait EncryptedDataStore: Default + Serialize {
43 fn open<P>(db: P, password: &str) -> io::Result<EncryptedAtomicDatabase<Self>>
46 where
47 P: AsRef<Path>,
48 Self: DeserializeOwned,
49 {
50 let db_path = db.as_ref();
51 if db_path.exists() {
52 EncryptedAtomicDatabase::load(db_path, password)
53 } else {
54 EncryptedAtomicDatabase::create_new(db_path, password)
55 }
56 }
57
58 fn create_from_str<P>(
62 data: &str,
63 path: P,
64 password: &str,
65 ) -> io::Result<EncryptedAtomicDatabase<Self>>
66 where
67 P: AsRef<Path>,
68 Self: DeserializeOwned,
69 {
70 let db_path = path.as_ref();
71 if !db_path.exists() {
72 EncryptedAtomicDatabase::create_from_str(data, path, password)
73 } else {
74 Err(io::Error::new(
75 io::ErrorKind::AlreadyExists,
76 "A file already exists at the provided path!",
77 ))
78 }
79 }
80
81 fn load_encrypted(file: &mut impl Read, key: &Key<Aes256Gcm>) -> io::Result<Self>
83 where
84 Self: DeserializeOwned,
85 {
86 let encrypted: EncryptedData = decode_from_std_read(file, bincode_cfg()).map_err(|e| {
87 io::Error::new(
88 io::ErrorKind::InvalidData,
89 format!("Failed to decode encrypted data: {e}"),
90 )
91 })?;
92 Self::decrypt(&encrypted, key)
93 }
94
95 fn save_encrypted(
97 &self,
98 mut file: impl Write,
99 key: &Key<Aes256Gcm>,
100 salt: [u8; SALT_LEN],
101 ) -> io::Result<usize> {
102 let encrypted = self.encrypt(key, salt)?;
103 encode_into_std_write(encrypted, &mut file, bincode_cfg()).map_err(|e| {
104 io::Error::new(
105 io::ErrorKind::Other,
106 format!("Failed to write encrypted data to file: {e}"),
107 )
108 })
109 }
110
111 fn encrypt(&self, key: &Key<Aes256Gcm>, salt: [u8; SALT_LEN]) -> io::Result<EncryptedData> {
113 let mut nonce = [0u8; NONCE_LEN];
115 OsRng.fill_bytes(&mut nonce);
116
117 let plaintext = encode_to_vec(self, bincode_cfg()).map_err(|e| {
119 io::Error::new(io::ErrorKind::InvalidData, format!("Encoding failed: {e}"))
120 })?;
121
122 let cipher = Aes256Gcm::new(key);
123 let ct = cipher
124 .encrypt(Nonce::from_slice(&nonce), plaintext.as_ref())
125 .map_err(|e| io::Error::new(io::ErrorKind::Other, format!("Encryption failed: {e}")))?;
126
127 Ok(EncryptedData {
128 salt,
129 nonce,
130 ciphertext: ct,
131 })
132 }
133
134 fn decrypt(encrypted: &EncryptedData, key: &Key<Aes256Gcm>) -> io::Result<Self>
136 where
137 Self: DeserializeOwned,
138 {
139 let cipher = Aes256Gcm::new(key);
140 let pt = cipher
141 .decrypt(
142 Nonce::from_slice(&encrypted.nonce),
143 encrypted.ciphertext.as_ref(),
144 )
145 .map_err(|e| {
146 io::Error::new(
147 io::ErrorKind::InvalidData,
148 format!("Decryption failed: Incorrect password or corrupted data. {e}"),
149 )
150 })?;
151
152 let (data, _) = decode_from_slice(&pt, bincode_cfg()).map_err(|e| {
153 io::Error::new(
154 io::ErrorKind::InvalidData,
155 format!("Failed to decode decrypted data: {e}"),
156 )
157 })?;
158
159 Ok(data)
160 }
161}
162
163fn derive_key(password: &str, salt: &[u8]) -> io::Result<Key<Aes256Gcm>> {
165 let mut key = [0u8; 32];
166 Argon2::default()
167 .hash_password_into(password.as_bytes(), salt, &mut key)
168 .map_err(|_| io::Error::new(io::ErrorKind::Other, "Key derivation failed"))?;
169
170 let out = *Key::<Aes256Gcm>::from_slice(&key);
171 key.zeroize(); Ok(out)
173}
174
175pub struct EncryptedAtomicDatabase<T: EncryptedDataStore> {
177 path: PathBuf,
178 tmp: PathBuf,
179 data: RwLock<T>,
180 key: RwLock<Key<Aes256Gcm>>,
181 salt: RwLock<[u8; SALT_LEN]>,
182}
183
184impl<T: EncryptedDataStore + DeserializeOwned> EncryptedAtomicDatabase<T> {
185 pub fn load<P: AsRef<Path>>(path: P, password: &str) -> io::Result<Self> {
187 let new_path = path.as_ref().to_path_buf();
188 let tmp = Self::tmp_path(&new_path)?;
189
190 let mut file = File::open(&new_path)?;
192 let encrypted: EncryptedData =
193 decode_from_std_read(&mut file, bincode_cfg()).map_err(|e| {
194 io::Error::new(
195 io::ErrorKind::InvalidData,
196 format!("Failed to decode encrypted data: {e}"),
197 )
198 })?;
199 let key = derive_key(password, &encrypted.salt)?;
200 let data = T::decrypt(&encrypted, &key)?;
201
202 Ok(Self {
203 path: new_path,
204 tmp,
205 data: RwLock::new(data),
206 key: RwLock::new(key),
207 salt: RwLock::new(encrypted.salt),
208 })
209 }
210
211 pub fn create_from_str<P: AsRef<Path>>(
214 data: &str,
215 path: P,
216 password: &str,
217 ) -> io::Result<Self> {
218 let new_path = path.as_ref().to_path_buf();
219 let tmp = Self::tmp_path(&new_path)?;
220
221 let (encrypted, _): (EncryptedData, usize) =
222 decode_from_slice(data.as_bytes(), bincode_cfg()).map_err(|e| {
223 io::Error::new(
224 io::ErrorKind::InvalidData,
225 format!("Failed to decode 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}