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 {
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 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 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 fn encrypt(&self, key: &Key<Aes256Gcm>, salt: [u8; SALT_LEN]) -> io::Result<EncryptedData> {
105 let mut nonce = [0u8; NONCE_LEN];
107 OsRng.fill_bytes(&mut nonce);
108
109 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 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
158fn 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(); Ok(out)
168}
169
170pub 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 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 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 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 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 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 pub fn read(&self) -> EncryptedAtomicDatabaseRead<'_, T> {
258 EncryptedAtomicDatabaseRead {
259 data: self.data.read(),
260 }
261 }
262
263 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 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
316fn 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}