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 serde_with::{DisplayFromStr, serde_as};
28use shared_crypto::intent::{Intent, IntentMessage};
29use tracing::{debug, info};
30
31use crate::{
32    key_derive::{derive_key_pair_from_path, generate_new_key},
33    random_names::{random_name, random_names},
34    serde_iota_keypair, serde_public_key,
35};
36
37#[derive(Serialize, Deserialize)]
38#[enum_dispatch(AccountKeystore)]
39pub enum Keystore {
40    File(FileBasedKeystore),
41    InMem(InMemKeystore),
42}
43
44#[enum_dispatch]
45pub trait AccountKeystore: Send + Sync {
46    fn add_key(
47        &mut self,
48        alias: Option<String>,
49        key: impl Into<StoredKey>,
50    ) -> Result<(), anyhow::Error>;
51    fn remove_key(&mut self, address: &IotaAddress) -> Result<(), anyhow::Error>;
52    fn keys(&self) -> Vec<&StoredKey>;
53    fn get_key(&self, address: &IotaAddress) -> Result<&StoredKey, anyhow::Error>;
54
55    fn sign_hashed(&self, address: &IotaAddress, msg: &[u8])
56    -> Result<Signature, signature::Error>;
57
58    fn sign_secure<T>(
59        &self,
60        address: &IotaAddress,
61        msg: &T,
62        intent: Intent,
63    ) -> Result<Signature, signature::Error>
64    where
65        T: Serialize;
66    fn addresses(&self) -> Vec<IotaAddress> {
67        self.keys().into_iter().map(|k| k.address()).collect()
68    }
69    fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)>;
70    fn aliases(&self) -> Vec<&Alias>;
71    fn aliases_mut(&mut self) -> Vec<&mut Alias>;
72    fn alias_names(&self) -> Vec<&str> {
73        self.aliases()
74            .into_iter()
75            .map(|a| a.alias.as_str())
76            .collect()
77    }
78    /// Get alias of address
79    fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error>;
80    fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error>;
81    /// Check if an alias exists by its name
82    fn alias_exists(&self, alias: &str) -> bool {
83        self.alias_names().contains(&alias)
84    }
85
86    fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error>;
87
88    fn update_alias(
89        &mut self,
90        old_alias: &str,
91        new_alias: Option<&str>,
92    ) -> Result<String, anyhow::Error>;
93
94    // Internal function. Use update_alias instead
95    fn update_alias_value(
96        &mut self,
97        old_alias: &str,
98        new_alias: Option<&str>,
99    ) -> Result<String, anyhow::Error> {
100        if !self.alias_exists(old_alias) {
101            bail!("The provided alias {old_alias} does not exist");
102        }
103
104        let new_alias_name = self.create_alias(new_alias.map(str::to_string))?;
105
106        for a in self.aliases_mut() {
107            if a.alias == old_alias {
108                let pk = &a.public_key_base64;
109                *a = Alias {
110                    alias: new_alias_name.clone(),
111                    public_key_base64: pk.clone(),
112                };
113            }
114        }
115        Ok(new_alias_name)
116    }
117
118    fn generate_and_add_new_key(
119        &mut self,
120        key_scheme: SignatureScheme,
121        alias: Option<String>,
122        derivation_path: Option<DerivationPath>,
123        word_length: Option<String>,
124    ) -> Result<(IotaAddress, String, SignatureScheme), anyhow::Error> {
125        let (address, kp, scheme, phrase) =
126            generate_new_key(key_scheme, derivation_path, word_length)?;
127        self.add_key(alias, kp)?;
128        Ok((address, phrase, scheme))
129    }
130
131    fn import_from_mnemonic(
132        &mut self,
133        phrase: &str,
134        key_scheme: SignatureScheme,
135        derivation_path: Option<DerivationPath>,
136        alias: Option<String>,
137    ) -> Result<IotaAddress, anyhow::Error> {
138        let mnemonic = Mnemonic::from_phrase(phrase, Language::English)
139            .map_err(|e| anyhow::anyhow!("Invalid mnemonic phrase: {:?}", e))?;
140        let seed = Seed::new(&mnemonic, "");
141        self.import_from_seed(seed.as_bytes(), key_scheme, derivation_path, alias)
142    }
143
144    fn import_from_seed(
145        &mut self,
146        seed: &[u8],
147        key_scheme: SignatureScheme,
148        derivation_path: Option<DerivationPath>,
149        alias: Option<String>,
150    ) -> Result<IotaAddress, anyhow::Error> {
151        match derive_key_pair_from_path(seed, derivation_path, &key_scheme) {
152            Ok((address, kp)) => {
153                self.add_key(alias, kp)?;
154                Ok(address)
155            }
156            Err(e) => Err(anyhow!("error getting keypair {:?}", e)),
157        }
158    }
159
160    fn import_from_external(
161        &mut self,
162        source: &str,
163        public_key: PublicKey,
164        derivation_path: Option<DerivationPath>,
165        alias: Option<String>,
166    ) -> Result<(), anyhow::Error> {
167        self.add_key(
168            alias,
169            StoredKey::External {
170                derivation_path,
171                public_key,
172                source: source.to_string(),
173            },
174        )
175    }
176}
177
178impl Display for Keystore {
179    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
180        let mut writer = String::new();
181        match self {
182            Keystore::File(file) => {
183                writeln!(writer, "Keystore Type: File")?;
184                write!(writer, "Keystore Path : {:?}", file.path)?;
185                write!(f, "{writer}")
186            }
187            Keystore::InMem(_) => {
188                writeln!(writer, "Keystore Type: InMem")?;
189                write!(f, "{writer}")
190            }
191        }
192    }
193}
194
195#[derive(Serialize, Deserialize, Clone, Debug)]
196pub struct Alias {
197    pub alias: String,
198    pub public_key_base64: String,
199}
200
201#[expect(clippy::large_enum_variant)]
202#[serde_as]
203#[derive(Serialize, Deserialize, Debug, Clone)]
204#[serde(
205    tag = "type",            // this makes {"type": "...", "value": …}
206    content = "value",       // name the payload field "value"
207    rename_all = "snake_case"
208)]
209pub enum StoredKey {
210    #[serde(with = "serde_iota_keypair")]
211    KeyPair(IotaKeyPair),
212    External {
213        source: String,
214        #[serde_as(as = "Option<DisplayFromStr>")]
215        #[serde(skip_serializing_if = "Option::is_none")]
216        derivation_path: Option<DerivationPath>,
217        #[serde(rename = "public_key_base64_with_flag", with = "serde_public_key")]
218        public_key: PublicKey,
219    },
220}
221
222impl From<IotaKeyPair> for StoredKey {
223    fn from(keypair: IotaKeyPair) -> Self {
224        StoredKey::KeyPair(keypair)
225    }
226}
227
228impl StoredKey {
229    pub fn address(&self) -> IotaAddress {
230        IotaAddress::from(&self.public())
231    }
232
233    pub fn public(&self) -> PublicKey {
234        match self {
235            StoredKey::KeyPair(keypair) => keypair.public(),
236            StoredKey::External { public_key, .. } => public_key.clone(),
237        }
238    }
239
240    pub fn derivation_path(&self) -> Option<DerivationPath> {
241        match self {
242            StoredKey::KeyPair(_) => None,
243            StoredKey::External {
244                derivation_path, ..
245            } => derivation_path.clone(),
246        }
247    }
248
249    pub fn external_source(&self) -> Option<String> {
250        match self {
251            StoredKey::KeyPair(_) => None,
252            StoredKey::External { source, .. } => Some(source.clone()),
253        }
254    }
255
256    pub fn as_keypair(&self) -> Result<&IotaKeyPair, anyhow::Error> {
257        match self {
258            StoredKey::KeyPair(keypair) => Ok(keypair),
259            StoredKey::External { .. } => bail!("Cannot get key pair for External keys."),
260        }
261    }
262}
263
264#[derive(Default)]
265pub struct FileBasedKeystore {
266    keys: BTreeMap<IotaAddress, StoredKey>,
267    aliases: BTreeMap<IotaAddress, Alias>,
268    path: PathBuf,
269}
270
271impl Serialize for FileBasedKeystore {
272    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
273    where
274        S: Serializer,
275    {
276        serializer.serialize_str(self.path.to_str().unwrap_or(""))
277    }
278}
279
280impl<'de> Deserialize<'de> for FileBasedKeystore {
281    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
282    where
283        D: Deserializer<'de>,
284    {
285        use serde::de::Error;
286        FileBasedKeystore::new(&PathBuf::from(String::deserialize(deserializer)?))
287            .map_err(D::Error::custom)
288    }
289}
290#[derive(Serialize, Deserialize, Debug)]
291pub struct FileBasedKeystoreFile {
292    pub version: u8,
293    pub keys: Vec<AliasedKey>,
294}
295
296#[derive(Serialize, Deserialize, Debug)]
297pub struct AliasedKey {
298    pub alias: String,
299    pub address: IotaAddress,
300    pub key: StoredKey,
301}
302
303impl AccountKeystore for FileBasedKeystore {
304    fn sign_hashed(
305        &self,
306        address: &IotaAddress,
307        msg: &[u8],
308    ) -> Result<Signature, signature::Error> {
309        let stored_key = self.keys.get(address).ok_or_else(|| {
310            signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
311        })?;
312
313        match stored_key {
314            StoredKey::KeyPair(keypair) => Ok(Signature::new_hashed(msg, keypair)),
315            StoredKey::External { source, .. } => Err(signature::Error::from_source(format!(
316                "sign_hashed is not supported for external type: {source} [{address}]"
317            ))),
318        }
319    }
320    fn sign_secure<T>(
321        &self,
322        address: &IotaAddress,
323        msg: &T,
324        intent: Intent,
325    ) -> Result<Signature, signature::Error>
326    where
327        T: Serialize,
328    {
329        let stored_key = self.keys.get(address).ok_or_else(|| {
330            signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
331        })?;
332
333        let intent_msg = &IntentMessage::new(intent, msg);
334        match stored_key {
335            StoredKey::KeyPair(keypair) => Ok(Signature::new_secure(intent_msg, keypair)),
336            StoredKey::External { source, .. } => Err(signature::Error::from_source(format!(
337                "sign_secure is not supported for external type: {source} [{address}]",
338            ))),
339        }
340    }
341
342    fn add_key(
343        &mut self,
344        alias: Option<String>,
345        key: impl Into<StoredKey>,
346    ) -> Result<(), anyhow::Error> {
347        let key = key.into();
348        let address: IotaAddress = (&key.public()).into();
349        let alias = self.create_alias(alias)?;
350        self.aliases.insert(
351            address,
352            Alias {
353                alias,
354                public_key_base64: key.public().encode_base64(),
355            },
356        );
357        self.keys.insert(address, key);
358        self.save()?;
359        Ok(())
360    }
361
362    fn remove_key(&mut self, address: &IotaAddress) -> Result<(), anyhow::Error> {
363        self.aliases.remove(address);
364        self.keys.remove(address);
365        self.save()?;
366        Ok(())
367    }
368
369    /// Return an array of `Alias`, consisting of every alias and its
370    /// corresponding public key.
371    fn aliases(&self) -> Vec<&Alias> {
372        self.aliases.values().collect()
373    }
374
375    fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)> {
376        self.aliases.iter().collect::<Vec<_>>()
377    }
378
379    /// Return an array of `Alias`, consisting of every alias and its
380    /// corresponding public key.
381    fn aliases_mut(&mut self) -> Vec<&mut Alias> {
382        self.aliases.values_mut().collect()
383    }
384
385    fn keys(&self) -> Vec<&StoredKey> {
386        self.keys.values().collect()
387    }
388
389    /// This function returns an error if the provided alias already exists. If
390    /// the alias has not already been used, then it returns the alias.
391    /// If no alias has been passed, it will generate a new alias.
392    fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error> {
393        match alias {
394            Some(a) if self.alias_exists(&a) => {
395                bail!("Alias {a} already exists. Please choose another alias.")
396            }
397            Some(a) => validate_alias(&a),
398            None => Ok(random_name(
399                &self
400                    .alias_names()
401                    .into_iter()
402                    .map(|x| x.to_string())
403                    .collect::<HashSet<_>>(),
404            )),
405        }
406    }
407
408    /// Get the address by its alias
409    fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error> {
410        self.addresses_with_alias()
411            .iter()
412            .find(|x| x.1.alias == alias)
413            .ok_or_else(|| anyhow!("Cannot resolve alias {alias} to an address"))
414            .map(|x| x.0)
415    }
416
417    /// Get the alias if it exists, or return an error if it does not exist.
418    fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error> {
419        match self.aliases.get(address) {
420            Some(alias) => Ok(alias.alias.clone()),
421            None => bail!("Cannot find alias for address {address}"),
422        }
423    }
424
425    fn get_key(&self, address: &IotaAddress) -> Result<&StoredKey, anyhow::Error> {
426        match self.keys.get(address) {
427            Some(key) => Ok(key),
428            None => Err(anyhow!("Cannot find key for address: [{address}]")),
429        }
430    }
431
432    /// Updates an old alias to the new alias and saves it to the alias file.
433    /// If the new_alias is None, it will generate a new random alias.
434    fn update_alias(
435        &mut self,
436        old_alias: &str,
437        new_alias: Option<&str>,
438    ) -> Result<String, anyhow::Error> {
439        let new_alias_name = self.update_alias_value(old_alias, new_alias)?;
440        self.save()?;
441        Ok(new_alias_name)
442    }
443}
444
445impl FileBasedKeystore {
446    pub fn new_from_v1(path: &PathBuf) -> Result<Self, anyhow::Error> {
447        let keys = if path.exists() {
448            let reader =
449                BufReader::new(File::open(path).with_context(|| {
450                    format!("Cannot open the keystore file: {}", path.display())
451                })?);
452            let kp_strings: Vec<String> = serde_json::from_reader(reader).with_context(|| {
453                format!("Cannot deserialize the keystore file: {}", path.display(),)
454            })?;
455            kp_strings
456                .iter()
457                .map(|kpstr| {
458                    let key = IotaKeyPair::decode(kpstr);
459                    key.map(|k| (IotaAddress::from(&k.public()), StoredKey::KeyPair(k)))
460                })
461                .collect::<Result<BTreeMap<_, _>, _>>()
462                .map_err(|e| anyhow!("Invalid keystore file: {}. {}", path.display(), e))?
463        } else {
464            BTreeMap::new()
465        };
466
467        // check aliases
468        let mut aliases_path = path.clone();
469        aliases_path.set_extension("aliases");
470
471        let aliases = if aliases_path.exists() {
472            let reader = BufReader::new(File::open(&aliases_path).with_context(|| {
473                format!(
474                    "Cannot open aliases file in keystore: {}",
475                    aliases_path.display()
476                )
477            })?);
478
479            let aliases: Vec<Alias> = serde_json::from_reader(reader).with_context(|| {
480                format!(
481                    "Cannot deserialize aliases file in keystore: {}",
482                    aliases_path.display(),
483                )
484            })?;
485
486            aliases
487                .into_iter()
488                .map(|alias| {
489                    let key = PublicKey::decode_base64(&alias.public_key_base64);
490                    key.map(|k| (Into::<IotaAddress>::into(&k), alias))
491                })
492                .collect::<Result<BTreeMap<_, _>, _>>()
493                .map_err(|e| {
494                    anyhow!(
495                        "Invalid aliases file in keystore: {}. {}",
496                        aliases_path.display(),
497                        e
498                    )
499                })?
500        } else if keys.is_empty() {
501            BTreeMap::new()
502        } else {
503            let names: Vec<String> = random_names(HashSet::new(), keys.len());
504            let aliases = keys
505                .iter()
506                .zip(names)
507                .map(|((iota_address, ikp), alias)| {
508                    let public_key_base64 = ikp.public().encode_base64();
509                    (
510                        *iota_address,
511                        Alias {
512                            alias,
513                            public_key_base64,
514                        },
515                    )
516                })
517                .collect::<BTreeMap<_, _>>();
518            let aliases_store = serde_json::to_string_pretty(&aliases.values().collect::<Vec<_>>())
519                .with_context(|| {
520                    format!(
521                        "Cannot serialize aliases to file in keystore: {}",
522                        aliases_path.display()
523                    )
524                })?;
525            fs::write(aliases_path, aliases_store)?;
526            aliases
527        };
528
529        Ok(Self {
530            keys,
531            aliases,
532            path: path.to_path_buf(),
533        })
534    }
535
536    fn needs_migration(path: &PathBuf) -> Result<bool, anyhow::Error> {
537        let mut aliases_path = path.clone();
538        aliases_path.set_extension("aliases");
539        if aliases_path.exists() {
540            // If the aliases file exists, we assume that the keystore is in v1 format
541            debug!(
542                "An alias file exists at {}, assuming keystore is in v1 format",
543                aliases_path.display()
544            );
545            return Ok(true);
546        }
547        // If the aliases file does not exist, we check if the keystore file exists and
548        // it has the old format
549        if path.exists() {
550            let reader =
551                BufReader::new(File::open(path).with_context(|| {
552                    format!("Cannot open the keystore file: {}", path.display())
553                })?);
554            // If we can deserialize the keystore file as a Vec<String>, it is in v1 format
555            // If it fails, it is in v2 format or invalid
556            let is_v1_format = serde_json::from_reader::<_, Vec<String>>(reader).is_ok();
557            debug!(
558                "Keystore file at {} is in v1 format: {}",
559                path.display(),
560                is_v1_format,
561            );
562            Ok(is_v1_format)
563        } else {
564            // If the keystore file does not exist, no migration is needed
565            Ok(false)
566        }
567    }
568
569    fn migrate_v1_to_v2(path: &PathBuf) -> Result<Self, anyhow::Error> {
570        let migrated = Self::new_from_v1(path)?;
571        // If the migration was successful, we rename the <path> and <path>.aliases
572        // files as a backup and append .migrated
573        let mut backup_path = path.clone();
574        backup_path.set_extension(
575            // We append .migrated to the original file extension
576            // add_extension is still experimental, so we do it manually
577            path.extension()
578                .and_then(|ext| Some(ext.to_str()?.to_owned() + ".migrated"))
579                .unwrap_or(String::from("migrated")),
580        );
581        fs::rename(path, &backup_path).with_context(|| {
582            format!(
583                "Failed to rename the old keystore file to {}",
584                backup_path.display()
585            )
586        })?;
587        let mut aliases_path = path.clone();
588        aliases_path.set_extension("aliases");
589        let mut backup_aliases_path = aliases_path.clone();
590        backup_aliases_path.set_extension("aliases.migrated");
591        fs::rename(&aliases_path, &backup_aliases_path).with_context(|| {
592            format!(
593                "Failed to rename the old aliases file to {}",
594                backup_aliases_path.display()
595            )
596        })?;
597
598        info!(
599            "Migrated {} keys in keystore from v1 to v2 format. Old files have been renamed to {} and {}",
600            migrated.keys.len(),
601            backup_path.display(),
602            backup_aliases_path.display()
603        );
604
605        // Now we save the migrated keystore to the original path
606        migrated.save()?;
607
608        Ok(migrated)
609    }
610
611    pub fn new(path: &PathBuf) -> Result<Self, anyhow::Error> {
612        if Self::needs_migration(path)? {
613            return Self::migrate_v1_to_v2(path);
614        }
615
616        let (keys, aliases) = if path.exists() {
617            let reader =
618                BufReader::new(File::open(path).with_context(|| {
619                    format!("Cannot open the keystore file: {}", path.display())
620                })?);
621
622            let file: FileBasedKeystoreFile = serde_json::from_reader(reader)
623                .with_context(|| {
624                    format!("Cannot deserialize the keystore file: {}", path.display(),)
625                })
626                .map_err(|e| anyhow!("Invalid keystore file: {}. {}", path.display(), e))?;
627
628            let aliases = file
629                .keys
630                .iter()
631                .map(|aliased| {
632                    let public_key = aliased.key.public();
633                    (
634                        IotaAddress::from(&public_key),
635                        Alias {
636                            alias: aliased.alias.clone(),
637                            public_key_base64: public_key.encode_base64(),
638                        },
639                    )
640                })
641                .collect::<BTreeMap<_, _>>();
642
643            let keys = file
644                .keys
645                .into_iter()
646                .map(|aliased| (IotaAddress::from(&aliased.key.public()), aliased.key))
647                .collect::<BTreeMap<_, _>>();
648
649            (keys, aliases)
650        } else {
651            (BTreeMap::new(), BTreeMap::new())
652        };
653
654        Ok(Self {
655            keys,
656            aliases,
657            path: path.to_path_buf(),
658        })
659    }
660
661    pub fn set_path(&mut self, path: &Path) {
662        self.path = path.to_path_buf();
663    }
664
665    pub fn save(&self) -> Result<(), anyhow::Error> {
666        let file = FileBasedKeystoreFile {
667            version: 2,
668            keys: self
669                .keys
670                .iter()
671                .map(|(address, key)| AliasedKey {
672                    alias: self
673                        .aliases
674                        .get(address)
675                        .map_or_else(|| self.create_alias(None).unwrap(), |a| a.alias.clone()),
676                    address: *address,
677                    key: key.clone(),
678                })
679                .collect(),
680        };
681
682        let store = serde_json::to_string_pretty(&file).with_context(|| {
683            format!("Cannot serialize keystore to file: {}", self.path.display())
684        })?;
685        fs::write(&self.path, store)
686            .map_err(|e| anyhow!("Couldn't save keystore to {}: {e}", self.path.display()))?;
687        Ok(())
688    }
689}
690
691#[derive(Default, Serialize, Deserialize)]
692pub struct InMemKeystore {
693    aliases: BTreeMap<IotaAddress, Alias>,
694    keys: BTreeMap<IotaAddress, StoredKey>,
695}
696
697impl AccountKeystore for InMemKeystore {
698    fn sign_hashed(
699        &self,
700        address: &IotaAddress,
701        msg: &[u8],
702    ) -> Result<Signature, signature::Error> {
703        let stored_key = self.keys.get(address).ok_or_else(|| {
704            signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
705        })?;
706
707        match stored_key {
708            StoredKey::KeyPair(keypair) => Ok(Signature::new_hashed(msg, keypair)),
709            StoredKey::External { source, .. } => Err(signature::Error::from_source(format!(
710                "sign_hashed is not supported for external type: {source} [{address}]"
711            ))),
712        }
713    }
714    fn sign_secure<T>(
715        &self,
716        address: &IotaAddress,
717        msg: &T,
718        intent: Intent,
719    ) -> Result<Signature, signature::Error>
720    where
721        T: Serialize,
722    {
723        let stored_key = self.keys.get(address).ok_or_else(|| {
724            signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
725        })?;
726
727        let intent_msg = &IntentMessage::new(intent, msg);
728        match stored_key {
729            StoredKey::KeyPair(keypair) => Ok(Signature::new_secure(intent_msg, keypair)),
730            StoredKey::External { source, .. } => Err(signature::Error::from_source(format!(
731                "sign_secure is not supported for external type: {source} [{address}]",
732            ))),
733        }
734    }
735
736    fn add_key(
737        &mut self,
738        alias: Option<String>,
739        key: impl Into<StoredKey>,
740    ) -> Result<(), anyhow::Error> {
741        let key = key.into();
742        let address: IotaAddress = (&key.public()).into();
743        let alias = alias.unwrap_or_else(|| {
744            random_name(
745                &self
746                    .aliases()
747                    .iter()
748                    .map(|x| x.alias.clone())
749                    .collect::<HashSet<_>>(),
750            )
751        });
752
753        let public_key_base64 = key.public().encode_base64();
754        let alias = Alias {
755            alias,
756            public_key_base64,
757        };
758        self.aliases.insert(address, alias);
759        self.keys.insert(address, key);
760        Ok(())
761    }
762
763    fn remove_key(&mut self, address: &IotaAddress) -> Result<(), anyhow::Error> {
764        self.aliases.remove(address);
765        self.keys.remove(address);
766        Ok(())
767    }
768
769    /// Get all aliases objects
770    fn aliases(&self) -> Vec<&Alias> {
771        self.aliases.values().collect()
772    }
773
774    fn addresses_with_alias(&self) -> Vec<(&IotaAddress, &Alias)> {
775        self.aliases.iter().collect::<Vec<_>>()
776    }
777
778    fn keys(&self) -> Vec<&StoredKey> {
779        self.keys.values().collect()
780    }
781
782    fn get_key(&self, address: &IotaAddress) -> Result<&StoredKey, anyhow::Error> {
783        match self.keys.get(address) {
784            Some(key) => Ok(key),
785            None => Err(anyhow!("Cannot find key for address: [{address}]")),
786        }
787    }
788
789    /// Get alias of address
790    fn get_alias_by_address(&self, address: &IotaAddress) -> Result<String, anyhow::Error> {
791        match self.aliases.get(address) {
792            Some(alias) => Ok(alias.alias.clone()),
793            None => bail!("Cannot find alias for address {address}"),
794        }
795    }
796
797    /// Get the address by its alias
798    fn get_address_by_alias(&self, alias: String) -> Result<&IotaAddress, anyhow::Error> {
799        self.addresses_with_alias()
800            .iter()
801            .find(|x| x.1.alias == alias)
802            .ok_or_else(|| anyhow!("Cannot resolve alias {alias} to an address"))
803            .map(|x| x.0)
804    }
805
806    /// This function returns an error if the provided alias already exists. If
807    /// the alias has not already been used, then it returns the alias.
808    /// If no alias has been passed, it will generate a new alias.
809    fn create_alias(&self, alias: Option<String>) -> Result<String, anyhow::Error> {
810        match alias {
811            Some(a) if self.alias_exists(&a) => {
812                bail!("Alias {a} already exists. Please choose another alias.")
813            }
814            Some(a) => validate_alias(&a),
815            None => Ok(random_name(
816                &self
817                    .alias_names()
818                    .into_iter()
819                    .map(|x| x.to_string())
820                    .collect::<HashSet<_>>(),
821            )),
822        }
823    }
824
825    fn aliases_mut(&mut self) -> Vec<&mut Alias> {
826        self.aliases.values_mut().collect()
827    }
828
829    /// Updates an old alias to the new alias. If the new_alias is None,
830    /// it will generate a new random alias.
831    fn update_alias(
832        &mut self,
833        old_alias: &str,
834        new_alias: Option<&str>,
835    ) -> Result<String, anyhow::Error> {
836        self.update_alias_value(old_alias, new_alias)
837    }
838}
839
840impl InMemKeystore {
841    pub fn new_insecure_for_tests(initial_key_number: usize) -> Self {
842        let mut rng = StdRng::from_seed([0; 32]);
843        let keys = (0..initial_key_number)
844            .map(|_| get_key_pair_from_rng(&mut rng))
845            .map(|(ad, k)| (ad, IotaKeyPair::Ed25519(k).into()))
846            .collect::<BTreeMap<IotaAddress, StoredKey>>();
847
848        let aliases = keys
849            .iter()
850            .zip(random_names(HashSet::new(), keys.len()))
851            .map(|((iota_address, ikp), alias)| {
852                let public_key_base64 = ikp.public().encode_base64();
853                (
854                    *iota_address,
855                    Alias {
856                        alias,
857                        public_key_base64,
858                    },
859                )
860            })
861            .collect::<BTreeMap<_, _>>();
862
863        Self { aliases, keys }
864    }
865}
866
867fn validate_alias(alias: &str) -> Result<String, anyhow::Error> {
868    let re = Regex::new(r"^[A-Za-z][A-Za-z0-9-_\.]*$")
869        .map_err(|_| anyhow!("Cannot build the regex needed to validate the alias naming"))?;
870    let alias = alias.trim();
871    ensure!(
872        re.is_match(alias),
873        "Invalid alias. A valid alias must start with a letter and can contain only letters, digits, hyphens (-), dots (.), or underscores (_)."
874    );
875    Ok(alias.to_string())
876}
877
878#[cfg(test)]
879mod tests {
880    use crate::keystore::validate_alias;
881
882    #[test]
883    fn validate_alias_test() {
884        // OK
885        assert!(validate_alias("A.B_dash").is_ok());
886        assert!(validate_alias("A.B-C1_dash").is_ok());
887        assert!(validate_alias("abc_123.iota").is_ok());
888        // Not allowed
889        assert!(validate_alias("A.B-C_dash!").is_err());
890        assert!(validate_alias(".B-C_dash!").is_err());
891        assert!(validate_alias("_test").is_err());
892        assert!(validate_alias("123").is_err());
893        assert!(validate_alias("@@123").is_err());
894        assert!(validate_alias("@_Ab").is_err());
895        assert!(validate_alias("_Ab").is_err());
896        assert!(validate_alias("^A").is_err());
897        assert!(validate_alias("-A").is_err());
898    }
899}