Skip to main content

forest/key_management/
keystore.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use 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
36/// Environmental variable which holds the `KeyStore` encryption phrase.
37pub const FOREST_KEYSTORE_PHRASE_ENV: &str = "FOREST_KEYSTORE_PHRASE";
38
39type SaltByteArray = [u8; RECOMMENDED_SALT_LEN];
40
41/// `KeyInfo` structure, this contains the type of key (stored as a string) and
42/// the private key. Note how the private key is stored as a byte vector
43#[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    // Vec<u8> is used because The private keys for BLS and SECP256K1 are not of the same type
48    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    /// Return a reference to the key's signature type
59    pub fn key_type(&self) -> &SignatureType {
60        &self.key_type
61    }
62
63    /// Return a reference to the private key
64    pub fn private_key(&self) -> &Vec<u8> {
65        &self.private_key
66    }
67}
68
69/// `KeyStore` structure, this contains a set of `KeyInfos` indexed by address.
70#[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/// Persistent `KeyStore` in JSON clear text in `KEYSTORE_LOCATION`
84#[derive(Clone, PartialEq, Debug, Eq)]
85struct PersistentKeyStore {
86    file_path: PathBuf,
87}
88
89/// Encrypted `KeyStore`
90/// `Argon2id` hash key derivation
91/// `XSalsa20Poly1305` authenticated encryption
92/// CBOR encoding
93#[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    /// An error occurred while encrypting keys
102    #[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                        // Existing cleartext JSON keystore
122                        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                            // New encrypted keystore if file exists but is zero bytes (i.e., touch)
184                            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                            // Existing encrypted keystore
206                            // Split off data from prepended salt
207                            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                        // Flush For EncryptedKeyStore
269                        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                        // Flush for PersistentKeyStore
294                        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                // NoOp for MemKeyStore
304                Ok(())
305            }
306        }
307    }
308
309    /// Return all of the keys that are stored in the `KeyStore`
310    pub fn list(&self) -> Vec<String> {
311        self.key_info.keys().cloned().collect()
312    }
313
314    /// Return `KeyInfo` that corresponds to a given key
315    pub fn get(&self, k: &str) -> Result<KeyInfo, Error> {
316        self.key_info.get(k).cloned().ok_or(Error::KeyInfo)
317    }
318
319    /// Save a key/`KeyInfo` pair to the `KeyStore`
320    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    /// Insert or update the default `KeyInfo` stored in the `KeyStore`
334    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    /// Remove the key and corresponding `KeyInfo` from the `KeyStore`
345    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        // #define crypto_pwhash_argon2id_MEMLIMIT_INTERACTIVE 67108864U
372        // see <https://github.com/jedisct1/libsodium/blob/089f850608737f9d969157092988cb274fe7f8d4/src/libsodium/include/sodium/crypto_pwhash_argon2id.h#L70>
373        const CRYPTO_PWHASH_ARGON2ID_MEMLIMIT_INTERACTIVE: u32 = 67108864;
374        // #define crypto_pwhash_argon2id_OPSLIMIT_INTERACTIVE 2U
375        // see <https://github.com/jedisct1/libsodium/blob/089f850608737f9d969157092988cb274fe7f8d4/src/libsodium/include/sodium/crypto_pwhash_argon2id.h#L66>
376        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        // https://docs.rs/sodiumoxide/latest/sodiumoxide/crypto/secretbox/xsalsa20poly1305/constant.KEYBYTES.html
381        // KEYBYTES = 0x20
382        // param_builder.output_len(32)?;
383        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        // Manually parse keystore.json
513        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        // Read existing keystore.json
529        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}