New shielder smart contract (#42)

shielder/contract/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "shielder-contract"
+version = "0.1.0"
+authors = ["Cardinal"]
+homepage = "https://alephzero.org"
+repository = "https://github.com/Cardinal-Cryptography/zk-apps"
+edition = "2021"
+
+[dependencies]
+ink = { version = "5.0.0-rc", default-features = false }
+
+[dev-dependencies]
+
+[lib]
+path = "lib.rs"
+
+[features]
+default = ["std"]
+std = [
+    "ink/std",
+]
+ink-as-dependency = []
+e2e-tests = []

shielder/contract/errors.rs

@@ -0,0 +1,19 @@
+use crate::traits::psp22::PSP22Error;
+
+#[ink::scale_derive(Encode, Decode, TypeInfo)]
+#[derive(PartialEq, Debug)]
+pub enum ShielderError {
+    PSP22(PSP22Error),
+    NullifierIsInSet,
+    MerkleTreeVerificationFail,
+    MerkleTreeLimitExceeded,
+    MerkleTreeProofGenFail,
+    ZkpVerificationFail,
+    ArithmeticError,
+}
+
+impl From<PSP22Error> for ShielderError {
+    fn from(inner: PSP22Error) -> Self {
+        ShielderError::PSP22(inner)
+    }
+}

shielder/contract/lib.rs

@@ -0,0 +1,125 @@
+//! Smart contract implementing shielder specification
+//! https://docs.alephzero.org/aleph-zero/shielder/introduction-informal
+
+#![cfg_attr(not(feature = "std"), no_std, no_main)]
+#![deny(missing_docs)]
+
+mod errors;
+mod merkle;
+mod mocked_zk;
+mod traits;
+mod types;
+
+/// Contract module
+#[ink::contract]
+pub mod contract {
+
+    use crate::{
+        errors::ShielderError,
+        merkle::MerkleTree,
+        mocked_zk::relations::ZkProof,
+        traits::psp22::PSP22,
+        types::{Scalar, Set},
+    };
+
+    /// Enum
+    #[ink::scale_derive(Encode, Decode, TypeInfo)]
+    #[derive(Clone, Copy)]
+    pub enum OpPub {
+        /// Deposit PSP-22 token
+        Deposit {
+            /// amount of deposit
+            amount: u128,
+            /// PSP-22 token address
+            token: AccountId,
+            /// User address, from whom tokens are transferred
+            user: AccountId,
+        },
+        /// Withdraw PSP-22 token
+        Withdraw {
+            /// amount of withdrawal
+            amount: u128,
+            /// PSP-22 token address
+            token: AccountId,
+            /// User address, from whom tokens are transferred
+            user: AccountId,
+        },
+    }
+
+    /// Contract storage
+    #[ink(storage)]
+    #[derive(Default)]
+    pub struct Contract {
+        nullifier_set: Set<Scalar>,
+        notes: MerkleTree,
+    }
+
+    impl Contract {
+        /// Constructor
+        #[ink(constructor)]
+        pub fn new() -> Self {
+            Self::default()
+        }
+
+        /// Adds empty note to shielder storage
+        /// Registers new account with empty balance
+        #[ink(message)]
+        pub fn add_note(
+            &mut self,
+            h_note_new: Scalar,
+            proof: ZkProof,
+        ) -> Result<(), ShielderError> {
+            proof.verify_creation(h_note_new)?;
+            self.notes.add_leaf(h_note_new)?;
+            Ok(())
+        }
+
+        /// Updates existing note
+        /// Applies operation to private account stored in shielder
+        #[ink(message)]
+        pub fn update_note(
+            &mut self,
+            op_pub: OpPub,
+            h_note_new: Scalar,
+            merkle_root: Scalar,
+            nullifier_old: Scalar,
+            proof: ZkProof,
+        ) -> Result<(), ShielderError> {
+            self.notes.is_historical_root(merkle_root)?;
+            self.nullify(nullifier_old)?;
+            proof.verify_update(op_pub, h_note_new, merkle_root, nullifier_old)?;
+            self.notes.add_leaf(h_note_new)?;
+            self.process_operation(op_pub)?;
+            Ok(())
+        }
+
+        fn process_operation(&mut self, op_pub: OpPub) -> Result<(), ShielderError> {
+            match op_pub {
+                OpPub::Deposit {
+                    amount,
+                    token,
+                    user,
+                } => {
+                    let mut psp22: ink::contract_ref!(PSP22) = token.into();
+                    psp22.transfer_from(user, self.env().account_id(), amount, [].to_vec())?;
+                }
+                OpPub::Withdraw {
+                    amount,
+                    token,
+                    user,
+                } => {
+                    let mut psp22: ink::contract_ref!(PSP22) = token.into();
+                    psp22.transfer(user, amount, [].to_vec())?;
+                }
+            };
+            Ok(())
+        }
+
+        fn nullify(&mut self, nullifier: Scalar) -> Result<(), ShielderError> {
+            self.nullifier_set
+                .insert(nullifier, &())
+                .map(|_| {})
+                .ok_or(ShielderError::NullifierIsInSet)
+        }
+    }
+}

shielder/contract/merkle.rs

@@ -0,0 +1,81 @@
+use ink::{
+    env::hash::{CryptoHash, Sha2x256},
+    storage::Mapping,
+};
+
+use crate::{
+    errors::ShielderError,
+    types::{Scalar, Set},
+};
+
+/// depth of the tree
+pub const DEPTH: usize = 10;
+
+#[ink::storage_item]
+#[derive(Debug)]
+pub struct MerkleTree {
+    /// mapping of tree indexes to values held in nodes
+    nodes: Mapping<u32, Scalar>,
+    /// set of historical roots (nodes[1]) of tree
+    roots_log: Set<Scalar>,
+    /// index of next available leaf
+    next_leaf_idx: u32,
+    /// number of leaves in the tree, should be equal to 2^DEPTH
+    size: u32,
+}
+
+pub fn compute_hash(first: Scalar, second: Scalar) -> Scalar {
+    let mut res = [0x0; 32];
+    Sha2x256::hash([first.bytes, second.bytes].concat().as_slice(), &mut res);
+    Scalar { bytes: res }
+}
+
+impl Default for MerkleTree {
+    fn default() -> Self {
+        Self {
+            nodes: Default::default(),
+            roots_log: Default::default(),
+            next_leaf_idx: 0,
+            size: (1 << DEPTH),
+        }
+    }
+}
+
+impl MerkleTree {
+    pub fn add_leaf(&mut self, leaf_value: Scalar) -> Result<(), ShielderError> {
+        if self.next_leaf_idx == self.size {
+            return Err(ShielderError::MerkleTreeLimitExceeded);
+        }
+        let mut id = self
+            .next_leaf_idx
+            .checked_add(self.size)
+            .ok_or(ShielderError::ArithmeticError)?;
+        self.nodes.insert(id, &leaf_value);
+
+        id /= 2;
+        while id > 0 {
+            let id_mul_2 = id.checked_mul(2).ok_or(ShielderError::ArithmeticError)?;
+            let left_n = self.node_value(id_mul_2);
+            let right_n = self.node_value(id.checked_add(1).ok_or(ShielderError::ArithmeticError)?);
+            let hash = compute_hash(left_n, right_n);
+            self.nodes.insert(id, &hash);
+            id /= 2;
+        }
+        self.next_leaf_idx = self
+            .next_leaf_idx
+            .checked_add(1)
+            .ok_or(ShielderError::ArithmeticError)?;
+        Ok(())
+    }
+
+    pub fn is_historical_root(&self, merkle_root_possible: Scalar) -> Result<(), ShielderError> {
+        self.roots_log
+            .contains(merkle_root_possible)
+            .then_some(())
+            .ok_or(ShielderError::MerkleTreeVerificationFail)
+    }
+
+    fn node_value(&self, id: u32) -> Scalar {
+        self.nodes.get(id).unwrap_or_default()
+    }
+}

shielder/contract/mocked_zk/account.rs

@@ -0,0 +1,63 @@
+use ink::env::hash::{CryptoHash, Sha2x256};
+
+use super::{ops::Operation, traits::Hashable, USDT_TOKEN};
+use crate::{contract::OpPub, errors::ShielderError, types::Scalar};
+
+#[ink::scale_derive(Encode, Decode, TypeInfo)]
+#[derive(Default, Clone, Copy)]
+pub struct Account {
+    balance_aleph: Scalar,
+    balance_usdt: Scalar,
+}
+
+impl Hashable for Account {
+    fn hash(&self) -> Scalar {
+        let mut res = [0x0; 32];
+        Sha2x256::hash(
+            [self.balance_aleph.bytes, self.balance_usdt.bytes]
+                .concat()
+                .as_slice(),
+            &mut res,
+        );
+        Scalar { bytes: res }
+    }
+}
+
+impl Account {
+    pub fn new() -> Self {
+        Self {
+            balance_aleph: 0_u128.into(),
+            balance_usdt: 0_u128.into(),
+        }
+    }
+    pub fn update(&self, operation: Operation) -> Result<Self, ShielderError> {
+        match operation.op_pub {
+            OpPub::Deposit { amount, token, .. } => {
+                let mut balance_usdt = self.balance_usdt;
+                if token.as_ref() == USDT_TOKEN {
+                    balance_usdt = (u128::from(balance_usdt)
+                        .checked_add(amount)
+                        .ok_or(ShielderError::ArithmeticError)?)
+                    .into();
+                }
+                Ok(Self {
+                    balance_aleph: self.balance_aleph,
+                    balance_usdt,
+                })
+            }
+            OpPub::Withdraw { amount, token, .. } => {
+                let mut balance_usdt = self.balance_usdt;
+                if token.as_ref() == USDT_TOKEN {
+                    balance_usdt = (u128::from(balance_usdt)
+                        .checked_sub(amount)
+                        .ok_or(ShielderError::ArithmeticError)?)
+                    .into();
+                }
+                Ok(Self {
+                    balance_aleph: self.balance_aleph,
+                    balance_usdt,
+                })
+            }
+        }
+    }
+}

shielder/contract/mocked_zk/mod.rs

@@ -0,0 +1,9 @@
+mod account;
+mod note;
+mod ops;
+pub mod relations;
+#[cfg(test)]
+mod tests;
+mod traits;
+
+const USDT_TOKEN: [u8; 32] = [0x2_u8; 32];

shielder/contract/mocked_zk/note.rs

@@ -0,0 +1,42 @@
+use ink::env::hash::{CryptoHash, Sha2x256};
+
+use super::traits::Hashable;
+use crate::types::Scalar;
+
+#[ink::scale_derive(Encode, Decode, TypeInfo)]
+#[derive(Clone, Copy)]
+pub struct Note {
+    id: Scalar,
+    trapdoor: Scalar,
+    nullifier: Scalar,
+    account_hash: Scalar,
+}
+
+impl Note {
+    pub fn new(id: Scalar, trapdoor: Scalar, nullifier: Scalar, account_hash: Scalar) -> Self {
+        Self {
+            id,
+            trapdoor,
+            nullifier,
+            account_hash,
+        }
+    }
+}
+
+impl Hashable for Note {
+    fn hash(&self) -> Scalar {
+        let mut res = [0x0; 32];
+        Sha2x256::hash(
+            [
+                self.id.bytes,
+                self.trapdoor.bytes,
+                self.nullifier.bytes,
+                self.account_hash.bytes,
+            ]
+            .concat()
+            .as_slice(),
+            &mut res,
+        );
+        Scalar { bytes: res }
+    }
+}