Skip to main content

nklave_core/keys/
keystore.rs

1//! EIP-2335 Keystore support
2//!
3//! Implements loading of validator keystores according to EIP-2335
4//! https://eips.ethereum.org/EIPS/eip-2335
5
6use crate::keys::bls::{BlsKeypair, BlsSecretKey};
7use aes::cipher::{KeyIvInit, StreamCipher};
8use pbkdf2::pbkdf2_hmac;
9use scrypt::{scrypt, Params as ScryptParams};
10use serde::{Deserialize, Serialize};
11use sha2::{Digest, Sha256};
12use std::fs;
13use std::path::Path;
14use thiserror::Error;
15use uuid::Uuid;
16
17type Aes128Ctr = ctr::Ctr128BE<aes::Aes128>;
18
19/// EIP-2335 Keystore JSON structure
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct Keystore {
22    pub crypto: KeystoreCrypto,
23    #[serde(default)]
24    pub description: Option<String>,
25    pub pubkey: String,
26    pub path: String,
27    pub uuid: Uuid,
28    pub version: u32,
29}
30
31/// Crypto section of the keystore
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct KeystoreCrypto {
34    pub kdf: KeystoreKdf,
35    pub checksum: KeystoreChecksum,
36    pub cipher: KeystoreCipher,
37}
38
39/// KDF parameters
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct KeystoreKdf {
42    pub function: String,
43    pub params: KdfParams,
44    pub message: String,
45}
46
47/// KDF parameters (either scrypt or pbkdf2)
48#[derive(Debug, Clone, Serialize, Deserialize)]
49#[serde(untagged)]
50pub enum KdfParams {
51    Scrypt(ScryptKdfParams),
52    Pbkdf2(Pbkdf2KdfParams),
53}
54
55/// Scrypt KDF parameters
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct ScryptKdfParams {
58    pub dklen: u32,
59    pub n: u32,
60    pub p: u32,
61    pub r: u32,
62    pub salt: String,
63}
64
65/// PBKDF2 KDF parameters
66#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct Pbkdf2KdfParams {
68    pub dklen: u32,
69    pub c: u32,
70    pub prf: String,
71    pub salt: String,
72}
73
74/// Checksum section
75#[derive(Debug, Clone, Serialize, Deserialize)]
76pub struct KeystoreChecksum {
77    pub function: String,
78    pub params: serde_json::Value,
79    pub message: String,
80}
81
82/// Cipher section
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct KeystoreCipher {
85    pub function: String,
86    pub params: CipherParams,
87    pub message: String,
88}
89
90/// Cipher parameters
91#[derive(Debug, Clone, Serialize, Deserialize)]
92pub struct CipherParams {
93    pub iv: String,
94}
95
96impl Keystore {
97    /// Load a keystore from a JSON file
98    pub fn load(path: impl AsRef<Path>) -> Result<Self, KeystoreError> {
99        let contents = fs::read_to_string(path.as_ref())
100            .map_err(|e| KeystoreError::Io(e.to_string()))?;
101
102        serde_json::from_str(&contents)
103            .map_err(|e| KeystoreError::Parse(e.to_string()))
104    }
105
106    /// Decrypt the keystore and return the secret key
107    pub fn decrypt(&self, password: &str) -> Result<BlsSecretKey, KeystoreError> {
108        // Derive the decryption key
109        let decryption_key = self.derive_key(password)?;
110
111        // Verify checksum
112        self.verify_checksum(&decryption_key)?;
113
114        // Decrypt the secret key
115        let secret_key_bytes = self.decrypt_secret(&decryption_key)?;
116
117        // Create BLS secret key
118        BlsSecretKey::from_bytes(&secret_key_bytes)
119            .map_err(|e| KeystoreError::InvalidKey(e.to_string()))
120    }
121
122    /// Decrypt and return a full keypair
123    pub fn decrypt_keypair(&self, password: &str) -> Result<BlsKeypair, KeystoreError> {
124        let secret = self.decrypt(password)?;
125        Ok(BlsKeypair::from_secret(secret))
126    }
127
128    /// Derive the decryption key using the specified KDF
129    fn derive_key(&self, password: &str) -> Result<[u8; 32], KeystoreError> {
130        let kdf = &self.crypto.kdf;
131
132        match kdf.function.as_str() {
133            "scrypt" => {
134                let params = match &kdf.params {
135                    KdfParams::Scrypt(p) => p,
136                    _ => return Err(KeystoreError::InvalidKdf("Expected scrypt params".into())),
137                };
138
139                let salt = hex::decode(&params.salt)
140                    .map_err(|e| KeystoreError::Parse(format!("Invalid salt hex: {}", e)))?;
141
142                // Convert n to log2(n) for scrypt params
143                let log_n = (params.n as f64).log2() as u8;
144                let scrypt_params = ScryptParams::new(log_n, params.r, params.p, params.dklen as usize)
145                    .map_err(|e| KeystoreError::InvalidKdf(format!("Invalid scrypt params: {:?}", e)))?;
146
147                let mut dk = [0u8; 32];
148                scrypt(password.as_bytes(), &salt, &scrypt_params, &mut dk)
149                    .map_err(|e| KeystoreError::InvalidKdf(format!("Scrypt failed: {:?}", e)))?;
150
151                Ok(dk)
152            }
153            "pbkdf2" => {
154                let params = match &kdf.params {
155                    KdfParams::Pbkdf2(p) => p,
156                    _ => return Err(KeystoreError::InvalidKdf("Expected pbkdf2 params".into())),
157                };
158
159                let salt = hex::decode(&params.salt)
160                    .map_err(|e| KeystoreError::Parse(format!("Invalid salt hex: {}", e)))?;
161
162                let mut dk = [0u8; 32];
163                pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, params.c, &mut dk);
164
165                Ok(dk)
166            }
167            other => Err(KeystoreError::UnsupportedKdf(other.to_string())),
168        }
169    }
170
171    /// Verify the checksum
172    fn verify_checksum(&self, decryption_key: &[u8; 32]) -> Result<(), KeystoreError> {
173        let cipher_message = hex::decode(&self.crypto.cipher.message)
174            .map_err(|e| KeystoreError::Parse(format!("Invalid cipher message hex: {}", e)))?;
175
176        // Checksum = SHA256(decryption_key[16:32] || cipher_message)
177        let mut hasher = Sha256::new();
178        hasher.update(&decryption_key[16..32]);
179        hasher.update(&cipher_message);
180        let computed_checksum = hasher.finalize();
181
182        let expected_checksum = hex::decode(&self.crypto.checksum.message)
183            .map_err(|e| KeystoreError::Parse(format!("Invalid checksum hex: {}", e)))?;
184
185        if computed_checksum.as_slice() != expected_checksum.as_slice() {
186            return Err(KeystoreError::InvalidPassword);
187        }
188
189        Ok(())
190    }
191
192    /// Decrypt the secret key bytes
193    fn decrypt_secret(&self, decryption_key: &[u8; 32]) -> Result<[u8; 32], KeystoreError> {
194        let cipher = &self.crypto.cipher;
195
196        if cipher.function != "aes-128-ctr" {
197            return Err(KeystoreError::UnsupportedCipher(cipher.function.clone()));
198        }
199
200        let iv = hex::decode(&cipher.params.iv)
201            .map_err(|e| KeystoreError::Parse(format!("Invalid IV hex: {}", e)))?;
202
203        let ciphertext = hex::decode(&cipher.message)
204            .map_err(|e| KeystoreError::Parse(format!("Invalid ciphertext hex: {}", e)))?;
205
206        // Use first 16 bytes of decryption key for AES-128
207        let aes_key: [u8; 16] = decryption_key[0..16].try_into().unwrap();
208        let iv_array: [u8; 16] = iv.try_into()
209            .map_err(|_| KeystoreError::Parse("Invalid IV length".into()))?;
210
211        let mut cipher = Aes128Ctr::new(&aes_key.into(), &iv_array.into());
212        let mut plaintext = ciphertext;
213        cipher.apply_keystream(&mut plaintext);
214
215        if plaintext.len() != 32 {
216            return Err(KeystoreError::InvalidKey(format!(
217                "Expected 32 bytes, got {}",
218                plaintext.len()
219            )));
220        }
221
222        let mut secret = [0u8; 32];
223        secret.copy_from_slice(&plaintext);
224        Ok(secret)
225    }
226
227    /// Get the public key as bytes
228    pub fn public_key_bytes(&self) -> Result<[u8; 48], KeystoreError> {
229        let pubkey = self.pubkey.strip_prefix("0x").unwrap_or(&self.pubkey);
230        let bytes = hex::decode(pubkey)
231            .map_err(|e| KeystoreError::Parse(format!("Invalid pubkey hex: {}", e)))?;
232
233        if bytes.len() != 48 {
234            return Err(KeystoreError::Parse(format!(
235                "Expected 48 bytes for pubkey, got {}",
236                bytes.len()
237            )));
238        }
239
240        let mut arr = [0u8; 48];
241        arr.copy_from_slice(&bytes);
242        Ok(arr)
243    }
244}
245
246/// Load all keystores from a directory
247pub fn load_keystores_from_dir(
248    dir: impl AsRef<Path>,
249    password: &str,
250) -> Result<Vec<BlsKeypair>, KeystoreError> {
251    let dir = dir.as_ref();
252
253    if !dir.exists() {
254        return Err(KeystoreError::Io(format!(
255            "Directory does not exist: {}",
256            dir.display()
257        )));
258    }
259
260    let mut keypairs = Vec::new();
261
262    for entry in fs::read_dir(dir).map_err(|e| KeystoreError::Io(e.to_string()))? {
263        let entry = entry.map_err(|e| KeystoreError::Io(e.to_string()))?;
264        let path = entry.path();
265
266        if path.extension().and_then(|s| s.to_str()) == Some("json") {
267            match Keystore::load(&path) {
268                Ok(keystore) => {
269                    match keystore.decrypt_keypair(password) {
270                        Ok(keypair) => {
271                            tracing::info!(
272                                pubkey = %keystore.pubkey,
273                                path = %path.display(),
274                                "Loaded validator key"
275                            );
276                            keypairs.push(keypair);
277                        }
278                        Err(e) => {
279                            tracing::warn!(
280                                path = %path.display(),
281                                error = %e,
282                                "Failed to decrypt keystore"
283                            );
284                        }
285                    }
286                }
287                Err(e) => {
288                    tracing::debug!(
289                        path = %path.display(),
290                        error = %e,
291                        "Failed to load keystore"
292                    );
293                }
294            }
295        }
296    }
297
298    Ok(keypairs)
299}
300
301/// Errors related to keystore operations
302#[derive(Debug, Error)]
303pub enum KeystoreError {
304    #[error("I/O error: {0}")]
305    Io(String),
306
307    #[error("Parse error: {0}")]
308    Parse(String),
309
310    #[error("Invalid password")]
311    InvalidPassword,
312
313    #[error("Unsupported KDF: {0}")]
314    UnsupportedKdf(String),
315
316    #[error("Invalid KDF parameters: {0}")]
317    InvalidKdf(String),
318
319    #[error("Unsupported cipher: {0}")]
320    UnsupportedCipher(String),
321
322    #[error("Invalid key: {0}")]
323    InvalidKey(String),
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use std::io::Write;
330    use tempfile::TempDir;
331
332    // Example keystore JSON (scrypt)
333    const TEST_KEYSTORE_SCRYPT: &str = r#"{
334        "crypto": {
335            "kdf": {
336                "function": "scrypt",
337                "params": {
338                    "dklen": 32,
339                    "n": 262144,
340                    "p": 1,
341                    "r": 8,
342                    "salt": "d4e56740f876aef8c010b86a40d5f56745a118d0906a34e69aec8c0db1cb8fa3"
343                },
344                "message": ""
345            },
346            "checksum": {
347                "function": "sha256",
348                "params": {},
349                "message": "d2217fe5f3e9a1e34581ef8a78f7c9928e436d36dacc5e846690a5581e8ea484"
350            },
351            "cipher": {
352                "function": "aes-128-ctr",
353                "params": {
354                    "iv": "264daa3f303d7259501c93d997d84fe6"
355                },
356                "message": "06ae90d55fe0a6e9c5c3bc5b170827b2e5cce3929ed3f116c2811e6366dfe20f"
357            }
358        },
359        "description": "Test keystore",
360        "pubkey": "0x9612d7a727c9d0a22e185a1c768478dfe919cada9266988cb32359c11f2b7b27f4ae4040902382ae2910c15e2b420d07",
361        "path": "m/12381/3600/0/0/0",
362        "uuid": "1d85ae20-35c5-4611-8e4c-00a8c6c5fb55",
363        "version": 4
364    }"#;
365
366    #[test]
367    fn test_parse_keystore() {
368        let keystore: Keystore = serde_json::from_str(TEST_KEYSTORE_SCRYPT).unwrap();
369        assert_eq!(keystore.version, 4);
370        assert!(keystore.pubkey.starts_with("0x"));
371    }
372
373    #[test]
374    fn test_load_keystore_from_file() {
375        let dir = TempDir::new().unwrap();
376        let path = dir.path().join("keystore.json");
377
378        let mut file = fs::File::create(&path).unwrap();
379        file.write_all(TEST_KEYSTORE_SCRYPT.as_bytes()).unwrap();
380
381        let keystore = Keystore::load(&path).unwrap();
382        assert_eq!(keystore.version, 4);
383    }
384
385    #[test]
386    fn test_public_key_bytes() {
387        let keystore: Keystore = serde_json::from_str(TEST_KEYSTORE_SCRYPT).unwrap();
388        let pubkey = keystore.public_key_bytes().unwrap();
389        assert_eq!(pubkey.len(), 48);
390    }
391}