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