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

iota_keys/
keystore.rs
