libzkbob_rs/client/
mod.rs

1use kvdb::KeyValueDB;
2use libzeropool::{
3    constants,
4    fawkes_crypto::ff_uint::PrimeField,
5    fawkes_crypto::{
6        core::sizedvec::SizedVec,
7        ff_uint::Num,
8        ff_uint::{NumRepr, Uint},
9        rand::Rng,
10    },
11    native::{
12        account::Account,
13        boundednum::BoundedNum,
14        cipher,
15        key::derive_key_p_d,
16        note::Note,
17        params::PoolParams,
18        tx::{
19            make_delta, nullifier, nullifier_intermediate_hash,
20            out_commitment_hash, tx_hash, tx_sign,
21            TransferPub, TransferSec,
22            Tx,
23        },
24    },
25};
26
27use serde::{Deserialize, Serialize};
28use std::{convert::TryInto, io::Write, ops::Range};
29use thiserror::Error;
30
31use self::state::{State, Transaction};
32use crate::{
33    address::{format_address, parse_address, AddressParseError, AddressFormat},
34    keys::{reduce_sk, Keys},
35    random::CustomRng,
36    merkle::Hash,
37    utils::{keccak256, zero_note, zero_proof},
38};
39
40pub mod state;
41
42
43
44#[derive(Debug, Error)]
45pub enum CreateTxError {
46    #[error("Too many outputs: expected {max} max got {got}")]
47    TooManyOutputs { max: usize, got: usize },
48    #[error("Too few outputs: expected {min} min got {got}")]
49    TooFewOutputs { min: usize, got: usize },
50    #[error("Could not get merkle proof for leaf {0}")]
51    ProofNotFound(u64),
52    #[error("Failed to parse address: {0}")]
53    AddressParseError(#[from] AddressParseError),
54    #[error("Insufficient balance: sum of outputs is greater than sum of inputs: {0} > {1}")]
55    InsufficientBalance(String, String),
56    #[error("Insufficient energy: available {0}, received {1}")]
57    InsufficientEnergy(String, String),
58    #[error("Failed to serialize transaction: {0}")]
59    IoError(#[from] std::io::Error),
60}
61
62#[derive(Serialize, Deserialize, Default, Clone)]
63pub struct StateFragment<Fr: PrimeField> {
64    pub new_leafs: Vec<(u64, Vec<Hash<Fr>>)>,
65    pub new_commitments: Vec<(u64, Hash<Fr>)>,
66    pub new_accounts: Vec<(u64, Account<Fr>)>,
67    pub new_notes: Vec<(u64, Note<Fr>)>,
68}
69
70#[derive(Serialize, Deserialize, Clone)]
71pub struct TransactionData<Fr: PrimeField> {
72    pub public: TransferPub<Fr>,
73    pub secret: TransferSec<Fr>,
74    pub ciphertext: Vec<u8>,
75    pub memo: Vec<u8>,
76    pub commitment_root: Num<Fr>,
77    pub out_hashes: SizedVec<Num<Fr>, { constants::OUT + 1 }>,
78}
79
80#[derive(Serialize, Deserialize, Clone)]
81pub struct TransactionInputs<Fr: PrimeField> {
82    pub account: (u64, Account<Fr>),
83    pub intermediate_nullifier: Num<Fr>,   // intermediate nullifier hash
84    pub notes: Vec<(u64, Note<Fr>)>,
85}
86
87pub type TokenAmount<Fr> = BoundedNum<Fr, { constants::BALANCE_SIZE_BITS }>;
88
89#[derive(Debug, Serialize, Deserialize, Clone)]
90pub struct TxOutput<Fr: PrimeField> {
91    pub to: String,
92    pub amount: TokenAmount<Fr>,
93}
94
95#[derive(Debug, Serialize, Deserialize, Clone)]
96pub enum TxType<Fr: PrimeField> {
97    // fee, data, tx_outputs
98    Transfer(TokenAmount<Fr>, Vec<u8>, Vec<TxOutput<Fr>>),
99    // fee, data, deposit_amount
100    Deposit(TokenAmount<Fr>, Vec<u8>, TokenAmount<Fr>),
101    // fee, data, deposit_amount, deadline, holder
102    DepositPermittable(
103        TokenAmount<Fr>,
104        Vec<u8>,
105        TokenAmount<Fr>,
106        u64,
107        Vec<u8>
108    ),
109    // fee, data, withdraw_amount, to, native_amount, energy_amount
110    Withdraw(
111        TokenAmount<Fr>,
112        Vec<u8>,
113        TokenAmount<Fr>,
114        Vec<u8>,
115        TokenAmount<Fr>,
116        TokenAmount<Fr>,
117    ),
118}
119
120pub struct UserAccount<D: KeyValueDB, P: PoolParams> {
121    pub pool_id: u32,
122    pub keys: Keys<P>,
123    pub params: P,
124    pub state: State<D, P>,
125    pub sign_callback: Option<Box<dyn Fn(&[u8]) -> Vec<u8>>>, // TODO: Find a way to make it async
126}
127
128impl<D, P> UserAccount<D, P>
129where
130    D: KeyValueDB,
131    P: PoolParams,
132    P::Fr: 'static,
133{
134    /// Initializes UserAccount with a spending key that has to be an element of the prime field Fs (p = 6554484396890773809930967563523245729705921265872317281365359162392183254199).
135    pub fn new(sk: Num<P::Fs>, pool_id: u32, state: State<D, P>, params: P) -> Self {
136        let keys = Keys::derive(sk, &params);
137
138        UserAccount {
139            pool_id,
140            keys,
141            state,
142            params,
143            sign_callback: None,
144        }
145    }
146
147    /// Same as constructor but accepts arbitrary data as spending key.
148    pub fn from_seed(seed: &[u8], pool_id: u32, state: State<D, P>, params: P) -> Self {
149        let sk = reduce_sk(seed);
150        Self::new(sk, pool_id, state, params)
151    }
152
153    fn generate_address_components(
154        &self,
155    ) -> (
156        BoundedNum<P::Fr, { constants::DIVERSIFIER_SIZE_BITS }>,
157        Num<P::Fr>,
158    ) {
159        let mut rng = CustomRng;
160
161        let d: BoundedNum<_, { constants::DIVERSIFIER_SIZE_BITS }> = rng.gen();
162        let pk_d = derive_key_p_d(d.to_num(), self.keys.eta, &self.params);
163        (d, pk_d.x)
164    }
165
166    /// Generates a new private address for the current pool
167    pub fn generate_address(&self) -> String {
168        let (d, p_d) = self.generate_address_components();
169
170        format_address::<P>(d, p_d, Some(self.pool_id))
171    }
172
173    /// Generates a new private generic address (for all pools)
174    pub fn generate_universal_address(&self) -> String {
175        let (d, p_d) = self.generate_address_components();
176
177        format_address::<P>(d, p_d, None)
178    }
179
180    pub fn generate_address_from_components(
181        &self,
182        d: BoundedNum<P::Fr, { constants::DIVERSIFIER_SIZE_BITS }>,
183        p_d: Num<P::Fr>
184    ) -> String {
185        format_address::<P>(d, p_d, Some(self.pool_id))
186    }
187
188    pub fn generate_universal_address_from_components(
189        &self,
190        d: BoundedNum<P::Fr, { constants::DIVERSIFIER_SIZE_BITS }>,
191        p_d: Num<P::Fr>
192    ) -> String {
193        format_address::<P>(d, p_d, None)
194    }
195
196    pub fn gen_address_for_seed(&self, seed: &[u8]) -> String {
197        self.gen_address_for_seed_and_pool_id(seed, Some(self.pool_id))
198    }
199
200    pub fn gen_universal_address_for_seed(&self, seed: &[u8]) -> String {
201        self.gen_address_for_seed_and_pool_id(seed, None)
202    }
203
204    fn gen_address_for_seed_and_pool_id(&self, seed: &[u8], pool_id: Option<u32>) -> String {
205        let mut rng = CustomRng;
206
207        let sk = reduce_sk(seed);
208        let keys = Keys::derive(sk, &self.params);
209        let d: BoundedNum<_, { constants::DIVERSIFIER_SIZE_BITS }> = rng.gen();
210        let pk_d = derive_key_p_d(d.to_num(), keys.eta, &self.params);
211
212        format_address::<P>(d, pk_d.x, pool_id)
213    }
214
215    pub fn validate_address(&self, address: &str) -> bool {
216        match parse_address(address, &self.params, self.pool_id) {
217            Ok((_, _, format)) => format == AddressFormat::PoolSpecific,
218            Err(_) => false,
219        }
220    }
221
222    pub fn validate_universal_address(&self, address: &str) -> bool {
223        match parse_address(address, &self.params, self.pool_id) {
224            Ok((_, _, format)) => format == AddressFormat::Generic,
225            Err(_) => false,
226        }
227    }
228
229    pub fn is_own_address(&self, address: &str) -> bool {
230        match parse_address::<P>(address, &self.params, self.pool_id) {
231            Ok((d, p_d, _)) => {
232                self.is_derived_from_our_sk(d, p_d)
233            },
234            Err(_) => false
235        }
236    }
237
238    pub fn is_derived_from_our_sk(&self, d: BoundedNum<P::Fr, { constants::DIVERSIFIER_SIZE_BITS }>, p_d: Num<P::Fr>) -> bool {
239        derive_key_p_d(d.to_num(), self.keys.eta, &self.params).x == p_d
240    }
241
242    /// Attempts to decrypt notes.
243    pub fn decrypt_notes(&self, data: Vec<u8>) -> Vec<Option<Note<P::Fr>>> {
244        cipher::decrypt_in(self.keys.eta, &data, &self.params)
245    }
246
247    /// Attempts to decrypt account and notes.
248    pub fn decrypt_pair(&self, data: Vec<u8>) -> Option<(Account<P::Fr>, Vec<Note<P::Fr>>)> {
249        cipher::decrypt_out(self.keys.eta, &data, &self.params)
250    }
251
252    pub fn initial_account(&self) -> Account<P::Fr> {
253        // Initial account should have d = pool_id to protect from reply attacks
254        let d = Num::<P::Fr>::from_uint(NumRepr(Uint::from_u64(self.pool_id as u64))).unwrap();
255        let p_d = derive_key_p_d(d, self.keys.eta, &self.params).x;
256        Account {
257            d: BoundedNum::new(d),
258            p_d,
259            i: BoundedNum::new(Num::ZERO),
260            b: BoundedNum::new(Num::ZERO),
261            e: BoundedNum::new(Num::ZERO),
262        }
263    }
264
265    /// Constructs a transaction.
266    pub fn create_tx(
267        &self,
268        tx: TxType<P::Fr>,
269        delta_index: Option<u64>,
270        extra_state: Option<StateFragment<P::Fr>>
271    ) -> Result<TransactionData<P::Fr>, CreateTxError> {
272        let mut rng = CustomRng;
273        let keys = self.keys.clone();
274        let state = &self.state;
275
276        let extra_state = extra_state.unwrap_or(
277            StateFragment {
278                new_leafs: [].to_vec(),
279                new_commitments: [].to_vec(),
280                new_accounts: [].to_vec(),
281                new_notes: [].to_vec(),
282            }
283        );
284
285        // initial input account (from optimistic state)
286        let (in_account_optimistic_index, in_account_optimistic) = {
287            let last_acc = extra_state.new_accounts.last();
288            match last_acc {
289                Some(last_acc) => (Some(last_acc.0), Some(last_acc.1)),
290                _ => (None, None),
291            }
292        };
293
294        // initial input account (from non-optimistic state)
295        let in_account = in_account_optimistic.unwrap_or_else(|| {
296            state.latest_account.unwrap_or_else(|| self.initial_account())
297        });
298
299        let tree = &self.state.tree;
300
301        let in_account_index = in_account_optimistic_index.or(state.latest_account_index);
302
303        // initial usable note index
304        let next_usable_index = state.earliest_usable_index_optimistic(&extra_state.new_accounts, &extra_state.new_notes);
305
306        let latest_note_index_optimistic = extra_state.new_notes
307            .last()
308            .map(|indexed_note| indexed_note.0)
309            .unwrap_or(state.latest_note_index);
310
311        // Should be provided by relayer together with note proofs, but as a fallback
312        // take the next index of the tree (optimistic part included).
313        let delta_index = Num::from(delta_index.unwrap_or_else( || {
314            let next_by_optimistic_leaf = extra_state.new_leafs
315                .last()
316                .map(|leafs| {
317                    (((leafs.0 + (leafs.1.len() as u64) - 1) >> constants::OUTPLUSONELOG) + 1) << constants::OUTPLUSONELOG
318                });
319            let next_by_optimistic_commitment = extra_state.new_commitments
320                .last()
321                .map(|commitment| {
322                    ((commitment.0  >> constants::OUTPLUSONELOG) + 1) << constants::OUTPLUSONELOG
323                });
324            next_by_optimistic_leaf
325                .into_iter()
326                .chain(next_by_optimistic_commitment)
327                .max()
328                .unwrap_or_else(|| self.state.tree.next_index())
329        }));
330
331        let (fee, tx_data, user_data) = {
332            let mut tx_data: Vec<u8> = vec![];
333            match &tx {
334                TxType::Deposit(fee, user_data, _) => {
335                    let raw_fee: u64 = fee.to_num().try_into().unwrap();
336                    tx_data.write_all(&raw_fee.to_be_bytes()).unwrap();
337                    (fee, tx_data, user_data)
338                }
339                TxType::DepositPermittable(fee, user_data, _, deadline, holder) => {
340                    let raw_fee: u64 = fee.to_num().try_into().unwrap();
341
342                    tx_data.write_all(&raw_fee.to_be_bytes()).unwrap();
343                    tx_data.write_all(&deadline.to_be_bytes()).unwrap();
344                    tx_data.append(&mut holder.clone());
345                    
346                    (fee, tx_data, user_data)
347                }
348                TxType::Transfer(fee, user_data, _) => {
349                    let raw_fee: u64 = fee.to_num().try_into().unwrap();
350                    tx_data.write_all(&raw_fee.to_be_bytes()).unwrap();
351                    (fee, tx_data, user_data)
352                }
353                TxType::Withdraw(fee, user_data, _, reciever, native_amount, _) => {
354                    let raw_fee: u64 = fee.to_num().try_into().unwrap();
355                    let raw_native_amount: u64 = native_amount.to_num().try_into().unwrap();
356
357                    tx_data.write_all(&raw_fee.to_be_bytes()).unwrap();
358                    tx_data.write_all(&raw_native_amount.to_be_bytes()).unwrap();
359                    tx_data.append(&mut reciever.clone());
360
361                    (fee, tx_data, user_data)
362                }
363            }
364        };
365
366        // Optimistic available notes
367        let optimistic_available_notes = extra_state.new_notes
368            .into_iter()
369            .filter(|indexed_note| indexed_note.0 >= next_usable_index);
370
371        // Fetch constants::IN usable notes from state
372        let in_notes_original: Vec<(u64, Note<P::Fr>)> = state
373            .txs
374            .iter_slice(next_usable_index..=state.latest_note_index)
375            .filter_map(|(index, tx)| match tx {
376                Transaction::Note(note) => Some((index, note)),
377                _ => None,
378            })
379            .chain(optimistic_available_notes)
380            .take(constants::IN)
381            .collect();
382        
383        let spend_interval_index = in_notes_original
384            .last()
385            .map(|(index, _)| *index + 1)
386            .unwrap_or(if latest_note_index_optimistic > 0 { latest_note_index_optimistic + 1 } else { 0 });
387
388        // Calculate total balance (account + constants::IN notes).
389        let mut input_value = in_account.b.to_num();
390        for (_index, note) in &in_notes_original {
391            input_value += note.b.to_num();
392        }
393
394        let mut output_value = Num::ZERO;
395
396        let (num_real_out_notes, out_notes): (_, SizedVec<_, { constants::OUT }>) =
397            if let TxType::Transfer(_, _, outputs) = &tx {
398                if outputs.len() > constants::OUT {
399                    return Err(CreateTxError::TooManyOutputs {
400                        max: constants::OUT,
401                        got: outputs.len(),
402                    });
403                }
404
405                let out_notes = outputs
406                    .iter()
407                    .map(|dest| {
408                        let (to_d, to_p_d, _) = parse_address::<P>(&dest.to, &self.params, self.pool_id)?;
409                        output_value += dest.amount.to_num();
410                        Ok(Note {
411                            d: to_d,
412                            p_d: to_p_d,
413                            b: dest.amount,
414                            t: rng.gen(),
415                        })
416                    })
417                    // fill out remaining output notes with zeroes
418                    .chain((0..).map(|_| Ok(zero_note())))
419                    .take(constants::OUT)
420                    .collect::<Result<SizedVec<_, { constants::OUT }>, CreateTxError>>()?;
421
422                (outputs.len(), out_notes)
423            } else {
424                (0, (0..).map(|_| zero_note()).take(constants::OUT).collect())
425            };
426
427        let mut delta_value = -fee.as_num();
428        // By default all account energy will be withdrawn on withdraw tx
429        let mut delta_energy = Num::ZERO;
430
431        let in_account_pos = in_account_index.unwrap_or(0);
432
433        let mut input_energy = in_account.e.to_num();
434        input_energy += in_account.b.to_num() * (delta_index - Num::from(in_account_pos));
435
436        for (note_index, note) in &in_notes_original {
437            input_energy += note.b.to_num() * (delta_index - Num::from(*note_index));
438        }
439        let new_balance = match &tx {
440            TxType::Transfer(_, _, _) => {
441                if input_value.to_uint() >= (output_value + fee.as_num()).to_uint() {
442                    input_value - output_value - fee.as_num()
443                } else {
444                    return Err(CreateTxError::InsufficientBalance(
445                        (output_value + fee.as_num()).to_string(),
446                        input_value.to_string(),
447                    ));
448                }
449            }
450            TxType::Withdraw(_, _, amount, _, _, energy) => {
451                let amount = amount.to_num();
452                let energy = energy.to_num();
453
454                if energy.to_uint() > input_energy.to_uint() {
455                    return Err(CreateTxError::InsufficientEnergy(
456                        input_energy.to_string(),
457                        energy.to_string(),
458                    ));
459                }
460
461                delta_energy -= energy;
462                delta_value -= amount;
463
464                if input_value.to_uint() >= amount.to_uint() {
465                    input_value + delta_value
466                } else {
467                    return Err(CreateTxError::InsufficientBalance(
468                        delta_value.to_string(),
469                        input_value.to_string(),
470                    ));
471                }
472            }
473            TxType::Deposit(_, _, amount) | TxType::DepositPermittable(_, _, amount, _, _) => {
474                delta_value += amount.to_num();
475                input_value + delta_value
476            }
477        };
478
479        let (d, p_d) = self.generate_address_components();
480        let out_account = Account {
481            d,
482            p_d,
483            i: BoundedNum::new(Num::from(spend_interval_index)),
484            b: BoundedNum::new(new_balance),
485            e: BoundedNum::new(delta_energy + input_energy),
486        };
487
488        let in_account_hash = in_account.hash(&self.params);
489        let nullifier = nullifier(
490            in_account_hash,
491            keys.eta,
492            in_account_pos.into(),
493            &self.params,
494        );
495
496        let ciphertext = {
497            let entropy: [u8; 32] = rng.gen();
498
499            // No need to include all the zero notes in the encrypted transaction
500            let out_notes = &out_notes[0..num_real_out_notes];
501
502            cipher::encrypt(&entropy, keys.eta, out_account, out_notes, &self.params)
503        };
504
505        // Hash input account + notes filling remaining space with non-hashed zeroes
506        let owned_zero_notes = (0..).map(|_| {
507            let d: BoundedNum<_, { constants::DIVERSIFIER_SIZE_BITS }> = rng.gen();
508            let p_d = derive_key_p_d::<P, P::Fr>(d.to_num(), keys.eta, &self.params).x;
509            Note {
510                d,
511                p_d,
512                b: BoundedNum::new(Num::ZERO),
513                t: rng.gen(),
514            }
515        });
516        let in_notes: SizedVec<Note<P::Fr>, { constants::IN }> = in_notes_original
517            .iter()
518            .map(|(_, note)| note)
519            .cloned()
520            .chain(owned_zero_notes)
521            .take(constants::IN)
522            .collect();
523        let in_note_hashes = in_notes.iter().map(|note| note.hash(&self.params));
524        let input_hashes: SizedVec<_, { constants::IN + 1 }> = [in_account_hash]
525            .iter()
526            .copied()
527            .chain(in_note_hashes)
528            .collect();
529
530        // Same with output
531        let out_account_hash = out_account.hash(&self.params);
532        let out_note_hashes = out_notes.iter().map(|n| n.hash(&self.params));
533        let out_hashes: SizedVec<Num<P::Fr>, { constants::OUT + 1 }> = [out_account_hash]
534            .iter()
535            .copied()
536            .chain(out_note_hashes)
537            .collect();
538
539        let out_commit = out_commitment_hash(out_hashes.as_slice(), &self.params);
540        let tx_hash = tx_hash(input_hashes.as_slice(), out_commit, &self.params);
541
542        let delta = make_delta::<P::Fr>(
543            delta_value,
544            delta_energy,
545            delta_index,
546            Num::<P::Fr>::from_uint(NumRepr(Uint::from_u64(self.pool_id as u64))).unwrap(),
547        );
548
549        // calculate virtual subtree from the optimistic state
550        let new_leafs = extra_state.new_leafs.iter().cloned();
551        let new_commitments = extra_state.new_commitments.iter().cloned();
552        let (mut virtual_nodes, update_boundaries) = tree.get_virtual_subtree(new_leafs, new_commitments);
553
554        let root: Num<P::Fr> = tree.get_root_optimistic(&mut virtual_nodes, &update_boundaries);
555
556        // memo = tx_specific_data, ciphertext, user_defined_data
557        let mut memo_data = {
558            let tx_data_size = tx_data.len();
559            let ciphertext_size = ciphertext.len();
560            let user_data_size = user_data.len();
561            Vec::with_capacity(tx_data_size + ciphertext_size + user_data_size)
562        };
563
564        #[allow(clippy::redundant_clone)]
565        memo_data.append(&mut tx_data.clone());
566        memo_data.extend(&ciphertext);
567        memo_data.append(&mut user_data.clone());
568
569        let memo_hash = keccak256(&memo_data);
570        let memo = Num::from_uint_reduced(NumRepr(Uint::from_big_endian(&memo_hash)));
571
572        let public = TransferPub::<P::Fr> {
573            root,
574            nullifier,
575            out_commit,
576            delta,
577            memo,
578        };
579
580        let tx = Tx {
581            input: (in_account, in_notes),
582            output: (out_account, out_notes),
583        };
584
585        let (eddsa_s, eddsa_r) = tx_sign(keys.sk, tx_hash, &self.params);
586
587        let account_proof = in_account_index.map_or_else(
588            || Ok(zero_proof()),
589            |i| {
590                tree.get_proof_optimistic_index(i, &mut virtual_nodes, &update_boundaries)
591                    .ok_or(CreateTxError::ProofNotFound(i))
592            },
593        )?;
594        let note_proofs = in_notes_original
595            .iter()
596            .copied()
597            .map(|(index, _note)| {
598                tree.get_proof_optimistic_index(index, &mut virtual_nodes, &update_boundaries)
599                    .ok_or(CreateTxError::ProofNotFound(index))
600            })
601            .chain((0..).map(|_| Ok(zero_proof())))
602            .take(constants::IN)
603            .collect::<Result<_, _>>()?;
604
605        
606        let secret = TransferSec::<P::Fr> {
607            tx,
608            in_proof: (account_proof, note_proofs),
609            eddsa_s: eddsa_s.to_other().unwrap(),
610            eddsa_r,
611            eddsa_a: keys.a,
612        };
613
614        Ok(TransactionData {
615            public,
616            secret,
617            ciphertext,
618            memo: memo_data,
619            commitment_root: out_commit,
620            out_hashes,
621        })
622    }
623
624    pub fn get_tx_input(&self, index: u64) -> Option<TransactionInputs<P::Fr>> {
625        let account = match self.state.get_account(index) {
626            Some(acc) => acc,
627            _ => return None,
628        };
629
630        let input_acc = self.state.get_previous_account(index).unwrap_or_else(|| (0, self.initial_account()));
631        let note_lower_bound = input_acc.1.i.to_num().try_into().unwrap();
632        let note_upper_bound = account.i.to_num().try_into().unwrap();
633        let notes_range: Range<u64> = note_lower_bound..note_upper_bound;
634        let input_notes = self.state.get_notes_in_range(notes_range);
635
636        let params = &self.params;
637        let inh = nullifier_intermediate_hash(input_acc.1.hash(params), self.keys.eta, index.into(), params);
638
639        Some(TransactionInputs {
640            account: input_acc,
641            intermediate_nullifier: inh,
642            notes: input_notes,
643        })
644    }
645}
646
647#[cfg(test)]
648mod tests {
649    use std::str::FromStr;
650
651    use super::*;
652    use libzeropool::POOL_PARAMS;
653    use crate::random::CustomRng;
654
655    #[test]
656    fn test_create_tx_deposit_zero() {
657        let state = State::init_test(POOL_PARAMS.clone());
658        let acc = UserAccount::new(Num::ZERO, 0, state, POOL_PARAMS.clone());
659
660        acc.create_tx(
661            TxType::Deposit(
662                BoundedNum::ZERO,
663                vec![],
664                BoundedNum::ZERO,
665            ),
666            None,
667            None,
668        )
669        .unwrap();
670    }
671
672    #[test]
673    fn test_create_tx_deposit_one() {
674        let state = State::init_test(POOL_PARAMS.clone());
675        let acc = UserAccount::new(Num::ZERO, 0, state, POOL_PARAMS.clone());
676
677        acc.create_tx(
678            TxType::Deposit(
679                BoundedNum::new(Num::ZERO),
680                vec![],
681                BoundedNum::new(Num::ONE),
682            ),
683            None,
684            None,
685        )
686        .unwrap();
687    }
688
689    // It's ok to transfer 0 while balance = 0
690    #[test]
691    fn test_create_tx_transfer_zero() {
692        let state = State::init_test(POOL_PARAMS.clone());
693        let acc = UserAccount::new(Num::ZERO, 0, state, POOL_PARAMS.clone());
694
695        let addr = acc.generate_address();
696
697        let out = TxOutput {
698            to: addr,
699            amount: BoundedNum::new(Num::ZERO),
700        };
701
702        acc.create_tx(
703            TxType::Transfer(BoundedNum::new(Num::ZERO), vec![], vec![out]),
704            None,
705            None,
706        )
707        .unwrap();
708    }
709
710    #[test]
711    #[should_panic]
712    fn test_create_tx_transfer_one_no_balance() {
713        let state = State::init_test(POOL_PARAMS.clone());
714        let acc = UserAccount::new(Num::ZERO, 0, state, POOL_PARAMS.clone());
715
716        let addr = acc.generate_address();
717
718        let out = TxOutput {
719            to: addr,
720            amount: BoundedNum::new(Num::ONE),
721        };
722
723        acc.create_tx(
724            TxType::Transfer(BoundedNum::new(Num::ZERO), vec![], vec![out]),
725            None,
726            None,
727        )
728        .unwrap();
729    }
730
731    #[test]
732    fn test_user_account_is_own_address() {
733        let acc_1 = UserAccount::new(
734            Num::ZERO,
735            0xffff02,
736            State::init_test(POOL_PARAMS.clone()),
737            POOL_PARAMS.clone(),
738        );
739        let acc_2 = UserAccount::new(
740            Num::ONE,
741            0xffff02,
742            State::init_test(POOL_PARAMS.clone()),
743            POOL_PARAMS.clone(),
744        );
745
746        let address_1 = acc_1.generate_address();
747        let address_2 = acc_2.generate_address();
748
749        assert!(acc_1.is_own_address(&address_1));
750        assert!(acc_2.is_own_address(&address_2));
751
752        assert!(!acc_1.is_own_address(&address_2));
753        assert!(!acc_2.is_own_address(&address_1));
754    }
755
756    #[test]
757    fn test_tx_inputs() {
758        let params = POOL_PARAMS.clone();
759        let mut rng = CustomRng;
760        let state = State::init_test(POOL_PARAMS.clone());
761        let mut user_account = UserAccount::new(
762            Num::ZERO,
763            1,
764            state,
765            POOL_PARAMS.clone()
766        );
767
768        let mut acc0 = Account::sample(&mut rng, &params);
769        acc0.i = BoundedNum::new(Num::from_str("0").unwrap());
770        let mut acc1 = Account::sample(&mut rng, &params);
771        acc1.i = BoundedNum::new(Num::from_str("51").unwrap());
772        let mut acc2 = Account::sample(&mut rng, &params);
773        acc2.i = BoundedNum::new(Num::from_str("259").unwrap());
774
775        let note0 = (50, Note::sample(&mut rng, &params));
776        let note1 = (257, Note::sample(&mut rng, &params));
777        let note2 = (258, Note::sample(&mut rng, &params));
778        let note3 = (259, Note::sample(&mut rng, &params));
779        let note4 = (300, Note::sample(&mut rng, &params));
780        let note5 = (666, Note::sample(&mut rng, &params));
781
782        user_account.state.add_account(0, acc0);
783        user_account.state.add_account(128, acc1);
784        user_account.state.add_note(note0.0, note0.1);
785        user_account.state.add_note(note1.0, note1.1);
786        user_account.state.add_note(note2.0, note2.1);
787        user_account.state.add_note(note3.0, note3.1);
788        user_account.state.add_note(note4.0, note4.1);
789        user_account.state.add_note(note5.0, note5.1);
790        user_account.state.add_account(1024, acc2);
791
792        (0..10).into_iter().for_each(|idx| {
793            user_account.state.add_note(1024 + idx + 1, Note::sample(&mut rng, &params))
794        });
795
796        let inputs0 = user_account.get_tx_input(0).unwrap();
797        assert!(inputs0.account.1 == user_account.initial_account());
798        assert_eq!(inputs0.notes.len(), 0);
799
800        let inputs1 = user_account.get_tx_input(128).unwrap();
801        assert!(inputs1.account.1 == acc0);
802        assert_eq!(inputs1.notes.len(), 1);
803        assert!(inputs1.notes.contains(&note0));
804
805        let inputs2 = user_account.get_tx_input(1024).unwrap();
806        assert!(inputs2.account.1 == acc1);
807        assert_eq!(inputs2.notes.len(), 2);
808        assert!(inputs2.notes.contains(&note1));
809        assert!(inputs2.notes.contains(&note2));
810
811    }
812
813    #[test]
814    fn test_chain_specific_addresses() {
815        let acc_polygon = UserAccount::new(Num::ZERO, 0, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone());
816        let acc_sepolia = UserAccount::new(Num::ZERO, 0, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone());
817        let acc_optimism = UserAccount::new(Num::ZERO, 1, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone());
818        let acc_optimism_eth = UserAccount::new(Num::ZERO, 2, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone());
819
820        assert!(acc_polygon.validate_universal_address("PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kxBwa"));   // generic without a prefix
821        assert!(acc_polygon.validate_universal_address("zkbob:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kxBwa")); // generic
822        assert!(acc_polygon.validate_address("zkbob_polygon:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kRF7i"));
823        assert!(!acc_polygon.validate_address("zkbob_optimism:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L"));
824        assert!(!acc_polygon.validate_address("zkbob_polygon:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L"));
825
826        assert!(acc_sepolia.validate_universal_address("PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kxBwa"));
827        assert!(acc_sepolia.validate_universal_address("zkbob:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kxBwa"));
828        assert!(acc_sepolia.validate_address("zkbob_polygon:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kRF7i")); // true due to the same pool_id
829        assert!(!acc_sepolia.validate_address("zkbob_optimism:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L"));
830        assert!(!acc_sepolia.validate_address("zkbob_polygon:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L"));
831        
832        assert!(acc_optimism.validate_universal_address("PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kxBwa"));
833        assert!(acc_optimism.validate_universal_address("zkbob:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kxBwa"));
834        assert!(!acc_optimism.validate_address("zkbob_polygon:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kRF7i"));
835        assert!(acc_optimism.validate_address("zkbob_optimism:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L"));
836        assert!(acc_optimism.validate_address("zkbob_polygon:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L")); // the lib cannot distinguish different address prefixes
837        assert!(!acc_optimism.validate_address("zkbob_optimism:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kRF7i"));
838        assert!(!acc_optimism.validate_address("zkbob_optimism:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kR*&**7i"));
839        assert!(acc_optimism.validate_address("zkbob_zkbober:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L"));
840        assert!(acc_optimism.validate_address("zkbob:optimism:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L"));
841        assert!(!acc_optimism.validate_universal_address("zkbob:"));
842        assert!(!acc_optimism.validate_address(":"));
843        assert!(!acc_optimism.validate_address(""));
844
845        assert!(acc_optimism_eth.validate_universal_address("PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kxBwa"));
846        assert!(acc_optimism_eth.validate_universal_address("zkbob:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kxBwa"));
847        assert!(acc_optimism_eth.validate_address("zkbob_optimism_eth:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse1j9diw"));
848        assert!(!acc_optimism_eth.validate_address("zkbob_polygon:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kRF7i"));
849        assert!(!acc_optimism_eth.validate_address("zkbob_optimism:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5vHs5L"));
850        assert!(!acc_optimism_eth.validate_address("zkbob_optimism_eth:PtfsqyJhA2yvmLtXBm55pkvFDX6XZrRMaib9F1GvwzmU8U4witUf8Jyse5kRF7i"));
851    }   
852
853    #[test]
854    fn test_chain_specific_address_ownable() {
855        let accs = [
856            UserAccount::new(Num::ZERO, 0, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone()),
857            UserAccount::new(Num::ZERO, 1, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone()),
858            UserAccount::new(Num::ZERO, 2, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone()),
859            UserAccount::new(Num::ZERO, 0xffff02, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone()),
860            UserAccount::new(Num::ZERO, 0xffff03, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone()),
861        ];
862        let acc2 = UserAccount::new(Num::ONE, 1, State::init_test(POOL_PARAMS.clone()), POOL_PARAMS.clone());
863        let pool_addresses: Vec<String> = accs.iter().map(|acc| acc.generate_address()).collect();
864        let universal_addresses: Vec<String> = accs.iter().map(|acc| acc.generate_universal_address()).collect();
865
866        accs.iter().enumerate().for_each(|(acc_idx, acc)| {
867            pool_addresses.iter().enumerate().for_each(|(addr_idx, addr)| {
868                if addr_idx == acc_idx {
869                    assert!(acc.is_own_address(&addr));
870                } else {
871                    assert!(!acc.is_own_address(&addr));
872                }
873                assert!(!acc2.is_own_address(&addr))
874            });
875
876            universal_addresses.iter().for_each(|addr| {
877                assert!(acc.is_own_address(&addr));
878                assert!(!acc2.is_own_address(&addr))
879            });
880        });
881    }
882}