iota_types/unit_tests/
utils.rs

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

iota_types/unit_tests/
utils.rs
