tx.rs: partially revived isolated testcase

src/bin/tx.rs

@@ -0,0 +1,309 @@
+use bellman::groth16;
+use bls12_381::Bls12;
+use ff::{Field, PrimeField};
+use rand::rngs::OsRng;
+use std::path::Path;
+
+use async_std::sync::{Arc, Mutex};
+use drk::crypto::{
+    coin::Coin,
+    load_params,
+    merkle::{CommitmentTree, IncrementalWitness},
+    merkle_node::MerkleNode,
+    note::{EncryptedNote, Note},
+    nullifier::Nullifier,
+    save_params, setup_mint_prover, setup_spend_prover,
+};
+use drk::serial::{Decodable, Encodable};
+use drk::state::{state_transition, ProgramState, StateUpdate};
+use drk::tx;
+
+struct MemoryState {
+    // The entire merkle tree state
+    tree: CommitmentTree<MerkleNode>,
+    // List of all previous and the current merkle roots
+    // This is the hashed value of all the children.
+    merkle_roots: Vec<MerkleNode>,
+    // Nullifiers prevent double spending
+    nullifiers: Vec<Nullifier>,
+    // All received coins
+    // NOTE: we need maybe a flag to keep track of which ones are spent
+    // Maybe the spend field links to a tx hash:input index
+    // We should also keep track of the tx hash:output index where this
+    // coin was received
+    own_coins: Vec<(Coin, Note, jubjub::Fr, IncrementalWitness<MerkleNode>)>,
+
+    // Mint verifying key used by ZK
+    mint_pvk: groth16::PreparedVerifyingKey<Bls12>,
+    // Spend verifying key used by ZK
+    spend_pvk: groth16::PreparedVerifyingKey<Bls12>,
+
+    // Public key of the cashier
+    cashier_public: jubjub::SubgroupPoint,
+    // List of all our secret keys
+    secrets: Vec<jubjub::Fr>,
+}
+
+impl ProgramState for MemoryState {
+    fn is_valid_cashier_public_key(&self, public: &jubjub::SubgroupPoint) -> bool {
+        public == &self.cashier_public
+    }
+    fn is_valid_merkle(&self, merkle_root: &MerkleNode) -> bool {
+        self.merkle_roots.iter().any(|m| *m == *merkle_root)
+    }
+    fn nullifier_exists(&self, nullifier: &Nullifier) -> bool {
+        self.nullifiers.iter().any(|n| n.repr == nullifier.repr)
+    }
+
+    fn mint_pvk(&self) -> &groth16::PreparedVerifyingKey<Bls12> {
+        &self.mint_pvk
+    }
+    fn spend_pvk(&self) -> &groth16::PreparedVerifyingKey<Bls12> {
+        &self.spend_pvk
+    }
+}
+
+impl MemoryState {
+    fn apply(&mut self, mut update: StateUpdate) {
+        // Extend our list of nullifiers with the ones from the update
+        self.nullifiers.append(&mut update.nullifiers);
+
+        // Update merkle tree and witnesses
+        for (coin, enc_note) in update.coins.into_iter().zip(update.enc_notes.into_iter()) {
+            // Add the new coins to the merkle tree
+            let node = MerkleNode::from_coin(&coin);
+            self.tree.append(node).expect("Append to merkle tree");
+
+            // Keep track of all merkle roots that have existed
+            self.merkle_roots.push(self.tree.root());
+
+            // Also update all the coin witnesses
+            for (_, _, _, witness) in self.own_coins.iter_mut() {
+                witness.append(node).expect("append to witness");
+            }
+
+            if let Some((note, secret)) = self.try_decrypt_note(enc_note) {
+                // We need to keep track of the witness for this coin.
+                // This allows us to prove inclusion of the coin in the merkle tree with ZK.
+                // Just as we update the merkle tree with every new coin, so we do the same with
+                // the witness.
+
+                // Derive the current witness from the current tree.
+                // This is done right after we add our coin to the tree (but before any other
+                // coins are added)
+
+                // Make a new witness for this coin
+                let witness = IncrementalWitness::from_tree(&self.tree);
+                self.own_coins.push((coin, note, secret, witness));
+            }
+        }
+    }
+
+    fn try_decrypt_note(&self, ciphertext: EncryptedNote) -> Option<(Note, jubjub::Fr)> {
+        // Loop through all our secret keys...
+        for secret in &self.secrets {
+            // ... attempt to decrypt the note ...
+            match ciphertext.decrypt(secret) {
+                Ok(note) => {
+                    // ... and return the decrypted note for this coin.
+                    return Some((note, secret.clone()));
+                }
+                Err(_) => {}
+            }
+        }
+        // We weren't able to decrypt the note with any of our keys.
+        None
+    }
+}
+
+#[async_std::main]
+async fn main() {
+    // Auto create trusted ceremony parameters if they don't exist
+    if !Path::new("mint.params").exists() {
+        let params = setup_mint_prover();
+        save_params("mint.params", &params).unwrap();
+    }
+    if !Path::new("spend.params").exists() {
+        let params = setup_spend_prover();
+        save_params("spend.params", &params).unwrap();
+    }
+
+    // Load trusted setup parameters
+    let (mint_params, mint_pvk) = load_params("mint.params").expect("params should load");
+    let (spend_params, spend_pvk) = load_params("spend.params").expect("params should load");
+
+    // Cashier creates a secret key
+    let cashier_secret = jubjub::Fr::random(&mut OsRng);
+    // This is their public key
+    let cashier_public = zcash_primitives::constants::SPENDING_KEY_GENERATOR * cashier_secret;
+
+    // Wallet 1 creates a secret key
+    let secret = jubjub::Fr::random(&mut OsRng);
+    // This is their public key
+    let public = zcash_primitives::constants::SPENDING_KEY_GENERATOR * secret;
+
+    let mut state = MemoryState {
+        tree: CommitmentTree::empty(),
+        merkle_roots: vec![],
+        nullifiers: vec![],
+        own_coins: vec![],
+        mint_pvk,
+        spend_pvk,
+        cashier_public,
+        secrets: vec![secret.clone()],
+    };
+
+    // Step 1: Cashier deposits to wallet1's address
+
+    // Create the deposit for 110 BTC
+    // Clear inputs are visible to everyone on the network
+
+    let asset_id = jubjub::Fr::random(&mut OsRng);
+    let builder = tx::TransactionBuilder {
+        clear_inputs: vec![tx::TransactionBuilderClearInputInfo {
+            value: 110,
+            asset_id: asset_id,
+            signature_secret: cashier_secret,
+        }],
+        inputs: vec![],
+        outputs: vec![tx::TransactionBuilderOutputInfo {
+            value: 110,
+            asset_id: asset_id,
+            public,
+        }],
+    };
+
+    // We will 'compile' the tx, and then serialize it to this Vec<u8>
+    let mut tx_data = vec![];
+    {
+        // Build the tx
+        let tx = builder.build(&mint_params, &spend_params);
+        // Now serialize it
+        tx.encode(&mut tx_data).expect("encode tx");
+    }
+
+    // Step 1 is completed.
+    // Tx data is posted to the blockchain
+
+    // Step 2: wallet1 receive's payment from the cashier
+
+    // Wallet1 is receiving tx, and for every new coin it finds, it adds to its
+    // merkle tree
+    {
+        // Here we simulate 5 fake random coins, adding them to our tree.
+        let tree = &mut state.tree;
+        for _i in 0..5 {
+            // Don't worry about any of the code in this block
+            // We're just filling the tree with fake coins
+            let cmu = MerkleNode::new(bls12_381::Scalar::random(&mut OsRng).to_repr());
+            tree.append(cmu).unwrap();
+
+            let root = tree.root();
+            state.merkle_roots.push(root.into());
+        }
+    }
+
+    let state = Mutex::new(state);
+    let mut state = state.lock().await;
+    // Now we receive the tx data
+    {
+        let tx = tx::Transaction::decode(&tx_data[..]).unwrap();
+
+        let update = state_transition(&state, tx).expect("step 2 state transition failed");
+        // Our state impl is memory online for this demo
+        // but in the real version, this function will be async
+        // and using the databases.
+        state.apply(update);
+    }
+
+    //// Wallet1 has received payment from the cashier.
+    //// Step 2 is complete.
+    assert_eq!(state.own_coins.len(), 1);
+    ////let (coin, note, secret, witness) = &mut state.own_coins[0];
+
+    //let merkle_path = {
+    //    let tree = &mut state.tree;
+    //    let (coin, _, _, witness) = &mut state.own_coins[0];
+    //    // Check this is the 6th coin we added
+    //    assert_eq!(witness.position(), 5);
+    //    assert_eq!(tree.root(), witness.root());
+
+    //    // Add some more random coins in
+    //    for _i in 0..10 {
+    //        // Don't worry about any of the code in this block
+    //        // We're just filling the tree with fake coins
+    //        let cmu = MerkleNode::new(bls12_381::Scalar::random(&mut OsRng).to_repr());
+    //        tree.append(cmu).unwrap();
+    //        witness.append(cmu).unwrap();
+    //        assert_eq!(tree.root(), witness.root());
+
+    //        let root = tree.root();
+    //        state.merkle_roots.push(root.into());
+    //    }
+
+    //    assert_eq!(state.merkle_roots.len(), 16);
+
+    //    // This is the value we need to spend the coin
+    //    // We use the witness and the merkle root (both in sync with each other)
+    //    // to prove our coin exists inside the tree.
+    //    // The coin is not revealed publicly but is proved to exist inside
+    //    // a merkle tree. Only the root will be revealed, and then the
+    //    // verifier checks that merkle root actually existed before.
+    //    let merkle_path = witness.path().unwrap();
+
+    //    // Just test the path is good because we just added a bunch of fake coins
+    //    let node = MerkleNode::from_coin(&coin);
+    //    let root = tree.root();
+    //    drop(tree);
+    //    drop(witness);
+    //    assert_eq!(merkle_path.root(node), root);
+    //    let root = root.into();
+    //    assert!(state.is_valid_merkle(&root));
+
+    //    merkle_path
+    //};
+
+    // Step 3: wallet1 sends payment to wallet2
+
+    // Wallet1 now wishes to send the coin to wallet2
+
+    // The receiving wallet has a secret key
+    let secret2 = jubjub::Fr::random(&mut OsRng);
+    // This is their public key to receive payment
+    let public2 = zcash_primitives::constants::SPENDING_KEY_GENERATOR * secret2;
+
+    // Make a spend tx
+
+    //let inputs: Vec<tx::TransactionBuilderInputInfo> = vec![];
+    // Construct a new tx spending the coin
+    // We need the decrypted note and our private key
+    let builder = tx::TransactionBuilder {
+        clear_inputs: vec![],
+        inputs: vec![],
+        //inputs: vec![tx::TransactionBuilderInputInfo {
+        //    merkle_path,
+        //    secret: secret.clone(),
+        //    note: state.own_coins[0].1.clone(),
+        //}],
+        // We can add more outputs to this list.
+        // The only constraint is that sum(value in) == sum(value out)
+        outputs: vec![tx::TransactionBuilderOutputInfo {
+            value: 110,
+            asset_id: asset_id,
+            public: public2,
+        }],
+    };
+    // Build the tx
+    let mut tx_data = vec![];
+    {
+        let tx = builder.build(&mint_params, &spend_params);
+        tx.encode(&mut tx_data).expect("encode tx");
+    }
+    // Verify it's valid
+    {
+        let tx = tx::Transaction::decode(&tx_data[..]).unwrap();
+        //let state = state.lock().await;
+        //let update = state_transition(&state, tx).expect("step 3 state transition failed");
+        //state.apply(update);
+    }
+}