af_keys/
keystore.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::BTreeMap;
5use std::fmt::{Display, Formatter, Write};
6use std::fs::File;
7use std::io::BufReader;
8use std::path::PathBuf;
9
10use af_sui_types::Address as SuiAddress;
11use enum_dispatch::enum_dispatch;
12use eyre::{Context as _, bail, eyre};
13use fastcrypto::traits::EncodeDecodeBase64;
14use serde::{Deserialize, Deserializer, Serialize, Serializer};
15
16use crate::crypto::{PublicKey, Signature, SuiKeyPair};
17use crate::intent::{Intent, IntentMessage};
18
19pub type Error = eyre::Report;
20
21#[derive(Serialize, Deserialize)]
22#[enum_dispatch(ReadOnlyAccountKeystore)]
23pub enum Keystore {
24    File(FileBasedKeystore),
25    InMem(InMemKeystore),
26}
27
28impl Display for Keystore {
29    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
30        let mut writer = String::new();
31        match self {
32            Self::File(file) => {
33                writeln!(writer, "Keystore Type : File")?;
34                write!(writer, "Keystore Path : {:?}", file.path)?;
35                write!(f, "{}", writer)
36            }
37            Self::InMem(_) => {
38                writeln!(writer, "Keystore Type : InMem")?;
39                write!(f, "{}", writer)
40            }
41        }
42    }
43}
44
45/// Read-only version of [AccountKeystore].
46///
47/// Allows light-weight applications to use the same keystore file as the Sui CLI.
48///
49/// [AccountKeystore]: https://mystenlabs.github.io/sui/sui_keys/keystore/trait.AccountKeystore.html
50#[enum_dispatch]
51pub trait ReadOnlyAccountKeystore: Send + Sync {
52    fn keys(&self) -> Vec<PublicKey>;
53
54    fn get_key(&self, address: &SuiAddress) -> Result<&SuiKeyPair, Error>;
55
56    fn sign_hashed(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, signature::Error>;
57
58    fn sign_secure<T>(
59        &self,
60        address: &SuiAddress,
61        msg: &T,
62        intent: Intent,
63    ) -> Result<Signature, signature::Error>
64    where
65        T: Serialize;
66
67    fn addresses(&self) -> Vec<SuiAddress> {
68        self.keys().iter().map(|k| k.to_sui_address()).collect()
69    }
70
71    fn addresses_with_alias(&self) -> Vec<(&SuiAddress, &Alias)>;
72
73    fn aliases(&self) -> Vec<&Alias>;
74
75    fn alias_names(&self) -> Vec<&str> {
76        self.aliases()
77            .into_iter()
78            .map(|a| a.alias.as_str())
79            .collect()
80    }
81
82    /// Get alias of address
83    fn get_alias_by_address(&self, address: &SuiAddress) -> Result<String, Error>;
84
85    fn get_address_by_alias(&self, alias: String) -> Result<&SuiAddress, Error>;
86
87    /// Check if an alias exists by its name
88    fn alias_exists(&self, alias: &str) -> bool {
89        self.alias_names().contains(&alias)
90    }
91}
92
93#[derive(Default)]
94pub struct FileBasedKeystore {
95    keys: BTreeMap<SuiAddress, SuiKeyPair>,
96    aliases: BTreeMap<SuiAddress, Alias>,
97    path: Option<PathBuf>,
98}
99
100impl FileBasedKeystore {
101    pub fn new(path: PathBuf) -> Result<Self, Error> {
102        let keys = if path.exists() {
103            let reader =
104                BufReader::new(File::open(&path).with_context(|| {
105                    format!("Cannot open the keystore file: {}", path.display())
106                })?);
107            let kp_strings: Vec<String> = serde_json::from_reader(reader).with_context(|| {
108                format!("Cannot deserialize the keystore file: {}", path.display(),)
109            })?;
110            kp_strings
111                .iter()
112                .map(|kpstr| {
113                    let key = SuiKeyPair::decode_base64(kpstr);
114                    key.map(|k| (k.public().to_sui_address(), k))
115                })
116                .collect::<Result<BTreeMap<_, _>, _>>()
117                .map_err(|e| eyre!("Invalid keystore file: {}. {}", path.display(), e))?
118        } else {
119            BTreeMap::new()
120        };
121
122        // check aliases
123        let mut aliases_path = path.clone();
124        aliases_path.set_extension("aliases");
125
126        let aliases = if aliases_path.exists() {
127            let reader = BufReader::new(File::open(&aliases_path).with_context(|| {
128                format!(
129                    "Cannot open aliases file in keystore: {}",
130                    aliases_path.display()
131                )
132            })?);
133
134            let aliases: Vec<Alias> = serde_json::from_reader(reader).with_context(|| {
135                format!(
136                    "Cannot deserialize aliases file in keystore: {}",
137                    aliases_path.display(),
138                )
139            })?;
140
141            aliases
142                .into_iter()
143                .map(|alias| {
144                    let key = PublicKey::decode_base64(&alias.public_key_base64);
145                    key.map(|k| (k.to_sui_address(), alias))
146                })
147                .collect::<Result<BTreeMap<_, _>, _>>()
148                .map_err(|e| {
149                    eyre!(
150                        "Invalid aliases file in keystore: {}. {}",
151                        aliases_path.display(),
152                        e
153                    )
154                })?
155        } else {
156            BTreeMap::new()
157        };
158
159        Ok(Self {
160            keys,
161            aliases,
162            path: Some(path),
163        })
164    }
165
166    pub fn key_pairs(&self) -> Vec<&SuiKeyPair> {
167        self.keys.values().collect()
168    }
169}
170
171impl Serialize for FileBasedKeystore {
172    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
173    where
174        S: Serializer,
175    {
176        let path_default = PathBuf::default();
177        serializer.serialize_str(
178            self.path
179                .as_ref()
180                .unwrap_or(&path_default)
181                .to_str()
182                .unwrap_or(""),
183        )
184    }
185}
186
187impl<'de> Deserialize<'de> for FileBasedKeystore {
188    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
189    where
190        D: Deserializer<'de>,
191    {
192        use serde::de::Error;
193        Self::new(PathBuf::from(String::deserialize(deserializer)?)).map_err(D::Error::custom)
194    }
195}
196
197impl ReadOnlyAccountKeystore for FileBasedKeystore {
198    fn sign_hashed(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, signature::Error> {
199        Ok(Signature::new_hashed(
200            msg,
201            self.keys.get(address).ok_or_else(|| {
202                signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
203            })?,
204        ))
205    }
206    fn sign_secure<T>(
207        &self,
208        address: &SuiAddress,
209        msg: &T,
210        intent: Intent,
211    ) -> Result<Signature, signature::Error>
212    where
213        T: Serialize,
214    {
215        Ok(Signature::new_secure(
216            &IntentMessage::new(intent, msg),
217            self.keys.get(address).ok_or_else(|| {
218                signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
219            })?,
220        ))
221    }
222
223    /// Return an array of `Alias`, consisting of every alias and its corresponding public key.
224    fn aliases(&self) -> Vec<&Alias> {
225        self.aliases.values().collect()
226    }
227
228    fn addresses_with_alias(&self) -> Vec<(&SuiAddress, &Alias)> {
229        self.aliases.iter().collect::<Vec<_>>()
230    }
231
232    fn keys(&self) -> Vec<PublicKey> {
233        self.keys.values().map(|key| key.public()).collect()
234    }
235
236    /// Get the address by its alias
237    fn get_address_by_alias(&self, alias: String) -> Result<&SuiAddress, Error> {
238        self.addresses_with_alias()
239            .iter()
240            .find(|x| x.1.alias == alias)
241            .ok_or_else(|| eyre!("Cannot resolve alias {alias} to an address"))
242            .map(|x| x.0)
243    }
244
245    /// Get the alias if it exists, or return an error if it does not exist.
246    fn get_alias_by_address(&self, address: &SuiAddress) -> Result<String, Error> {
247        match self.aliases.get(address) {
248            Some(alias) => Ok(alias.alias.clone()),
249            None => bail!("Cannot find alias for address {address}"),
250        }
251    }
252
253    fn get_key(&self, address: &SuiAddress) -> Result<&SuiKeyPair, Error> {
254        #[allow(clippy::option_if_let_else)]
255        match self.keys.get(address) {
256            Some(key) => Ok(key),
257            None => Err(eyre!("Cannot find key for address: [{address}]")),
258        }
259    }
260}
261
262/// Carry-over from the original code, but un-initializable.
263#[derive(Default, Serialize, Deserialize)]
264pub struct InMemKeystore {
265    aliases: BTreeMap<SuiAddress, Alias>,
266    keys: BTreeMap<SuiAddress, SuiKeyPair>,
267}
268
269impl ReadOnlyAccountKeystore for InMemKeystore {
270    fn sign_hashed(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, signature::Error> {
271        Ok(Signature::new_hashed(
272            msg,
273            self.keys.get(address).ok_or_else(|| {
274                signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
275            })?,
276        ))
277    }
278    fn sign_secure<T>(
279        &self,
280        address: &SuiAddress,
281        msg: &T,
282        intent: Intent,
283    ) -> Result<Signature, signature::Error>
284    where
285        T: Serialize,
286    {
287        Ok(Signature::new_secure(
288            &IntentMessage::new(intent, msg),
289            self.keys.get(address).ok_or_else(|| {
290                signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
291            })?,
292        ))
293    }
294
295    /// Get all aliases objects
296    fn aliases(&self) -> Vec<&Alias> {
297        self.aliases.values().collect()
298    }
299
300    fn addresses_with_alias(&self) -> Vec<(&SuiAddress, &Alias)> {
301        self.aliases.iter().collect::<Vec<_>>()
302    }
303
304    fn keys(&self) -> Vec<PublicKey> {
305        self.keys.values().map(|key| key.public()).collect()
306    }
307
308    fn get_key(&self, address: &SuiAddress) -> Result<&SuiKeyPair, Error> {
309        #[allow(clippy::option_if_let_else)]
310        match self.keys.get(address) {
311            Some(key) => Ok(key),
312            None => Err(eyre!("Cannot find key for address: [{address}]")),
313        }
314    }
315
316    /// Get alias of address
317    fn get_alias_by_address(&self, address: &SuiAddress) -> Result<String, Error> {
318        match self.aliases.get(address) {
319            Some(alias) => Ok(alias.alias.clone()),
320            None => bail!("Cannot find alias for address {address}"),
321        }
322    }
323
324    /// Get the address by its alias
325    fn get_address_by_alias(&self, alias: String) -> Result<&SuiAddress, Error> {
326        self.addresses_with_alias()
327            .iter()
328            .find(|x| x.1.alias == alias)
329            .ok_or_else(|| eyre!("Cannot resolve alias {alias} to an address"))
330            .map(|x| x.0)
331    }
332}
333
334#[derive(Serialize, Deserialize, Clone, Debug)]
335pub struct Alias {
336    pub alias: String,
337    pub public_key_base64: String,
338}
339
340#[cfg(test)]
341mod tests {
342    use std::io::Write;
343
344    use super::*;
345
346    #[test]
347    fn new_file_keystore() -> eyre::Result<()> {
348        let temp_dir = tempfile::tempdir()?;
349        let path = temp_dir.path();
350        let mut keystore = File::create(path.join("sui.keystore"))?;
351        serde_json::to_writer(
352            &keystore,
353            &serde_json::json!([
354                "AKd4u480uT0eLUNe7vh2zHHYdbpUXY/fwcL13eJQ5/zs",
355                "AI1TKQ0qPLor32rdLOZiN0/J4qNPyypesT1eE+R/wSCB",
356                "AFHMjegm2IwuiLemXb6o7XvuDL7xn1JTHc66CZefYY+B",
357                "APhbsR3gpjBIRvZm5ZwMZhncejgYH/hGa6wHVtaTat22",
358                "ADO8QyYe0MM+HP0iLjHNLPAxZXNYyE1jieny3iN+fDCS",
359                "AKfLSiyx3pUSEpvn0tyY+17ef8AjN7izfQ9qm048BhqM",
360                "AOzplQlAK2Uznvog7xmcMtlFC+DfuJx3axo9lfyI876G",
361                "AI1I9i3mk2e1kAjPnB7fKiqquxc1OjjAkkpQPIk9Id5Q",
362                "AIUAgL5jYMzf0JPCmc263Ou6tH5Z/HuAdtWFFUiz8Zc0",
363                "AFmgBTlVGHfYieuSVmQ63BJ+zQSY8pNOUXH99Ucb1ZGl",
364                "AAu4ySMvq2wygxl/Ze6AGgkYfxg+rzUElj7UxxI6NHBI"
365            ]),
366        )?;
367        keystore.flush()?;
368        let mut aliases = File::create(path.join("sui.aliases"))?;
369        serde_json::to_writer(
370            &aliases,
371            &serde_json::json!([
372              {
373                "alias": "grace",
374                "public_key_base64": "ABhIIE33kaUT1rr9rNrh0XJNb7AC6EBSdh5Ku4a2B7wU"
375              },
376              {
377                "alias": "heidi",
378                "public_key_base64": "ACRAZZ+qMcBA7gJg6iacBSgB4S+DB3nHjk9E1237R4+h"
379              },
380              {
381                "alias": "admin",
382                "public_key_base64": "AONa32KBWXqsu6pksuwCLbA0v3JoSPbw8du45Rkw14nm"
383              },
384              {
385                "alias": "ivan",
386                "public_key_base64": "AKsTkJa8fJg2PJtUTUxIE+FHBBG6IFkHk4385yehR86L"
387              },
388              {
389                "alias": "judy",
390                "public_key_base64": "AEIcS8FhN0CjRUGjVHNmXOW6Rb+ootVN3a4kEbBoQ4R6"
391              },
392              {
393                "alias": "eve",
394                "public_key_base64": "AP0TE5MM1h7QSZrnlBcdQepKA/6Fh5pja3gjMNpL1fix"
395              },
396              {
397                "alias": "alice",
398                "public_key_base64": "AK9WofTFdyBcMpMxzYkbgNQiKLgr9qH8iz9ON6VFxwiW"
399              },
400              {
401                "alias": "bob",
402                "public_key_base64": "ALieneYHseSZILiNAda3z29Ob4lZKBAr3jEyP41WsJAG"
403              },
404              {
405                "alias": "charlie",
406                "public_key_base64": "ABm2kTdq/96JsbsTMunKZDqJbIsEa1lwIJ0cA2CJ4z5l"
407              },
408              {
409                "alias": "frank",
410                "public_key_base64": "ADSxYutFskDwLNnEto/E+KDJe4QXWHkO7d8Ha6nqBR0/"
411              },
412              {
413                "alias": "dave",
414                "public_key_base64": "ALmzETq2T6c06a+VXJzx1pkfuLBVetRs5q537l6UO4KI"
415              }
416            ]),
417        )?;
418        aliases.flush()?;
419        let keystore = FileBasedKeystore::new(path.join("sui.keystore"))?;
420        assert!(!keystore.key_pairs().is_empty());
421        assert!(
422            keystore
423                .get_key(
424                    &"0x98e9cafb116af9d69f77ce0d644c60e384f850f8af050b268377d8293d7fe7c6"
425                        .parse()?
426                )
427                .is_ok()
428        );
429        assert!(keystore.get_address_by_alias("alice".to_owned()).is_ok());
430        Ok(())
431    }
432
433    #[test]
434    fn new_file_keystore_no_aliases() -> eyre::Result<()> {
435        let temp_dir = tempfile::tempdir()?;
436        let path = temp_dir.path();
437        let keystore_path = path.join("sui.keystore");
438        serde_json::to_writer(
439            File::create(keystore_path.clone())?,
440            &serde_json::json!([
441                "AKd4u480uT0eLUNe7vh2zHHYdbpUXY/fwcL13eJQ5/zs",
442                "AI1TKQ0qPLor32rdLOZiN0/J4qNPyypesT1eE+R/wSCB",
443                "AFHMjegm2IwuiLemXb6o7XvuDL7xn1JTHc66CZefYY+B",
444                "APhbsR3gpjBIRvZm5ZwMZhncejgYH/hGa6wHVtaTat22",
445                "ADO8QyYe0MM+HP0iLjHNLPAxZXNYyE1jieny3iN+fDCS",
446                "AKfLSiyx3pUSEpvn0tyY+17ef8AjN7izfQ9qm048BhqM",
447                "AOzplQlAK2Uznvog7xmcMtlFC+DfuJx3axo9lfyI876G",
448                "AI1I9i3mk2e1kAjPnB7fKiqquxc1OjjAkkpQPIk9Id5Q",
449                "AIUAgL5jYMzf0JPCmc263Ou6tH5Z/HuAdtWFFUiz8Zc0",
450                "AFmgBTlVGHfYieuSVmQ63BJ+zQSY8pNOUXH99Ucb1ZGl",
451                "AAu4ySMvq2wygxl/Ze6AGgkYfxg+rzUElj7UxxI6NHBI"
452            ]),
453        )?;
454        let keystore = FileBasedKeystore::new(keystore_path)?;
455        assert!(!keystore.key_pairs().is_empty());
456        assert!(
457            keystore
458                .get_key(
459                    &"0x98e9cafb116af9d69f77ce0d644c60e384f850f8af050b268377d8293d7fe7c6"
460                        .parse()?
461                )
462                .is_ok()
463        );
464        assert!(keystore.get_address_by_alias("alice".to_owned()).is_err());
465        Ok(())
466    }
467}
af_keys/keystore.rs

af_keys/
keystore.rs
