rust_bls_bn254/keystores/
base_keystore.rs

1use crate::{
2    consts::UNICODE_CONTROL_CHARS,
3    errors::KeystoreError,
4    sk_to_pk_g2,
5    utils::{pbkdf2, scrypt_key},
6};
7use aes::{
8    cipher::{generic_array::GenericArray, KeyIvInit, StreamCipher},
9    Aes128,
10};
11use ctr::Ctr128BE;
12use num_traits::ToPrimitive;
13use rand::Rng;
14use serde::{Deserialize, Serialize};
15use serde_json::{Map, Value};
16use sha2::{Digest, Sha256};
17use std::{collections::HashMap, fs, io::Write, os::unix::fs::PermissionsExt};
18use unicode_normalization::UnicodeNormalization;
19use uuid::Uuid;
20
21#[derive(Serialize, Deserialize, Debug, Default, PartialEq)]
22pub struct KeystoreModule {
23    pub function: String,
24    #[serde(default)]
25    pub params: HashMap<String, serde_json::Value>,
26    pub(crate) message: String,
27}
28
29#[derive(Serialize, Deserialize, Debug, Default, PartialEq)]
30pub struct KeystoreCrypto {
31    pub kdf: KeystoreModule,
32    pub(crate) checksum: KeystoreModule,
33    pub cipher: KeystoreModule,
34}
35
36impl KeystoreCrypto {
37    fn from_json(json_dict: &Map<String, Value>) -> Result<Self, KeystoreError> {
38        let kdf: KeystoreModule =
39            serde_json::from_value(json_dict["kdf"].clone()).map_err(|e| KeystoreError::from(e))?;
40        let checksum: KeystoreModule = serde_json::from_value(json_dict["checksum"].clone())?;
41        let cipher: KeystoreModule = serde_json::from_value(json_dict["cipher"].clone())?;
42        Ok(Self {
43            kdf,
44            checksum,
45            cipher,
46        })
47    }
48}
49
50/// Keystore is an EIP 2335-compliant keystore. A keystore is a JSON file that
51/// stores an encrypted version of a private key under a user-supplied password.
52///
53/// Ref: https://github.com/ethereum/EIPs/blob/master/EIPS/eip-2335.md
54#[derive(Serialize, Deserialize, Debug, Default, PartialEq)]
55pub struct Keystore {
56    pub crypto: KeystoreCrypto,
57    pub(crate) description: String,
58    pub(crate) pubkey: String,
59    pub path: String,
60    pub(crate) uuid: String,
61    pub(crate) version: u32,
62}
63
64impl Keystore {
65    fn get_u32(param: Option<Value>) -> Result<u32, KeystoreError> {
66        param
67            .ok_or("Missing parameter".into())
68            .and_then(|v| v.as_u64().ok_or("Invalid 'n' parameter".into()))
69            .and_then(|v| v.to_u32().ok_or("Cannot convert 'n' to u32".into()))
70    }
71
72    #[allow(unused)]
73    fn get_u64(param: Option<Value>) -> Result<u64, KeystoreError> {
74        param
75            .ok_or("Missing parameter".into())
76            .and_then(|v| v.as_u64().ok_or("Invalid 'n' parameter".into()))
77    }
78
79    fn kdf(
80        &self,
81        password: &[u8],
82        salt: &[u8],
83        n: u32,
84        r: u32,
85        p: u32,
86        c: u32,
87        dklen: usize,
88    ) -> Result<Vec<u8>, KeystoreError> {
89        match self.crypto.kdf.function.as_str() {
90            "scrypt" => Ok(scrypt_key(password, salt, n, r, p, dklen)?),
91            "pbkdf2" => {
92                let prf = self
93                    .crypto
94                    .kdf
95                    .params
96                    .get("prf")
97                    .and_then(|v| v.as_str())
98                    .ok_or_else(|| {
99                        KeystoreError::GenericError("pubkey not found or not a string".into())
100                    })?;
101                Ok(pbkdf2(password, salt, dklen, c, prf)?)
102            },
103
104            _ => Err(KeystoreError::GenericError(format!(
105                "unsupported function {}",
106                self.crypto.kdf.function
107            ))),
108        }
109    }
110
111    /// save to self as a JSON keystore.
112    pub fn save(&self, file_path: &str) -> Result<(), KeystoreError> {
113        let json_data = serde_json::to_string(self)?;
114        let mut file = fs::File::create(file_path)?;
115        file.write_all(json_data.as_bytes())?;
116        if cfg!(unix) {
117            let mut perms = fs::metadata(file_path)?.permissions();
118            perms.set_mode(0o440);
119            fs::set_permissions(file_path, perms)?;
120        }
121        Ok(())
122    }
123
124    pub fn from_json(
125        json_dict: &HashMap<String, serde_json::Value>,
126    ) -> Result<Self, KeystoreError> {
127        let crypto_dict_object = match json_dict["crypto"].as_object() {
128            None => {
129                return Err(KeystoreError::GenericError(
130                    "crypto dict object not found".to_string(),
131                ))
132            },
133            Some(obj) => obj,
134        };
135        let crypto = KeystoreCrypto::from_json(crypto_dict_object)?;
136        let path = json_dict
137            .get("path")
138            .and_then(|v| v.as_str())
139            .ok_or_else(|| KeystoreError::GenericError("path not found or not a string".into()))?
140            .to_string();
141        let uuid = json_dict
142            .get("uuid")
143            .and_then(|v| v.as_str())
144            .ok_or_else(|| KeystoreError::GenericError("path not found or not a string".into()))?
145            .to_string();
146        let version = Self::get_u32(json_dict.get("version").cloned())?;
147        let description = json_dict
148            .get("description")
149            .and_then(|v| v.as_str())
150            .ok_or_else(|| {
151                KeystoreError::GenericError("Description not found or not a string".into())
152            })?
153            .to_string();
154        let pubkey = json_dict
155            .get("pubkey")
156            .and_then(|v| v.as_str())
157            .ok_or_else(|| KeystoreError::GenericError("pubkey not found or not a string".into()))?
158            .to_string();
159
160        Ok(Self {
161            crypto,
162            description,
163            pubkey,
164            path,
165            uuid,
166            version,
167        })
168    }
169
170    pub fn from_file(path: &str) -> Result<Self, KeystoreError> {
171        let file_content = fs::read_to_string(path).map_err(|e| KeystoreError::from(e))?;
172        let json_dict: HashMap<String, serde_json::Value> = serde_json::from_str(&file_content)?;
173        Ok(Self::from_json(&json_dict)?)
174    }
175
176    /// Encode password as NFKD UTF-8 as per:
177    /// https://github.com/ethereum/EIPs/blob/master/EIPS/eip-2335.md#password-requirements
178    pub fn process_password(password: &str) -> Vec<u8> {
179        let normalized: String = password.nfkd().collect();
180        let filtered: String = normalized
181            .chars()
182            .filter(|c| !UNICODE_CONTROL_CHARS.contains(c))
183            .collect();
184        filtered.as_bytes().to_vec()
185    }
186
187    /// Encrypt a secret (BLS SK) as an EIP 2335 Keystore.
188    pub fn encrypt(
189        &mut self,
190        secret: &[u8],
191        password: &str,
192        path: &str,
193        _kdf_salt: Option<Vec<u8>>,
194        _aes_iv: Option<Vec<u8>>,
195    ) -> Result<(), KeystoreError> {
196        let kdf_salt = match _kdf_salt {
197            Some(salt) => hex::decode(salt)?,
198            None => rand::thread_rng().gen::<[u8; 32]>().to_vec(),
199        };
200
201        let aes_iv = match _aes_iv {
202            Some(iv) => hex::decode(iv)?,
203            None => rand::thread_rng().gen::<[u8; 16]>().to_vec(),
204        };
205
206        self.uuid = Uuid::new_v4().to_string();
207
208        self.crypto.kdf.params.insert(
209            "salt".to_owned(),
210            serde_json::Value::String(hex::encode(&kdf_salt)),
211        );
212        self.crypto.cipher.params.insert(
213            "iv".to_string(),
214            serde_json::Value::String(hex::encode(&aes_iv)),
215        );
216
217        let decryption_key: Vec<u8>;
218        if !self.crypto.kdf.params.contains_key("n")
219            || !self.crypto.kdf.params.contains_key("r")
220            || !self.crypto.kdf.params.contains_key("p")
221        {
222            if !self.crypto.kdf.params.contains_key("c") {
223                return Err(KeystoreError::GenericError(
224                    "params didn't contain parameters for either scrypt or pbkdf2".into(),
225                ));
226            } else {
227                let c = Self::get_u32(self.crypto.kdf.params.get("c").cloned())?;
228                let dklen = Self::get_u32(self.crypto.kdf.params.get("dklen").cloned())? as usize;
229                decryption_key = self.kdf(
230                    &Self::process_password(password),
231                    &kdf_salt,
232                    0,
233                    0,
234                    0,
235                    c,
236                    dklen,
237                )?;
238            }
239        } else {
240            let n = Self::get_u32(self.crypto.kdf.params.get("n").cloned())?;
241            let r = Self::get_u32(self.crypto.kdf.params.get("r").cloned())?;
242            let p = Self::get_u32(self.crypto.kdf.params.get("p").cloned())?;
243            let dklen = Self::get_u32(self.crypto.kdf.params.get("dklen").cloned())? as usize;
244            decryption_key = self.kdf(
245                &Self::process_password(password),
246                &kdf_salt,
247                n,
248                r,
249                p,
250                0,
251                dklen,
252            )?;
253        }
254
255        let key = GenericArray::from_slice(&decryption_key[..16]);
256        let nonce = GenericArray::from_slice(&aes_iv);
257
258        let mut cipher = Ctr128BE::<Aes128>::new(key, nonce);
259        let mut encrypted_secret = secret.to_vec();
260        cipher.apply_keystream(&mut encrypted_secret);
261
262        self.crypto.cipher.message = hex::encode(&encrypted_secret);
263
264        let mut hasher = Sha256::new();
265        hasher.update(&decryption_key[16..32]);
266        hasher.update(&encrypted_secret);
267
268        self.crypto.checksum.message = hex::encode(hasher.finalize());
269
270        self.pubkey = hex::encode(sk_to_pk_g2(secret));
271        self.path = path.to_string();
272
273        Ok(())
274    }
275
276    /// Retrieve the secret (BLS SK) from the self keystore by decrypting it
277    /// with `password`
278    pub fn decrypt(&self, password: &str) -> Result<Vec<u8>, KeystoreError> {
279        let salt = hex::decode(
280            self.crypto
281                .kdf
282                .params
283                .get("salt")
284                .and_then(|v| v.as_str())
285                .ok_or_else(|| KeystoreError::GenericError("salt not found".into()))?,
286        )?;
287
288        let decryption_key: Vec<u8>;
289        if !self.crypto.kdf.params.contains_key("n")
290            || !self.crypto.kdf.params.contains_key("r")
291            || !self.crypto.kdf.params.contains_key("p")
292        {
293            if !self.crypto.kdf.params.contains_key("c") {
294                return Err(KeystoreError::DecryptionError(
295                    "params didn't contain parameters for either scrypt or pbkdf2".into(),
296                ));
297            } else {
298                let c = Self::get_u32(self.crypto.kdf.params.get("c").cloned())?;
299                let dklen = Self::get_u32(self.crypto.kdf.params.get("dklen").cloned())? as usize;
300                decryption_key =
301                    self.kdf(&Self::process_password(password), &salt, 0, 0, 0, c, dklen)?;
302            }
303        } else {
304            let n = Self::get_u32(self.crypto.kdf.params.get("n").cloned())?;
305            let r = Self::get_u32(self.crypto.kdf.params.get("r").cloned())?;
306            let p = Self::get_u32(self.crypto.kdf.params.get("p").cloned())?;
307            let dklen = Self::get_u32(self.crypto.kdf.params.get("dklen").cloned())? as usize;
308            decryption_key =
309                self.kdf(&Self::process_password(password), &salt, n, r, p, 0, dklen)?;
310        }
311
312        let mut hasher = Sha256::new();
313        hasher.update(&decryption_key[16..32]);
314        hasher.update(hex::decode(&self.crypto.cipher.message)?);
315
316        let calculated_checksum = hex::encode(hasher.finalize());
317        if calculated_checksum != self.crypto.checksum.message {
318            return Err(KeystoreError::DecryptionError(
319                "Checksum message error".into(),
320            ));
321        }
322
323        let key = GenericArray::from_slice(&decryption_key[..16]);
324        let iv_hex = self
325            .crypto
326            .cipher
327            .params
328            .get("iv")
329            .ok_or(KeystoreError::DecryptionError(
330                "IV not found in cipher params".into(),
331            ))?;
332        let iv = hex::decode(
333            iv_hex
334                .as_str()
335                .ok_or(KeystoreError::DecryptionError("IV decode error".into()))?,
336        )?;
337        let nonce = GenericArray::from_slice(&iv);
338        let mut cipher = Ctr128BE::<Aes128>::new(key, nonce);
339        let mut decrypted_message = hex::decode(&self.crypto.cipher.message)?;
340        cipher.apply_keystream(&mut decrypted_message);
341        Ok(decrypted_message.to_vec())
342    }
343}