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

iota_keys/
keystore.rs
