sui_types/unit_tests/
utils.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::SuiAddress;
5use crate::crypto::{Signer, SuiKeyPair};
6use crate::multisig::{MultiSig, MultiSigPublicKey};
7use crate::programmable_transaction_builder::ProgrammableTransactionBuilder;
8use crate::transaction::{SenderSignedData, TEST_ONLY_GAS_UNIT_FOR_TRANSFER};
9use crate::{
10    base_types::{ObjectID, dbg_addr},
11    committee::Committee,
12    crypto::{
13        AccountKeyPair, AuthorityKeyPair, AuthorityPublicKeyBytes, DefaultHash, Signature,
14        SignatureScheme, get_key_pair, get_key_pair_from_rng,
15    },
16    object::Object,
17    signature::GenericSignature,
18    transaction::{Transaction, TransactionData},
19    zk_login_authenticator::ZkLoginAuthenticator,
20};
21use fastcrypto::ed25519::Ed25519KeyPair;
22use fastcrypto::hash::HashFunction;
23use fastcrypto::traits::KeyPair as KeypairTraits;
24use rand::SeedableRng;
25use rand::rngs::StdRng;
26use serde::Deserialize;
27use shared_crypto::intent::{Intent, IntentMessage};
28use std::collections::BTreeMap;
29
30#[derive(Deserialize)]
31pub struct TestData {
32    pub zklogin_inputs: String,
33    pub kp: String,
34    pub pk_bigint: String,
35    pub salt: String,
36    pub address_seed: String,
37}
38
39pub fn make_committee_key<R>(rand: &mut R) -> (Vec<AuthorityKeyPair>, Committee)
40where
41    R: rand::CryptoRng + rand::RngCore,
42{
43    make_committee_key_num(4, rand)
44}
45
46pub fn make_committee_key_num<R>(num: usize, rand: &mut R) -> (Vec<AuthorityKeyPair>, Committee)
47where
48    R: rand::CryptoRng + rand::RngCore,
49{
50    let mut authorities: BTreeMap<AuthorityPublicKeyBytes, u64> = BTreeMap::new();
51    let mut keys = Vec::new();
52
53    for _ in 0..num {
54        let (_, inner_authority_key): (_, AuthorityKeyPair) = get_key_pair_from_rng(rand);
55        authorities.insert(
56            /* address */ AuthorityPublicKeyBytes::from(inner_authority_key.public()),
57            /* voting right */ 1,
58        );
59        keys.push(inner_authority_key);
60    }
61
62    let committee = Committee::new_for_testing_with_normalized_voting_power(0, authorities);
63    (keys, committee)
64}
65
66// Creates a fake sender-signed transaction for testing. This transaction will
67// not actually work.
68pub fn create_fake_transaction() -> Transaction {
69    let (sender, sender_key): (_, AccountKeyPair) = get_key_pair();
70    let recipient = dbg_addr(2);
71    let object_id = ObjectID::random();
72    let object = Object::immutable_with_id_for_testing(object_id);
73    let pt = {
74        let mut builder = ProgrammableTransactionBuilder::new();
75        builder.transfer_sui(recipient, None);
76        builder.finish()
77    };
78    let data = TransactionData::new_programmable(
79        sender,
80        vec![object.compute_object_reference()],
81        pt,
82        TEST_ONLY_GAS_UNIT_FOR_TRANSFER, // gas price is 1
83        1,
84    );
85    to_sender_signed_transaction(data, &sender_key)
86}
87
88pub fn make_transaction_data(sender: SuiAddress) -> TransactionData {
89    let object = Object::immutable_with_id_for_testing(ObjectID::random_from_rng(
90        &mut StdRng::from_seed([0; 32]),
91    ));
92    let pt = {
93        let mut builder = ProgrammableTransactionBuilder::new();
94        builder.transfer_sui(dbg_addr(2), None);
95        builder.finish()
96    };
97    TransactionData::new_programmable(
98        sender,
99        vec![object.compute_object_reference()],
100        pt,
101        TEST_ONLY_GAS_UNIT_FOR_TRANSFER, // gas price is 1
102        1,
103    )
104}
105
106/// Make a user signed transaction with the given sender and its keypair. This
107/// is not verified or signed by authority.
108pub fn make_transaction(sender: SuiAddress, kp: &SuiKeyPair) -> Transaction {
109    let data = make_transaction_data(sender);
110    Transaction::from_data_and_signer(data, vec![kp])
111}
112
113// This is used to sign transaction with signer using default Intent.
114pub fn to_sender_signed_transaction(
115    data: TransactionData,
116    signer: &dyn Signer<Signature>,
117) -> Transaction {
118    to_sender_signed_transaction_with_multi_signers(data, vec![signer])
119}
120
121pub fn to_sender_signed_transaction_with_multi_signers(
122    data: TransactionData,
123    signers: Vec<&dyn Signer<Signature>>,
124) -> Transaction {
125    Transaction::from_data_and_signer(data, signers)
126}
127
128mod zk_login {
129    use fastcrypto_zkp::bn254::zk_login::ZkLoginInputs;
130    use shared_crypto::intent::PersonalMessage;
131
132    use crate::{crypto::PublicKey, zk_login_util::get_zklogin_inputs};
133
134    use super::*;
135    pub static DEFAULT_ADDRESS_SEED: &str =
136        "20794788559620669596206457022966176986688727876128223628113916380927502737911";
137    pub static SHORT_ADDRESS_SEED: &str =
138        "380704556853533152350240698167704405529973457670972223618755249929828551006";
139
140    pub fn load_test_vectors(path: &str) -> Vec<(SuiKeyPair, PublicKey, ZkLoginInputs)> {
141        // read in test files that has a list of matching zklogin_inputs and its ephemeral private keys.
142        let file = std::fs::File::open(path).expect("Unable to open file");
143
144        let test_datum: Vec<TestData> = serde_json::from_reader(file).unwrap();
145        let mut res = vec![];
146        for test in test_datum {
147            let kp = SuiKeyPair::decode(&test.kp).unwrap();
148            let inputs =
149                ZkLoginInputs::from_json(&test.zklogin_inputs, &test.address_seed).unwrap();
150            let pk_zklogin = PublicKey::from_zklogin_inputs(&inputs).unwrap();
151            res.push((kp, pk_zklogin, inputs));
152        }
153        res
154    }
155    pub fn get_one_zklogin_inputs(path: &str) -> String {
156        let file = std::fs::File::open(path).expect("Unable to open file");
157
158        let test_data: Vec<TestData> = serde_json::from_reader(file).unwrap();
159        test_data[1].zklogin_inputs.clone()
160    }
161
162    pub fn get_zklogin_user_address() -> SuiAddress {
163        thread_local! {
164            static USER_ADDRESS: SuiAddress = {
165                // Derive user address manually: Blake2b_256 hash of [zklogin_flag || iss_bytes_length || iss_bytes || address seed in bytes])
166                let mut hasher = DefaultHash::default();
167                hasher.update([SignatureScheme::ZkLoginAuthenticator.flag()]);
168                let inputs = get_zklogin_inputs();
169                let iss_bytes = inputs.get_iss().as_bytes();
170                hasher.update([iss_bytes.len() as u8]);
171                hasher.update(iss_bytes);
172                hasher.update(inputs.get_address_seed().unpadded());
173                SuiAddress::from_bytes(hasher.finalize().digest).unwrap()
174            };
175        }
176        USER_ADDRESS.with(|a| *a)
177    }
178
179    fn get_zklogin_user_key() -> SuiKeyPair {
180        SuiKeyPair::Ed25519(Ed25519KeyPair::generate(&mut StdRng::from_seed([0; 32])))
181    }
182
183    fn get_inputs_with_bad_address_seed() -> ZkLoginInputs {
184        thread_local! {
185        static ZKLOGIN_INPUTS: ZkLoginInputs = ZkLoginInputs::from_json("{\"proofPoints\":{\"a\":[\"17276311605393076686048412951904952585208929623427027497902331765285829154985\",\"2195957390349729412627479867125563520760023859523358729791332629632025124364\",\"1\"],\"b\":[[\"10285059021604767951039627893758482248204478992077021270802057708215366770814\",\"20086937595807139308592304218494658586282197458549968652049579308384943311509\"],[\"7481123765095657256931104563876569626157448050870256177668773471703520958615\",\"11912752790863530118410797223176516777328266521602785233083571774104055633375\"],[\"1\",\"0\"]],\"c\":[\"15742763887654796666500488588763616323599882100448686869458326409877111249163\",\"6112916537574993759490787691149241262893771114597679488354854987586060572876\",\"1\"]},\"issBase64Details\":{\"value\":\"wiaXNzIjoiaHR0cHM6Ly9pZC50d2l0Y2gudHYvb2F1dGgyIiw\",\"indexMod4\":2},\"headerBase64\":\"eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6IjEifQ\"}", SHORT_ADDRESS_SEED).unwrap(); }
186        ZKLOGIN_INPUTS.with(|a| a.clone())
187    }
188
189    pub fn get_legacy_zklogin_user_address() -> SuiAddress {
190        thread_local! {
191            static USER_ADDRESS: SuiAddress = {
192                let inputs = get_inputs_with_bad_address_seed();
193                SuiAddress::from(&PublicKey::from_zklogin_inputs(&inputs).unwrap())
194            };
195        }
196        USER_ADDRESS.with(|a| *a)
197    }
198
199    pub fn sign_zklogin_personal_msg(data: PersonalMessage) -> (SuiAddress, GenericSignature) {
200        let inputs = get_zklogin_inputs();
201        let msg = IntentMessage::new(Intent::personal_message(), data);
202        let s = Signature::new_secure(&msg, &get_zklogin_user_key());
203        let authenticator =
204            GenericSignature::ZkLoginAuthenticator(ZkLoginAuthenticator::new(inputs, 10, s));
205        let address = get_zklogin_user_address();
206        (address, authenticator)
207    }
208
209    pub fn sign_zklogin_tx_with_default_proof(
210        data: TransactionData,
211        legacy: bool,
212    ) -> (SuiAddress, Transaction, GenericSignature) {
213        let inputs = if legacy {
214            get_inputs_with_bad_address_seed()
215        } else {
216            get_zklogin_inputs()
217        };
218
219        sign_zklogin_tx(&get_zklogin_user_key(), inputs, data)
220    }
221
222    pub fn sign_zklogin_tx(
223        user_key: &SuiKeyPair,
224        proof: ZkLoginInputs,
225        data: TransactionData,
226    ) -> (SuiAddress, Transaction, GenericSignature) {
227        let tx = Transaction::from_data_and_signer(data.clone(), vec![user_key]);
228
229        let s = match tx.inner().tx_signatures.first().unwrap() {
230            GenericSignature::Signature(s) => s,
231            _ => panic!("Expected a signature"),
232        };
233
234        // Construct the authenticator with all user submitted components.
235        let authenticator =
236            GenericSignature::ZkLoginAuthenticator(ZkLoginAuthenticator::new(proof, 10, s.clone()));
237
238        let tx = Transaction::new(SenderSignedData::new(
239            tx.transaction_data().clone(),
240            vec![authenticator.clone()],
241        ));
242        (data.execution_parts().1, tx, authenticator)
243    }
244
245    pub fn make_zklogin_tx(
246        address: SuiAddress,
247        legacy: bool,
248    ) -> (SuiAddress, Transaction, GenericSignature) {
249        let data = make_transaction_data(address);
250        sign_zklogin_tx_with_default_proof(data, legacy)
251    }
252
253    pub fn keys() -> Vec<SuiKeyPair> {
254        let mut seed = StdRng::from_seed([0; 32]);
255        let kp1: SuiKeyPair = SuiKeyPair::Ed25519(get_key_pair_from_rng(&mut seed).1);
256        let kp2: SuiKeyPair = SuiKeyPair::Secp256k1(get_key_pair_from_rng(&mut seed).1);
257        let kp3: SuiKeyPair = SuiKeyPair::Secp256r1(get_key_pair_from_rng(&mut seed).1);
258        vec![kp1, kp2, kp3]
259    }
260
261    pub fn make_upgraded_multisig_tx() -> Transaction {
262        let keys = keys();
263        let pk1 = &keys[0].public();
264        let pk2 = &keys[1].public();
265        let pk3 = &keys[2].public();
266
267        let multisig_pk = MultiSigPublicKey::new(
268            vec![pk1.clone(), pk2.clone(), pk3.clone()],
269            vec![1, 1, 1],
270            2,
271        )
272        .unwrap();
273        let addr = SuiAddress::from(&multisig_pk);
274        let tx = make_transaction(addr, &keys[0]);
275
276        let msg = IntentMessage::new(Intent::sui_transaction(), tx.transaction_data().clone());
277        let sig1 = Signature::new_secure(&msg, &keys[0]).into();
278        let sig2 = Signature::new_secure(&msg, &keys[1]).into();
279
280        // Any 2 of 3 signatures verifies ok.
281        let multi_sig1 = MultiSig::combine(vec![sig1, sig2], multisig_pk).unwrap();
282        Transaction::new(SenderSignedData::new(
283            tx.transaction_data().clone(),
284            vec![GenericSignature::MultiSig(multi_sig1)],
285        ))
286    }
287}
288pub use zk_login::*;