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

iota_keys/
keystore.rs
