Skip to main content

af_keys/
keystore.rs

1use std::collections::BTreeMap;
2use std::fs::File;
3use std::io::BufReader;
4use std::path::{Path, PathBuf};
5
6use anyhow::{Context, Error, bail};
7use serde::{Deserialize, Serialize};
8use sui_crypto::Signer as _;
9use sui_crypto::ed25519::Ed25519PrivateKey;
10use sui_crypto::secp256k1::Secp256k1PrivateKey;
11use sui_crypto::secp256r1::Secp256r1PrivateKey;
12use sui_crypto::simple::SimpleKeypair;
13use sui_sdk_types::bcs::FromBcs;
14use sui_sdk_types::{
15    Address, MultisigAggregatedSignature, MultisigCommittee, MultisigMemberPublicKey,
16    MultisigMemberSignature, SignatureScheme, SimpleSignature, Transaction,
17};
18
19use crate::PublicKey;
20
21#[derive(Clone, Debug, Deserialize, Serialize)]
22pub struct Alias {
23    pub alias: String,
24    pub public_key_base64: String,
25}
26
27#[derive(Debug)]
28pub struct Keystore {
29    path: PathBuf,
30    keys: BTreeMap<Address, SimpleKeypair>,
31    aliases: BTreeMap<Address, Alias>,
32}
33
34impl Keystore {
35    /// Loads a keystore from the default path: `$HOME/.sui/sui_config/sui.keystore`.
36    pub fn new_default() -> Result<Self, Error> {
37        let keystore_path = match std::env::home_dir() {
38            Some(v) => v.join(".sui/sui_config/sui.keystore"),
39            None => bail!("cannot obtain home directory path"),
40        };
41        Self::new(keystore_path)
42    }
43    pub fn new(path: PathBuf) -> Result<Self, Error> {
44        let keys = if path.exists() {
45            let path_display = path.display();
46            let f = File::open(&path)
47                .with_context(|| format!("unable to open the keystore file \"{path_display}\""))?;
48            let reader = BufReader::new(f);
49            let kp_strings: Vec<String> = serde_json::from_reader(reader).with_context(|| {
50                format!("unable to deserialize the keystore file \"{path_display}\"")
51            })?;
52            kp_strings
53                .iter()
54                .map(|kpstr| {
55                    let key = keypair_from_base64(kpstr)?;
56                    let address = PublicKey::from(key.public_key()).address()?;
57                    Ok((address, key))
58                })
59                .collect::<Result<BTreeMap<_, _>, Error>>()
60                .with_context(|| format!("invalid keystore file \"{path_display}\""))?
61        } else {
62            BTreeMap::new()
63        };
64
65        let mut aliases_path = path.clone();
66        aliases_path.set_extension("aliases");
67        let aliases = if aliases_path.exists() {
68            let path_display = aliases_path.display();
69            let reader = BufReader::new(
70                File::open(&aliases_path)
71                    .with_context(|| format!("unable to open aliases file \"{path_display}\""))?,
72            );
73
74            let aliases: Vec<Alias> = serde_json::from_reader(reader).with_context(|| {
75                format!("unable to deserialize aliases file \"{path_display}\"")
76            })?;
77
78            aliases
79                .into_iter()
80                .map(|alias| {
81                    let key = PublicKey::from_base64(&alias.public_key_base64)?;
82                    let address = key.address()?;
83                    Ok((address, alias))
84                })
85                .collect::<Result<BTreeMap<_, _>, Error>>()
86                .with_context(|| format!("invalid aliases file \"{path_display}\""))?
87        } else {
88            BTreeMap::new()
89        };
90
91        Ok(Self {
92            path,
93            keys,
94            aliases,
95        })
96    }
97    pub fn path(&self) -> &Path {
98        &self.path
99    }
100    pub const fn aliases(&self) -> &BTreeMap<Address, Alias> {
101        &self.aliases
102    }
103    pub fn get_public_key(&self, address: Address) -> Option<PublicKey> {
104        self.keys
105            .get(&address)
106            .map(|keypair| keypair.public_key().into())
107    }
108    pub fn sign_message(&self, message: &[u8], signer: Address) -> Result<SimpleSignature, Error> {
109        let Some(key_pair) = self.keys.get(&signer) else {
110            bail!("unable to find keypair for signer address {signer}")
111        };
112        Ok(key_pair.try_sign(message)?)
113    }
114    /// Sign the `transaction` for a simple address. Fails if the keystore lacks the private
115    /// key for it.
116    pub fn sign_tx(
117        &self,
118        transaction: &Transaction,
119        signer: Address,
120    ) -> Result<SimpleSignature, Error> {
121        let message = transaction.signing_digest();
122        self.sign_message(&message, signer)
123    }
124    /// Sign the `transaction` for a native Sui multisig address. Fails if the keystore lacks the
125    /// private keys for the signers with the given `indices`.
126    pub fn multisign_tx(
127        &self,
128        transaction: &Transaction,
129        committee: MultisigCommittee,
130        indices: &[usize],
131    ) -> Result<MultisigAggregatedSignature, Error> {
132        let message = transaction.signing_digest();
133        let mut total_weight = 0;
134        let mut signatures = vec![];
135        let mut bitmap = 0;
136
137        for index in indices.iter().copied() {
138            let Some(member) = committee.members().get(index) else {
139                bail!("signer index {index} out of bounds for multisig {committee:?}");
140            };
141            total_weight += member.weight() as u16;
142            let address = match member.public_key() {
143                MultisigMemberPublicKey::Ed25519(public_key) => public_key.derive_address(),
144                MultisigMemberPublicKey::Secp256k1(public_key) => public_key.derive_address(),
145                MultisigMemberPublicKey::Secp256r1(public_key) => public_key.derive_address(),
146                _ => bail!("unsupported public key scheme for multisig member {member:?}"),
147            };
148            signatures.push(self.sign_message(&message, address)?);
149            bitmap |= 1 << index;
150        }
151
152        if total_weight < committee.threshold() {
153            bail!("signers do not have enough weight to sign for multisig");
154        }
155
156        let signatures = signatures
157            .into_iter()
158            .map(member_signature_from_simple)
159            .collect::<Result<Vec<_>, Error>>()?;
160
161        Ok(MultisigAggregatedSignature::new(
162            committee, signatures, bitmap,
163        ))
164    }
165}
166
167pub fn member_signature_from_simple(
168    signature: SimpleSignature,
169) -> Result<MultisigMemberSignature, Error> {
170    match signature {
171        SimpleSignature::Ed25519 { signature, .. } => {
172            Ok(MultisigMemberSignature::Ed25519(signature))
173        }
174        SimpleSignature::Secp256k1 { signature, .. } => {
175            Ok(MultisigMemberSignature::Secp256k1(signature))
176        }
177        SimpleSignature::Secp256r1 { signature, .. } => {
178            Ok(MultisigMemberSignature::Secp256r1(signature))
179        }
180        _ => bail!(
181            "unsupported signature scheme for multisig: {}",
182            signature.scheme().name(),
183        ),
184    }
185}
186
187#[derive(Deserialize)]
188struct KeyBytes {
189    flag: u8,
190    key: [u8; 32],
191}
192
193pub fn keypair_from_base64(base64: &str) -> Result<SimpleKeypair, Error> {
194    let KeyBytes { flag, key } = KeyBytes::from_bcs_base64(base64)?;
195    let scheme = SignatureScheme::from_byte(flag).map_err(|err| anyhow::anyhow!(err))?;
196    Ok(match scheme {
197        SignatureScheme::Ed25519 => SimpleKeypair::from(Ed25519PrivateKey::new(key)),
198        SignatureScheme::Secp256k1 => SimpleKeypair::from(Secp256k1PrivateKey::new(key)?),
199        SignatureScheme::Secp256r1 => SimpleKeypair::from(Secp256r1PrivateKey::new(key)),
200        _ => {
201            bail!(
202                "unsupported signature scheme {} for a base64-encoded private key",
203                scheme.name(),
204            );
205        }
206    })
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    const TESTING_PRIVATE_KEYS: [&str; 13] = [
214        "AI1TKQ0qPLor32rdLOZiN0/J4qNPyypesT1eE+R/wSCB",
215        "AFHMjegm2IwuiLemXb6o7XvuDL7xn1JTHc66CZefYY+B",
216        "APhbsR3gpjBIRvZm5ZwMZhncejgYH/hGa6wHVtaTat22",
217        "ADO8QyYe0MM+HP0iLjHNLPAxZXNYyE1jieny3iN+fDCS",
218        "AKfLSiyx3pUSEpvn0tyY+17ef8AjN7izfQ9qm048BhqM",
219        "AOzplQlAK2Uznvog7xmcMtlFC+DfuJx3axo9lfyI876G",
220        "AI1I9i3mk2e1kAjPnB7fKiqquxc1OjjAkkpQPIk9Id5Q",
221        "AIUAgL5jYMzf0JPCmc263Ou6tH5Z/HuAdtWFFUiz8Zc0",
222        "AFmgBTlVGHfYieuSVmQ63BJ+zQSY8pNOUXH99Ucb1ZGl",
223        "AAu4ySMvq2wygxl/Ze6AGgkYfxg+rzUElj7UxxI6NHBI",
224        "Aoa82Y+xoAzdBLBehaon2kdDst6DNlSOhu+0E43iIfpL",
225        "AHAlBn/RWkr6ATvorp6pABpBxy2mRBUNV9RmcU5naeFr",
226        "ARn1JTV9CB6x++N/3+BucJFw58vE7p16i1Exd6MOhwnT",
227    ];
228
229    const TESTING_PUBLIC_KEYS: [&str; 13] = [
230        "ACRAZZ+qMcBA7gJg6iacBSgB4S+DB3nHjk9E1237R4+h",
231        "AONa32KBWXqsu6pksuwCLbA0v3JoSPbw8du45Rkw14nm",
232        "AKsTkJa8fJg2PJtUTUxIE+FHBBG6IFkHk4385yehR86L",
233        "AEIcS8FhN0CjRUGjVHNmXOW6Rb+ootVN3a4kEbBoQ4R6",
234        "AP0TE5MM1h7QSZrnlBcdQepKA/6Fh5pja3gjMNpL1fix",
235        "AK9WofTFdyBcMpMxzYkbgNQiKLgr9qH8iz9ON6VFxwiW",
236        "ALieneYHseSZILiNAda3z29Ob4lZKBAr3jEyP41WsJAG",
237        "ABm2kTdq/96JsbsTMunKZDqJbIsEa1lwIJ0cA2CJ4z5l",
238        "ADSxYutFskDwLNnEto/E+KDJe4QXWHkO7d8Ha6nqBR0/",
239        "ALmzETq2T6c06a+VXJzx1pkfuLBVetRs5q537l6UO4KI",
240        "AgJFm9OwmeDknCkUElQlg0e0fJmZg/McSUm6UJH37r61uQ==",
241        "ABAiMvjSzayOOYjqNhi2vSgc0qasEQbdJI8ponQ6scXI",
242        "AQIu5EC7mUXcgF3oVvqIuCzbp562mUtBqQ/sG+tUqo5KVQ==",
243    ];
244
245    const TESTING_ADDRESSES: [&str; 13] = [
246        "0x02ef3105413b0bd2ea2f1eee19df48ef4b873694e75b36eaa81c1d1e7d9cf13c",
247        "0x2d78d396d59080e2ee66d73cb09ce28b70708b0672c390bcb68cff529e298964",
248        "0x43bb1276973beb02c31854145b5c726715c27d8cd49c99534b504ef49951b5fa",
249        "0x8cecbc32959d9f610c19b96fe134aa53aa6ba608afaac4e081cd71b30de3459a",
250        "0x93d43128794ae9ace7aa8f456ab42281b322c48c0785589ef729bfc3fbd0cda5",
251        "0x98e9cafb116af9d69f77ce0d644c60e384f850f8af050b268377d8293d7fe7c6",
252        "0xb7d13dae9aec267ae30bbb0811247c032647bf07d4025e97a576dd5a055a713e",
253        "0xc4cc77d7de4418d1b84c04e1061f43b74ff2b1e39a85551a3a72fcfe5b8198b5",
254        "0xe3a9692f8423d893f87201445a07e24d7d29f997d7ecf8ae880bd635c9845ed4",
255        "0xe94c6ed879599794a241d748d714e130da5401489be5d44868377e8c66b620e2",
256        "0x0e1205897f909f80a4c6f199abf201cec0f1198ffe4d3c99944be0c7ecb2e2f0",
257        "0x3fa7feadd12495a52edc6228cdc1447a8930824dd0fc44eca5909263ec4aa211",
258        "0xe0d6507e453e43b6e71400083ef56b658c04de4ab49344d9ceb16f8275845231",
259    ];
260
261    #[test]
262    fn test_keypair_from_base64_and_address_from_public_key() -> Result<(), Error> {
263        let iter = std::iter::zip(
264            std::iter::zip(TESTING_PRIVATE_KEYS, TESTING_PUBLIC_KEYS),
265            TESTING_ADDRESSES,
266        );
267        for ((encoded_private, encoded_public), encoded_address) in iter {
268            let keypair = keypair_from_base64(encoded_private)?;
269            let public_key = PublicKey::from_base64(encoded_public)?;
270            assert_eq!(PublicKey::from(keypair.public_key()), public_key);
271            let address = Address::from_hex(encoded_address)?;
272            assert_eq!(public_key.address()?, address);
273        }
274        Ok(())
275    }
276}