1use std::{
5 fmt::Display,
6 fs::{File, create_dir},
7 io::{BufReader, BufWriter, ErrorKind, Read, Write},
8 path::{Path, PathBuf},
9};
10
11use crate::{
12 shim::crypto::SignatureType,
13 utils::{encoding::from_slice_with_fallback, io::create_new_sensitive_file},
14};
15use ahash::{HashMap, HashMapExt};
16use argon2::{
17 Argon2, ParamsBuilder, PasswordHasher, RECOMMENDED_SALT_LEN, password_hash::SaltString,
18};
19use base64::{Engine, prelude::BASE64_STANDARD};
20use crypto_secretbox::{
21 KeyInit, SecretBox, XSalsa20Poly1305,
22 aead::{Aead, generic_array::GenericArray},
23};
24use rand::RngCore;
25use serde::{Deserialize, Serialize};
26use thiserror::Error;
27use tracing::{error, warn};
28
29use super::errors::Error;
30
31const NONCE_SIZE: usize = SecretBox::<Box<dyn std::any::Any>>::NONCE_SIZE;
32
33pub const KEYSTORE_NAME: &str = "keystore.json";
34pub const ENCRYPTED_KEYSTORE_NAME: &str = "keystore";
35
36pub const FOREST_KEYSTORE_PHRASE_ENV: &str = "FOREST_KEYSTORE_PHRASE";
38
39type SaltByteArray = [u8; RECOMMENDED_SALT_LEN];
40
41#[cfg_attr(test, derive(derive_quickcheck_arbitrary::Arbitrary))]
44#[derive(Clone, PartialEq, Debug, Eq, Serialize, Deserialize, derive_more::Constructor)]
45pub struct KeyInfo {
46 key_type: SignatureType,
47 private_key: Vec<u8>,
49}
50
51#[derive(Clone, PartialEq, Debug, Eq, Serialize, Deserialize)]
52pub struct PersistentKeyInfo {
53 key_type: SignatureType,
54 private_key: String,
55}
56
57impl KeyInfo {
58 pub fn key_type(&self) -> &SignatureType {
60 &self.key_type
61 }
62
63 pub fn private_key(&self) -> &Vec<u8> {
65 &self.private_key
66 }
67}
68
69#[derive(Clone, PartialEq, Debug, Eq)]
71pub struct KeyStore {
72 key_info: HashMap<String, KeyInfo>,
73 persistence: Option<PersistentKeyStore>,
74 encryption: Option<EncryptedKeyStore>,
75}
76
77pub enum KeyStoreConfig {
78 Memory,
79 Persistent(PathBuf),
80 Encrypted(PathBuf, String),
81}
82
83#[derive(Clone, PartialEq, Debug, Eq)]
85struct PersistentKeyStore {
86 file_path: PathBuf,
87}
88
89#[derive(Clone, PartialEq, Debug, Eq)]
94struct EncryptedKeyStore {
95 salt: SaltByteArray,
96 encryption_key: Vec<u8>,
97}
98
99#[derive(Debug, Error)]
100pub enum EncryptedKeyStoreError {
101 #[error("Error encrypting data")]
103 EncryptionError,
104}
105
106impl KeyStore {
107 pub fn new(config: KeyStoreConfig) -> Result<Self, Error> {
108 match config {
109 KeyStoreConfig::Memory => Ok(Self {
110 key_info: HashMap::new(),
111 persistence: None,
112 encryption: None,
113 }),
114 KeyStoreConfig::Persistent(location) => {
115 let file_path = location.join(KEYSTORE_NAME);
116
117 match File::open(&file_path) {
118 Ok(file) => {
119 let reader = BufReader::new(file);
120
121 let persisted_key_info: HashMap<String, PersistentKeyInfo> =
123 serde_json::from_reader(reader)
124 .inspect_err(|error| {
125 error!(%error, "failed to deserialize keyfile, initializing new keystore at: {}.", file_path.display());
126 })
127 .unwrap_or_default();
128
129 let mut key_info = HashMap::new();
130 for (key, value) in persisted_key_info.iter() {
131 key_info.insert(
132 key.to_string(),
133 KeyInfo {
134 private_key: BASE64_STANDARD
135 .decode(value.private_key.clone())
136 .map_err(|error| Error::Other(error.to_string()))?,
137 key_type: value.key_type,
138 },
139 );
140 }
141
142 Ok(Self {
143 key_info,
144 persistence: Some(PersistentKeyStore { file_path }),
145 encryption: None,
146 })
147 }
148 Err(e) => {
149 if e.kind() == ErrorKind::NotFound {
150 warn!(
151 "Keystore does not exist, initializing new keystore at: {}",
152 file_path.display()
153 );
154 Ok(Self {
155 key_info: HashMap::new(),
156 persistence: Some(PersistentKeyStore { file_path }),
157 encryption: None,
158 })
159 } else {
160 Err(Error::Other(e.to_string()))
161 }
162 }
163 }
164 }
165 KeyStoreConfig::Encrypted(location, passphrase) => {
166 if !location.exists() {
167 create_dir(location.clone())?;
168 }
169
170 let file_path = location.join(Path::new(ENCRYPTED_KEYSTORE_NAME));
171
172 if !file_path.exists() {
173 File::create(file_path.clone())?;
174 }
175
176 match File::open(&file_path) {
177 Ok(file) => {
178 let mut reader = BufReader::new(file);
179 let mut buf = vec![];
180 let read_bytes = reader.read_to_end(&mut buf)?;
181
182 if read_bytes == 0 {
183 warn!(
185 "Keystore does not exist, initializing new keystore at {:?}",
186 file_path
187 );
188
189 let (salt, encryption_key) =
190 EncryptedKeyStore::derive_key(&passphrase, None).map_err(
191 |error| {
192 error!("Failed to create key from passphrase");
193 Error::Other(error.to_string())
194 },
195 )?;
196 Ok(Self {
197 key_info: HashMap::new(),
198 persistence: Some(PersistentKeyStore { file_path }),
199 encryption: Some(EncryptedKeyStore {
200 salt,
201 encryption_key,
202 }),
203 })
204 } else {
205 let data = buf.split_off(RECOMMENDED_SALT_LEN);
208 let mut prev_salt = [0; RECOMMENDED_SALT_LEN];
209 prev_salt.copy_from_slice(&buf);
210 let (salt, encryption_key) =
211 EncryptedKeyStore::derive_key(&passphrase, Some(prev_salt))
212 .map_err(|error| {
213 error!("Failed to create key from passphrase");
214 Error::Other(error.to_string())
215 })?;
216
217 let decrypted_data = EncryptedKeyStore::decrypt(&encryption_key, &data)
218 .map_err(|error| Error::Other(error.to_string()))?;
219
220 let key_info = from_slice_with_fallback(&decrypted_data)
221 .inspect_err(|error| {
222 error!(%error, "Failed to deserialize keyfile, initializing new");
223 })
224 .unwrap_or_default();
225
226 Ok(Self {
227 key_info,
228 persistence: Some(PersistentKeyStore { file_path }),
229 encryption: Some(EncryptedKeyStore {
230 salt,
231 encryption_key,
232 }),
233 })
234 }
235 }
236 Err(_) => {
237 warn!("Encrypted keystore does not exist, initializing new keystore");
238
239 let (salt, encryption_key) =
240 EncryptedKeyStore::derive_key(&passphrase, None).map_err(|error| {
241 error!("Failed to create key from passphrase");
242 Error::Other(error.to_string())
243 })?;
244
245 Ok(Self {
246 key_info: HashMap::new(),
247 persistence: Some(PersistentKeyStore { file_path }),
248 encryption: Some(EncryptedKeyStore {
249 salt,
250 encryption_key,
251 }),
252 })
253 }
254 }
255 }
256 }
257 }
258
259 pub fn flush(&self) -> anyhow::Result<()> {
260 match &self.persistence {
261 Some(persistent_keystore) => {
262 let file = create_new_sensitive_file(&persistent_keystore.file_path)?;
263
264 let mut writer = BufWriter::new(file);
265
266 match &self.encryption {
267 Some(encrypted_keystore) => {
268 let data = serde_ipld_dagcbor::to_vec(&self.key_info).map_err(|e| {
270 Error::Other(format!("failed to serialize and write key info: {e}"))
271 })?;
272
273 let encrypted_data =
274 EncryptedKeyStore::encrypt(&encrypted_keystore.encryption_key, &data)?;
275 let mut salt_vec = encrypted_keystore.salt.to_vec();
276 salt_vec.extend(encrypted_data);
277 writer.write_all(&salt_vec)?;
278
279 Ok(())
280 }
281 None => {
282 let mut key_info: HashMap<String, PersistentKeyInfo> = HashMap::new();
283 for (key, value) in self.key_info.iter() {
284 key_info.insert(
285 key.to_string(),
286 PersistentKeyInfo {
287 private_key: BASE64_STANDARD.encode(value.private_key.clone()),
288 key_type: value.key_type,
289 },
290 );
291 }
292
293 serde_json::to_writer_pretty(writer, &key_info).map_err(|e| {
295 Error::Other(format!("failed to serialize and write key info: {e}"))
296 })?;
297
298 Ok(())
299 }
300 }
301 }
302 None => {
303 Ok(())
305 }
306 }
307 }
308
309 pub fn list(&self) -> Vec<String> {
311 self.key_info.keys().cloned().collect()
312 }
313
314 pub fn get(&self, k: &str) -> Result<KeyInfo, Error> {
316 self.key_info.get(k).cloned().ok_or(Error::KeyInfo)
317 }
318
319 pub fn put(&mut self, key: &str, key_info: KeyInfo) -> Result<(), Error> {
321 if self.key_info.contains_key(key) {
322 return Err(Error::KeyExists);
323 }
324 self.key_info.insert(key.to_string(), key_info);
325
326 if self.persistence.is_some() {
327 self.flush().map_err(|err| Error::Other(err.to_string()))?;
328 }
329
330 Ok(())
331 }
332
333 pub fn set_default(&mut self, key_info: KeyInfo) -> Result<(), Error> {
335 self.key_info.insert("default".to_owned(), key_info);
336
337 if self.persistence.is_some() {
338 self.flush().map_err(|err| Error::Other(err.to_string()))?;
339 }
340
341 Ok(())
342 }
343
344 pub fn remove(&mut self, key: &str) -> anyhow::Result<KeyInfo> {
346 let key_out = self.key_info.remove(key).ok_or(Error::KeyInfo)?;
347
348 if self.persistence.is_some() {
349 self.flush()?;
350 }
351
352 Ok(key_out)
353 }
354}
355
356impl EncryptedKeyStore {
357 fn derive_key(
358 passphrase: &str,
359 prev_salt: Option<SaltByteArray>,
360 ) -> anyhow::Result<(SaltByteArray, Vec<u8>)> {
361 let salt = match prev_salt {
362 Some(prev_salt) => prev_salt,
363 None => {
364 let mut salt = [0; RECOMMENDED_SALT_LEN];
365 crate::utils::rand::forest_os_rng().fill_bytes(&mut salt);
366 salt
367 }
368 };
369
370 let mut param_builder = ParamsBuilder::new();
371 const CRYPTO_PWHASH_ARGON2ID_MEMLIMIT_INTERACTIVE: u32 = 67108864;
374 const CRYPTO_PWHASH_ARGON2ID_OPSLIMIT_INTERACTIVE: u32 = 2;
377 param_builder
378 .m_cost(CRYPTO_PWHASH_ARGON2ID_MEMLIMIT_INTERACTIVE / 1024)
379 .t_cost(CRYPTO_PWHASH_ARGON2ID_OPSLIMIT_INTERACTIVE);
380 let hasher = Argon2::new(
384 argon2::Algorithm::Argon2id,
385 argon2::Version::V0x13,
386 param_builder.build().map_err(map_err_to_anyhow)?,
387 );
388 let salt_string = SaltString::encode_b64(&salt).map_err(map_err_to_anyhow)?;
389 let pw_hash = hasher
390 .hash_password(passphrase.as_bytes(), &salt_string)
391 .map_err(map_err_to_anyhow)?;
392 if let Some(hash) = pw_hash.hash {
393 Ok((salt, hash.as_bytes().to_vec()))
394 } else {
395 anyhow::bail!(EncryptedKeyStoreError::EncryptionError)
396 }
397 }
398
399 fn encrypt(encryption_key: &[u8], msg: &[u8]) -> anyhow::Result<Vec<u8>> {
400 let mut nonce = [0; NONCE_SIZE];
401 crate::utils::rand::forest_os_rng().fill_bytes(&mut nonce);
402 let nonce = GenericArray::from_slice(&nonce);
403 let key = GenericArray::from_slice(encryption_key);
404 let cipher = XSalsa20Poly1305::new(key);
405 let mut ciphertext = cipher.encrypt(nonce, msg).map_err(map_err_to_anyhow)?;
406 ciphertext.extend(nonce.iter());
407 Ok(ciphertext)
408 }
409
410 #[allow(clippy::indexing_slicing)]
411 fn decrypt(encryption_key: &[u8], msg: &[u8]) -> anyhow::Result<Vec<u8>> {
412 anyhow::ensure!(msg.len() > NONCE_SIZE);
413 let cyphertext_len = msg.len() - NONCE_SIZE;
414 let ciphertext = &msg[..cyphertext_len];
415 let nonce = GenericArray::from_slice(&msg[cyphertext_len..]);
416 let key = GenericArray::from_slice(encryption_key);
417 let cipher = XSalsa20Poly1305::new(key);
418 let plaintext = cipher
419 .decrypt(nonce, ciphertext)
420 .map_err(map_err_to_anyhow)?;
421 Ok(plaintext)
422 }
423}
424
425fn map_err_to_anyhow<T: Display>(e: T) -> anyhow::Error {
426 anyhow::Error::msg(e.to_string())
427}
428
429#[cfg(test)]
430mod test {
431 use base64::{Engine, prelude::BASE64_STANDARD};
432
433 use super::*;
434 use crate::key_management::wallet;
435
436 const PASSPHRASE: &str = "foobarbaz";
437
438 #[test]
439 fn test_generate_key() {
440 let (salt, encryption_key) = EncryptedKeyStore::derive_key(PASSPHRASE, None).unwrap();
441 let (second_salt, second_key) =
442 EncryptedKeyStore::derive_key(PASSPHRASE, Some(salt)).unwrap();
443
444 assert_eq!(
445 encryption_key, second_key,
446 "Derived key must be deterministic"
447 );
448 assert_eq!(salt, second_salt, "Salts must match");
449 }
450
451 #[test]
452 fn test_encrypt_message() {
453 let (_, private_key) = EncryptedKeyStore::derive_key(PASSPHRASE, None).unwrap();
454 let message = "foo is coming";
455 let ciphertext = EncryptedKeyStore::encrypt(&private_key, message.as_bytes()).unwrap();
456 let second_pass = EncryptedKeyStore::encrypt(&private_key, message.as_bytes()).unwrap();
457 assert_ne!(
458 ciphertext, second_pass,
459 "Ciphertexts use secure initialization vectors"
460 );
461 }
462
463 #[test]
464 fn test_decrypt_message() {
465 let (_, private_key) = EncryptedKeyStore::derive_key(PASSPHRASE, None).unwrap();
466 let message = "foo is coming";
467 let ciphertext = EncryptedKeyStore::encrypt(&private_key, message.as_bytes()).unwrap();
468 let plaintext = EncryptedKeyStore::decrypt(&private_key, &ciphertext).unwrap();
469 assert_eq!(plaintext, message.as_bytes());
470 }
471
472 #[test]
473 fn test_read_old_encrypted_keystore() {
474 let dir: PathBuf = "src/key_management/tests/keystore_encrypted_old".into();
475 assert!(dir.exists());
476 let ks = KeyStore::new(KeyStoreConfig::Encrypted(dir, PASSPHRASE.to_string())).unwrap();
477 assert!(ks.persistence.is_some());
478 }
479
480 #[test]
481 fn test_read_write_encrypted_keystore() {
482 let keystore_location = tempfile::tempdir().unwrap().keep();
483 let ks = KeyStore::new(KeyStoreConfig::Encrypted(
484 keystore_location.clone(),
485 PASSPHRASE.to_string(),
486 ))
487 .unwrap();
488 ks.flush().unwrap();
489
490 let ks_read = KeyStore::new(KeyStoreConfig::Encrypted(
491 keystore_location,
492 PASSPHRASE.to_string(),
493 ))
494 .unwrap();
495
496 assert_eq!(ks, ks_read);
497 }
498
499 #[test]
500 fn test_read_write_keystore() {
501 let keystore_location = tempfile::tempdir().unwrap().keep();
502 let mut ks = KeyStore::new(KeyStoreConfig::Persistent(keystore_location.clone())).unwrap();
503
504 let key = wallet::generate_key(SignatureType::Bls).unwrap();
505
506 let addr = format!("wallet-{}", key.address);
507 ks.put(&addr, key.key_info).unwrap();
508 ks.flush().unwrap();
509
510 let default = ks.get(&addr).unwrap();
511
512 let keystore_file = keystore_location.join(KEYSTORE_NAME);
514 let reader = BufReader::new(File::open(keystore_file).unwrap());
515 let persisted_keystore: HashMap<String, PersistentKeyInfo> =
516 serde_json::from_reader(reader).unwrap();
517
518 let default_key_info = persisted_keystore.get(&addr).unwrap();
519 let actual = BASE64_STANDARD
520 .decode(default_key_info.private_key.clone())
521 .unwrap();
522
523 assert_eq!(
524 default.private_key, actual,
525 "persisted key matches key from key store"
526 );
527
528 let ks_read = KeyStore::new(KeyStoreConfig::Persistent(keystore_location)).unwrap();
530 assert_eq!(ks, ks_read);
531 }
532
533 #[test]
534 fn test_set_default_replaces_and_persists() {
535 let keystore_location = tempfile::tempdir().unwrap().keep();
536 let mut ks = KeyStore::new(KeyStoreConfig::Persistent(keystore_location.clone())).unwrap();
537
538 let key_1 = wallet::generate_key(SignatureType::Secp256k1).unwrap();
539 let key_2 = wallet::generate_key(SignatureType::Bls).unwrap();
540
541 ks.set_default(key_1.key_info.clone()).unwrap();
542 assert_eq!(ks.get("default").unwrap(), key_1.key_info);
543
544 ks.set_default(key_2.key_info.clone()).unwrap();
545 assert_eq!(ks.get("default").unwrap(), key_2.key_info);
546
547 let ks_read = KeyStore::new(KeyStoreConfig::Persistent(keystore_location)).unwrap();
548 assert_eq!(ks_read.get("default").unwrap(), key_2.key_info);
549 }
550}