iota_keys/
keystore.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::{BTreeMap, HashSet},
7    fmt::{Display, Formatter, Write},
8    fs,
9    fs::File,
10    io::BufReader,
11    path::{Path, PathBuf},
12};
13
14use anyhow::{Context, anyhow, bail, ensure};
15use bip32::DerivationPath;
16use bip39::{Language, Mnemonic, Seed};
17use iota_types::{
18    base_types::IotaAddress,
19    crypto::{
20        EncodeDecodeBase64, IotaKeyPair, PublicKey, Signature, SignatureScheme, enum_dispatch,
21        get_key_pair_from_rng,
22    },
23};
24use rand::{SeedableRng, rngs::StdRng};
25use regex::Regex;
26use serde::{Deserialize, Deserializer, Serialize, Serializer};
27use shared_crypto::intent::{Intent, IntentMessage};
28
29use crate::{
30    key_derive::{derive_key_pair_from_path, generate_new_key},
31    random_names::{random_name, random_names},
32};
33
34#[derive(Serialize, Deserialize)]
35#[enum_dispatch(AccountKeystore)]
36pub enum Keystore {
37    File(FileBasedKeystore),
38    InMem(InMemKeystore),
39}
40#[enum_dispatch]
41pub trait AccountKeystore: Send + Sync {
42    fn add_key(&mut self, alias: Option<String>, keypair: IotaKeyPair)
43    -> Result<(), anyhow::Error>;
44    fn keys(&self) -> Vec<PublicKey>;
45    fn get_key(&self, address: &IotaAddress) -> Result<&IotaKeyPair, anyhow::Error>;
46
47    fn sign_hashed(&self, address: &IotaAddress, msg: &[u8])
48    -> Result<Signature, signature::Error>;
49
50    fn sign_secure<T>(
51        &self,
52        address: &IotaAddress,
53        msg: &T,
54        intent: Intent,
55    ) -> Result<Signature, signature::Error>
56    where
57        T: Serialize;
58    fn addresses(&self) -> Vec<IotaAddress> {
59        self.keys().iter().map(|k| k.into()).collect()
60    }
61    fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)>;
62    fn aliases(&self) -> Vec<&Alias>;
63    fn aliases_mut(&mut self) -> Vec<&mut Alias>;
64    fn alias_names(&self) -> Vec<&str> {
65        self.aliases()
66            .into_iter()
67            .map(|a| a.alias.as_str())
68            .collect()
69    }
70    /// Get alias of address
71    fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error>;
72    fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error>;
73    /// Check if an alias exists by its name
74    fn alias_exists(&self, alias: &str) -> bool {
75        self.alias_names().contains(&alias)
76    }
77
78    fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error>;
79
80    fn update_alias(
81        &mut self,
82        old_alias: &str,
83        new_alias: Option<&str>,
84    ) -> Result<String, anyhow::Error>;
85
86    // Internal function. Use update_alias instead
87    fn update_alias_value(
88        &mut self,
89        old_alias: &str,
90        new_alias: Option<&str>,
91    ) -> Result<String, anyhow::Error> {
92        if !self.alias_exists(old_alias) {
93            bail!("The provided alias {old_alias} does not exist");
94        }
95
96        let new_alias_name = self.create_alias(new_alias.map(str::to_string))?;
97
98        for a in self.aliases_mut() {
99            if a.alias == old_alias {
100                let pk = &a.public_key_base64;
101                *a = Alias {
102                    alias: new_alias_name.clone(),
103                    public_key_base64: pk.clone(),
104                };
105            }
106        }
107        Ok(new_alias_name)
108    }
109
110    fn generate_and_add_new_key(
111        &mut self,
112        key_scheme: SignatureScheme,
113        alias: Option<String>,
114        derivation_path: Option<DerivationPath>,
115        word_length: Option<String>,
116    ) -> Result<(IotaAddress, String, SignatureScheme), anyhow::Error> {
117        let (address, kp, scheme, phrase) =
118            generate_new_key(key_scheme, derivation_path, word_length)?;
119        self.add_key(alias, kp)?;
120        Ok((address, phrase, scheme))
121    }
122
123    fn import_from_mnemonic(
124        &mut self,
125        phrase: &str,
126        key_scheme: SignatureScheme,
127        derivation_path: Option<DerivationPath>,
128        alias: Option<String>,
129    ) -> Result<IotaAddress, anyhow::Error> {
130        let mnemonic = Mnemonic::from_phrase(phrase, Language::English)
131            .map_err(|e| anyhow::anyhow!("Invalid mnemonic phrase: {:?}", e))?;
132        let seed = Seed::new(&mnemonic, "");
133        self.import_from_seed(seed.as_bytes(), key_scheme, derivation_path, alias)
134    }
135
136    fn import_from_seed(
137        &mut self,
138        seed: &[u8],
139        key_scheme: SignatureScheme,
140        derivation_path: Option<DerivationPath>,
141        alias: Option<String>,
142    ) -> Result<IotaAddress, anyhow::Error> {
143        match derive_key_pair_from_path(seed, derivation_path, &key_scheme) {
144            Ok((address, kp)) => {
145                self.add_key(alias, kp)?;
146                Ok(address)
147            }
148            Err(e) => Err(anyhow!("error getting keypair {:?}", e)),
149        }
150    }
151}
152
153impl Display for Keystore {
154    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
155        let mut writer = String::new();
156        match self {
157            Keystore::File(file) => {
158                writeln!(writer, "Keystore Type: File")?;
159                write!(writer, "Keystore Path : {:?}", file.path)?;
160                write!(f, "{}", writer)
161            }
162            Keystore::InMem(_) => {
163                writeln!(writer, "Keystore Type: InMem")?;
164                write!(f, "{}", writer)
165            }
166        }
167    }
168}
169
170#[derive(Serialize, Deserialize, Clone, Debug)]
171pub struct Alias {
172    pub alias: String,
173    pub public_key_base64: String,
174}
175
176#[derive(Default)]
177pub struct FileBasedKeystore {
178    keys: BTreeMap<IotaAddress, IotaKeyPair>,
179    aliases: BTreeMap<IotaAddress, Alias>,
180    path: PathBuf,
181}
182
183impl Serialize for FileBasedKeystore {
184    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
185    where
186        S: Serializer,
187    {
188        serializer.serialize_str(self.path.to_str().unwrap_or(""))
189    }
190}
191
192impl<'de> Deserialize<'de> for FileBasedKeystore {
193    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
194    where
195        D: Deserializer<'de>,
196    {
197        use serde::de::Error;
198        FileBasedKeystore::new(&PathBuf::from(String::deserialize(deserializer)?))
199            .map_err(D::Error::custom)
200    }
201}
202
203impl AccountKeystore for FileBasedKeystore {
204    fn sign_hashed(
205        &self,
206        address: &IotaAddress,
207        msg: &[u8],
208    ) -> Result<Signature, signature::Error> {
209        Ok(Signature::new_hashed(
210            msg,
211            self.keys.get(address).ok_or_else(|| {
212                signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
213            })?,
214        ))
215    }
216    fn sign_secure<T>(
217        &self,
218        address: &IotaAddress,
219        msg: &T,
220        intent: Intent,
221    ) -> Result<Signature, signature::Error>
222    where
223        T: Serialize,
224    {
225        Ok(Signature::new_secure(
226            &IntentMessage::new(intent, msg),
227            self.keys.get(address).ok_or_else(|| {
228                signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
229            })?,
230        ))
231    }
232
233    fn add_key(
234        &mut self,
235        alias: Option<String>,
236        keypair: IotaKeyPair,
237    ) -> Result<(), anyhow::Error> {
238        let address: IotaAddress = (&keypair.public()).into();
239        let alias = self.create_alias(alias)?;
240        self.aliases.insert(
241            address,
242            Alias {
243                alias,
244                public_key_base64: keypair.public().encode_base64(),
245            },
246        );
247        self.keys.insert(address, keypair);
248        self.save()?;
249        Ok(())
250    }
251
252    /// Return an array of `Alias`, consisting of every alias and its
253    /// corresponding public key.
254    fn aliases(&self) -> Vec<&Alias> {
255        self.aliases.values().collect()
256    }
257
258    fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)> {
259        self.aliases.iter().collect::<Vec<_>>()
260    }
261
262    /// Return an array of `Alias`, consisting of every alias and its
263    /// corresponding public key.
264    fn aliases_mut(&mut self) -> Vec<&mut Alias> {
265        self.aliases.values_mut().collect()
266    }
267
268    fn keys(&self) -> Vec<PublicKey> {
269        self.keys.values().map(|key| key.public()).collect()
270    }
271
272    /// This function returns an error if the provided alias already exists. If
273    /// the alias has not already been used, then it returns the alias.
274    /// If no alias has been passed, it will generate a new alias.
275    fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error> {
276        match alias {
277            Some(a) if self.alias_exists(&a) => {
278                bail!("Alias {a} already exists. Please choose another alias.")
279            }
280            Some(a) => validate_alias(&a),
281            None => Ok(random_name(
282                &self
283                    .alias_names()
284                    .into_iter()
285                    .map(|x| x.to_string())
286                    .collect::<HashSet<_>>(),
287            )),
288        }
289    }
290
291    /// Get the address by its alias
292    fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error> {
293        self.addresses_with_alias()
294            .iter()
295            .find(|x| x.1.alias == alias)
296            .ok_or_else(|| anyhow!("Cannot resolve alias {alias} to an address"))
297            .map(|x| x.0)
298    }
299
300    /// Get the alias if it exists, or return an error if it does not exist.
301    fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error> {
302        match self.aliases.get(address) {
303            Some(alias) => Ok(alias.alias.clone()),
304            None => bail!("Cannot find alias for address {address}"),
305        }
306    }
307
308    fn get_key(&self, address: &IotaAddress) -> Result<&IotaKeyPair, anyhow::Error> {
309        match self.keys.get(address) {
310            Some(key) => Ok(key),
311            None => Err(anyhow!("Cannot find key for address: [{address}]")),
312        }
313    }
314
315    /// Updates an old alias to the new alias and saves it to the alias file.
316    /// If the new_alias is None, it will generate a new random alias.
317    fn update_alias(
318        &mut self,
319        old_alias: &str,
320        new_alias: Option<&str>,
321    ) -> Result<String, anyhow::Error> {
322        let new_alias_name = self.update_alias_value(old_alias, new_alias)?;
323        self.save_aliases()?;
324        Ok(new_alias_name)
325    }
326}
327
328impl FileBasedKeystore {
329    pub fn new(path: &PathBuf) -> Result<Self, anyhow::Error> {
330        let keys = if path.exists() {
331            let reader =
332                BufReader::new(File::open(path).with_context(|| {
333                    format!("Cannot open the keystore file: {}", path.display())
334                })?);
335            let kp_strings: Vec<String> = serde_json::from_reader(reader).with_context(|| {
336                format!("Cannot deserialize the keystore file: {}", path.display(),)
337            })?;
338            kp_strings
339                .iter()
340                .map(|kpstr| {
341                    let key = IotaKeyPair::decode(kpstr);
342                    key.map(|k| (IotaAddress::from(&k.public()), k))
343                })
344                .collect::<Result<BTreeMap<_, _>, _>>()
345                .map_err(|e| anyhow!("Invalid keystore file: {}. {}", path.display(), e))?
346        } else {
347            BTreeMap::new()
348        };
349
350        // check aliases
351        let mut aliases_path = path.clone();
352        aliases_path.set_extension("aliases");
353
354        let aliases = if aliases_path.exists() {
355            let reader = BufReader::new(File::open(&aliases_path).with_context(|| {
356                format!(
357                    "Cannot open aliases file in keystore: {}",
358                    aliases_path.display()
359                )
360            })?);
361
362            let aliases: Vec<Alias> = serde_json::from_reader(reader).with_context(|| {
363                format!(
364                    "Cannot deserialize aliases file in keystore: {}",
365                    aliases_path.display(),
366                )
367            })?;
368
369            aliases
370                .into_iter()
371                .map(|alias| {
372                    let key = PublicKey::decode_base64(&alias.public_key_base64);
373                    key.map(|k| (Into::<IotaAddress>::into(&k), alias))
374                })
375                .collect::<Result<BTreeMap<_, _>, _>>()
376                .map_err(|e| {
377                    anyhow!(
378                        "Invalid aliases file in keystore: {}. {}",
379                        aliases_path.display(),
380                        e
381                    )
382                })?
383        } else if keys.is_empty() {
384            BTreeMap::new()
385        } else {
386            let names: Vec<String> = random_names(HashSet::new(), keys.len());
387            let aliases = keys
388                .iter()
389                .zip(names)
390                .map(|((iota_address, ikp), alias)| {
391                    let public_key_base64 = ikp.public().encode_base64();
392                    (
393                        *iota_address,
394                        Alias {
395                            alias,
396                            public_key_base64,
397                        },
398                    )
399                })
400                .collect::<BTreeMap<_, _>>();
401            let aliases_store = serde_json::to_string_pretty(&aliases.values().collect::<Vec<_>>())
402                .with_context(|| {
403                    format!(
404                        "Cannot serialize aliases to file in keystore: {}",
405                        aliases_path.display()
406                    )
407                })?;
408            fs::write(aliases_path, aliases_store)?;
409            aliases
410        };
411
412        Ok(Self {
413            keys,
414            aliases,
415            path: path.to_path_buf(),
416        })
417    }
418
419    pub fn set_path(&mut self, path: &Path) {
420        self.path = path.to_path_buf();
421    }
422
423    pub fn save_aliases(&self) -> Result<(), anyhow::Error> {
424        let aliases_store = serde_json::to_string_pretty(
425            &self.aliases.values().collect::<Vec<_>>(),
426        )
427        .with_context(|| {
428            format!(
429                "Cannot serialize aliases to file in keystore: {}",
430                self.path.display()
431            )
432        })?;
433
434        let mut aliases_path = self.path.clone();
435        aliases_path.set_extension("aliases");
436        fs::write(&aliases_path, aliases_store)
437            .map_err(|e| anyhow!("Couldn't save aliases to {}: {e}", aliases_path.display()))?;
438        Ok(())
439    }
440
441    /// Keys saved as Base64 with 33 bytes `flag || privkey` ($BASE64_STR).
442    /// To see Bech32 format encoding, use `iota keytool export $IOTA_ADDRESS`
443    /// where $IOTA_ADDRESS can be found with `iota keytool list`. Or use
444    /// `iota keytool convert $BASE64_STR`
445    pub fn save_keystore(&self) -> Result<(), anyhow::Error> {
446        let store = serde_json::to_string_pretty(
447            &self
448                .keys
449                .values()
450                .map(|k| k.encode())
451                .collect::<Result<Vec<_>, _>>()
452                .map_err(|e| anyhow!(e))?,
453        )
454        .with_context(|| format!("Cannot serialize keystore to file: {}", self.path.display()))?;
455        fs::write(&self.path, store)
456            .map_err(|e| anyhow!("Couldn't save keystore to {}: {e}", self.path.display()))?;
457        Ok(())
458    }
459
460    pub fn save(&self) -> Result<(), anyhow::Error> {
461        self.save_aliases()?;
462        self.save_keystore()?;
463        Ok(())
464    }
465
466    pub fn key_pairs(&self) -> Vec<&IotaKeyPair> {
467        self.keys.values().collect()
468    }
469}
470
471#[derive(Default, Serialize, Deserialize)]
472pub struct InMemKeystore {
473    aliases: BTreeMap<IotaAddress, Alias>,
474    keys: BTreeMap<IotaAddress, IotaKeyPair>,
475}
476
477impl AccountKeystore for InMemKeystore {
478    fn sign_hashed(
479        &self,
480        address: &IotaAddress,
481        msg: &[u8],
482    ) -> Result<Signature, signature::Error> {
483        Ok(Signature::new_hashed(
484            msg,
485            self.keys.get(address).ok_or_else(|| {
486                signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
487            })?,
488        ))
489    }
490    fn sign_secure<T>(
491        &self,
492        address: &IotaAddress,
493        msg: &T,
494        intent: Intent,
495    ) -> Result<Signature, signature::Error>
496    where
497        T: Serialize,
498    {
499        Ok(Signature::new_secure(
500            &IntentMessage::new(intent, msg),
501            self.keys.get(address).ok_or_else(|| {
502                signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
503            })?,
504        ))
505    }
506
507    fn add_key(
508        &mut self,
509        alias: Option<String>,
510        keypair: IotaKeyPair,
511    ) -> Result<(), anyhow::Error> {
512        let address: IotaAddress = (&keypair.public()).into();
513        let alias = alias.unwrap_or_else(|| {
514            random_name(
515                &self
516                    .aliases()
517                    .iter()
518                    .map(|x| x.alias.clone())
519                    .collect::<HashSet<_>>(),
520            )
521        });
522
523        let public_key_base64 = keypair.public().encode_base64();
524        let alias = Alias {
525            alias,
526            public_key_base64,
527        };
528        self.aliases.insert(address, alias);
529        self.keys.insert(address, keypair);
530        Ok(())
531    }
532
533    /// Get all aliases objects
534    fn aliases(&self) -> Vec<&Alias> {
535        self.aliases.values().collect()
536    }
537
538    fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)> {
539        self.aliases.iter().collect::<Vec<_>>()
540    }
541
542    fn keys(&self) -> Vec<PublicKey> {
543        self.keys.values().map(|key| key.public()).collect()
544    }
545
546    fn get_key(&self, address: &IotaAddress) -> Result<&IotaKeyPair, anyhow::Error> {
547        match self.keys.get(address) {
548            Some(key) => Ok(key),
549            None => Err(anyhow!("Cannot find key for address: [{address}]")),
550        }
551    }
552
553    /// Get alias of address
554    fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error> {
555        match self.aliases.get(address) {
556            Some(alias) => Ok(alias.alias.clone()),
557            None => bail!("Cannot find alias for address {address}"),
558        }
559    }
560
561    /// Get the address by its alias
562    fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error> {
563        self.addresses_with_alias()
564            .iter()
565            .find(|x| x.1.alias == alias)
566            .ok_or_else(|| anyhow!("Cannot resolve alias {alias} to an address"))
567            .map(|x| x.0)
568    }
569
570    /// This function returns an error if the provided alias already exists. If
571    /// the alias has not already been used, then it returns the alias.
572    /// If no alias has been passed, it will generate a new alias.
573    fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error> {
574        match alias {
575            Some(a) if self.alias_exists(&a) => {
576                bail!("Alias {a} already exists. Please choose another alias.")
577            }
578            Some(a) => validate_alias(&a),
579            None => Ok(random_name(
580                &self
581                    .alias_names()
582                    .into_iter()
583                    .map(|x| x.to_string())
584                    .collect::<HashSet<_>>(),
585            )),
586        }
587    }
588
589    fn aliases_mut(&mut self) -> Vec<&mut Alias> {
590        self.aliases.values_mut().collect()
591    }
592
593    /// Updates an old alias to the new alias. If the new_alias is None,
594    /// it will generate a new random alias.
595    fn update_alias(
596        &mut self,
597        old_alias: &str,
598        new_alias: Option<&str>,
599    ) -> Result<String, anyhow::Error> {
600        self.update_alias_value(old_alias, new_alias)
601    }
602}
603
604impl InMemKeystore {
605    pub fn new_insecure_for_tests(initial_key_number: usize) -> Self {
606        let mut rng = StdRng::from_seed([0; 32]);
607        let keys = (0..initial_key_number)
608            .map(|_| get_key_pair_from_rng(&mut rng))
609            .map(|(ad, k)| (ad, IotaKeyPair::Ed25519(k)))
610            .collect::<BTreeMap<IotaAddress, IotaKeyPair>>();
611
612        let aliases = keys
613            .iter()
614            .zip(random_names(HashSet::new(), keys.len()))
615            .map(|((iota_address, ikp), alias)| {
616                let public_key_base64 = ikp.public().encode_base64();
617                (
618                    *iota_address,
619                    Alias {
620                        alias,
621                        public_key_base64,
622                    },
623                )
624            })
625            .collect::<BTreeMap<_, _>>();
626
627        Self { aliases, keys }
628    }
629}
630
631fn validate_alias(alias: &str) -> Result<String, anyhow::Error> {
632    let re = Regex::new(r"^[A-Za-z][A-Za-z0-9-_\.]*$")
633        .map_err(|_| anyhow!("Cannot build the regex needed to validate the alias naming"))?;
634    let alias = alias.trim();
635    ensure!(
636        re.is_match(alias),
637        "Invalid alias. A valid alias must start with a letter and can contain only letters, digits, hyphens (-), dots (.), or underscores (_)."
638    );
639    Ok(alias.to_string())
640}
641
642#[cfg(test)]
643mod tests {
644    use crate::keystore::validate_alias;
645
646    #[test]
647    fn validate_alias_test() {
648        // OK
649        assert!(validate_alias("A.B_dash").is_ok());
650        assert!(validate_alias("A.B-C1_dash").is_ok());
651        assert!(validate_alias("abc_123.iota").is_ok());
652        // Not allowed
653        assert!(validate_alias("A.B-C_dash!").is_err());
654        assert!(validate_alias(".B-C_dash!").is_err());
655        assert!(validate_alias("_test").is_err());
656        assert!(validate_alias("123").is_err());
657        assert!(validate_alias("@@123").is_err());
658        assert!(validate_alias("@_Ab").is_err());
659        assert!(validate_alias("_Ab").is_err());
660        assert!(validate_alias("^A").is_err());
661        assert!(validate_alias("-A").is_err());
662    }
663}
iota_keys/keystore.rs

iota_keys/
keystore.rs
