feat: adding delegated voting using notes

src/main.nr

@@ -4,16 +4,72 @@ use dep::aztec::macros::aztec;
 #[aztec]
 contract EasyPrivateVoting {
     use dep::aztec::{
-        keys::getters::get_public_keys,
-        macros::{functions::{initializer, internal, private, public}, storage::storage},
+        context::PrivateContext,
+        keys::getters::{get_nsk_app, get_public_keys},
+        macros::{functions::{initializer, internal, private, public}, storage::storage, notes::note},
+        note::{note_interface::NullifiableNote, utils::compute_note_hash_for_nullify},
+        prelude::{AztecAddress, Map, PublicImmutable, PublicMutable, PrivateSet, NoteHeader, NoteGetterOptions},
+        protocol_types::{hash::poseidon2_hash_with_separator, constants::GENERATOR_INDEX__NOTE_NULLIFIER},
+        encrypted_logs::encrypted_note_emission::encode_and_encrypt_note,
+        utils::comparison::Comparator,
     };
-    use dep::aztec::prelude::{AztecAddress, Map, PublicImmutable, PublicMutable};
+
+
     #[storage]
     struct Storage<Context> {
         admin: PublicMutable<AztecAddress, Context>, // admin can end vote
         tally: Map<Field, PublicMutable<Field, Context>, Context>, // we will store candidate as key and number of votes as value
         vote_ended: PublicMutable<bool, Context>, // vote_ended is boolean
         active_at_block: PublicImmutable<u32, Context>, // when people can start voting
+        delegations: PrivateSet<DelegateNote, Context>, // set of delegations
+    }
+
+    #[note]
+    pub struct DelegateNote {
+        delegator: AztecAddress,
+        delegatee: AztecAddress,
+        randomness: Field, // Added randomness to prevent privacy leaks
+    }
+
+    impl DelegateNote {
+        pub fn new(delegator: AztecAddress, delegatee: AztecAddress, randomness: Field) -> Self {
+            DelegateNote { delegator, delegatee, randomness, header: NoteHeader::empty() }
+        }
+    }
+
+    impl Eq for DelegateNote {
+        fn eq(self, other: Self) -> bool {
+            (self.delegator == other.delegator)
+                & (self.delegatee == other.delegatee)
+                // NOTE: we don't check for randomness in Eq (1 delegation per delegator)
+                // & (self.randomness == other.randomness) 
+        }
+    }
+
+    impl NullifiableNote for DelegateNote {
+        fn compute_nullifier(
+            self,
+            context: &mut PrivateContext,
+            note_hash_for_nullify: Field,
+        ) -> Field {
+            // the delegatee is the owner of the note (the delegator is used for desambiguation)
+            let owner_npk_m_hash = get_public_keys(self.delegatee).npk_m.hash();
+            let secret = context.request_nsk_app(owner_npk_m_hash);
+            poseidon2_hash_with_separator(
+                [note_hash_for_nullify, secret],
+                GENERATOR_INDEX__NOTE_NULLIFIER as Field,
+            )
+        }
+
+        unconstrained fn compute_nullifier_without_context(self) -> Field {
+            let note_hash_for_nullify = compute_note_hash_for_nullify(self);
+            let owner_npk_m_hash = get_public_keys(self.delegatee).npk_m.hash();
+            let secret = get_nsk_app(owner_npk_m_hash);
+            poseidon2_hash_with_separator(
+                [note_hash_for_nullify, secret],
+                GENERATOR_INDEX__NOTE_NULLIFIER as Field,
+            )
+        }
     }
 
     #[public]
@@ -33,16 +89,50 @@ contract EasyPrivateVoting {
         let secret = context.request_nsk_app(msg_sender_npk_m_hash); // get secret key of caller of function
         let nullifier = std::hash::pedersen_hash([context.msg_sender().to_field(), secret]); // derive nullifier from sender and secret
         context.push_nullifier(nullifier);
-        EasyPrivateVoting::at(context.this_address()).add_to_tally_public(candidate).enqueue(
+        EasyPrivateVoting::at(context.this_address()).add_to_tally_public(candidate, 1).enqueue(
             &mut context,
         );
     }
 
+    #[private]
+    // sample method for creating a note for someone else to vote on your behalf
+    fn delegate_vote(delegatee: AztecAddress, randomness: Field) {
+        let msg_sender_npk_m_hash = get_public_keys(context.msg_sender()).npk_m.hash();
+
+        let secret = context.request_nsk_app(msg_sender_npk_m_hash); // get secret key of caller of function
+        let nullifier = std::hash::pedersen_hash([context.msg_sender().to_field(), secret]); // derive nullifier from sender and secret
+        context.push_nullifier(nullifier);
+
+        let mut delegation = DelegateNote::new(context.msg_sender(), delegatee, randomness);
+        storage.delegations.insert(&mut delegation).emit(encode_and_encrypt_note(
+            &mut context,
+            delegatee,
+            context.msg_sender(),
+        ));
+    }
+
+    #[private]
+    fn cast_delegated_vote(candidate: Field) {
+        let mut options = NoteGetterOptions::new();
+        options = options
+            .select(DelegateNote::properties().delegatee, Comparator.EQ, context.msg_sender())
+            .set_limit(3); // low limit for testing purposes
+        let popped_notes = storage.delegations.pop_notes(options);
+        let num_delegated_votes = popped_notes.len() as Field;
+
+        EasyPrivateVoting::at(context.this_address()).add_to_tally_public(candidate, num_delegated_votes).enqueue(
+            &mut context,
+        );
+    }
+
+
     #[public]
     #[internal]
-    fn add_to_tally_public(candidate: Field) {
+    fn add_to_tally_public(candidate: Field, amount: Field) {
         assert(storage.vote_ended.read() == false, "Vote has ended"); // assert that vote has not ended
-        let new_tally = storage.tally.at(candidate).read() + 1;
+        let num_amount = amount;
+
+        let new_tally = storage.tally.at(candidate).read() + num_amount;
         storage.tally.at(candidate).write(new_tally);
     }
 

src/test/first.nr

@@ -94,6 +94,83 @@ unconstrained fn test_cast_vote_with_separate_accounts() {
     assert(vote_count == 2, "vote tally should be 2");
 }
 
+#[test]
+unconstrained fn test_cast_vote_with_delegation() {
+    let (env, voting_contract_address, _) = utils::setup();
+    let alice = env.create_account();
+    let bob = env.create_account();
+
+    let candidate = 101;
+    let random = 420;
+
+    env.impersonate(alice);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).delegate_vote(bob, random).call(&mut env.private());
+
+    env.impersonate(bob);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).cast_delegated_vote(candidate).call(&mut env.private());
+
+    // Read vote count from storage
+    let block_number = get_block_number();
+    let tally_slot = EasyPrivateVoting::storage_layout().tally.slot;
+    let candidate_tally_slot = derive_storage_slot_in_map(tally_slot, candidate);
+    let vote_count: u32 = storage_read(voting_contract_address, candidate_tally_slot, block_number);
+
+    assert(vote_count == 1, "vote tally should be 1");
+}
+
+#[test]
+unconstrained fn test_cast_vote_with_multiple_delegations() {
+    let (env, voting_contract_address, _) = utils::setup();
+    let alice = env.create_account();
+    let bob = env.create_account();
+    let carl = env.create_account();
+    let dave = env.create_account();
+    let delegatee = env.create_account();
+
+    let candidate = 101;
+    let random = 420;
+
+    env.impersonate(alice);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).delegate_vote(delegatee, random).call(&mut env.private());
+
+    env.impersonate(bob);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).delegate_vote(delegatee, random).call(&mut env.private());
+
+    env.impersonate(carl);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).delegate_vote(delegatee, random).call(&mut env.private());
+
+    // NOTE: the limit of delegations per `cast_delegated_vote` is set to 3
+    env.impersonate(dave);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).delegate_vote(delegatee, random).call(&mut env.private());
+
+    env.impersonate(delegatee);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).cast_delegated_vote(candidate).call(&mut env.private());
+
+    // Read vote count from storage
+    let mut block_number = get_block_number();
+    let tally_slot = EasyPrivateVoting::storage_layout().tally.slot;
+    let candidate_tally_slot = derive_storage_slot_in_map(tally_slot, candidate);
+    let mut vote_count: u32 = storage_read(voting_contract_address, candidate_tally_slot, block_number);
+
+    assert(vote_count == 3, "vote tally should be 3");
+
+    env.impersonate(delegatee);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).cast_delegated_vote(candidate).call(&mut env.private());
+
+    block_number = get_block_number();
+    vote_count = storage_read(voting_contract_address, candidate_tally_slot, block_number);
+
+    assert(vote_count == 4, "vote tally should be 4");
+}
+
 #[test(should_fail)]
 unconstrained fn test_fail_vote_twice() {
     let (env, voting_contract_address, _) = utils::setup();
@@ -109,3 +186,39 @@ unconstrained fn test_fail_vote_twice() {
     env.advance_block_by(1);
     EasyPrivateVoting::at(voting_contract_address).cast_vote(candidate).call(&mut env.private());
 }
+
+#[test(should_fail)]
+unconstrained fn test_fail_delegate_and_vote() {
+    let (env, voting_contract_address, _) = utils::setup();
+    let alice = env.create_account();
+    let bob = env.create_account();
+
+    let candidate = 101;
+    let random = 420;
+
+    env.impersonate(alice);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).delegate_vote(bob, random).call(&mut env.private());
+
+    // Vote again as alice
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).cast_vote(candidate).call(&mut env.private());
+}
+
+#[test(should_fail)]
+unconstrained fn test_fail_vote_and_delegate() {
+    let (env, voting_contract_address, _) = utils::setup();
+    let alice = env.create_account();
+    let bob = env.create_account();
+
+    let candidate = 101;
+    let random = 420;
+
+    env.impersonate(alice);
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).cast_vote(candidate).call(&mut env.private());
+
+    // Delegate vote as alice
+    env.advance_block_by(1);
+    EasyPrivateVoting::at(voting_contract_address).delegate_vote(bob, random).call(&mut env.private());
+}

src/test/index.test.ts

@@ -4,6 +4,11 @@ import { getInitialTestAccountsWallets } from "@aztec/accounts/testing"
 
 const setupSandbox = async () => {
     const { PXE_URL = 'http://localhost:8080' } = process.env;
+    // TODO: implement reading the DelegationNote from an isolated PXE
+    // 8080: cd ~/.aztec && docker-compose -f ./docker-compose.sandbox.yml up
+    // 8081: aztec start --port 8081 --pxe --pxe.nodeUrl http://host.docker.internal:8080/
+    // const DELEGATEE_PXE_URL = 'http://localhost:8081';
+
     const pxe = createPXEClient(PXE_URL);
     await waitForPXE(pxe);
     return pxe;
@@ -20,6 +25,7 @@ describe("Voting", () => {
         logger.info("Aztec-Starter tests running.")
 
         pxe = await setupSandbox();
+        // deployInitialTestAccounts(pxe); // NOTE: run at least once in sandbox to circumvent issue #9384
 
         wallets = await getInitialTestAccountsWallets(pxe);
         accounts = wallets.map(w => w.getCompleteAddress())
@@ -86,4 +92,34 @@ describe("Voting", () => {
 
     }, 300_000)
 
+    it("It casts a delegated vote", async () => {
+        const candidate = new Fr(1)
+        const delegatee = accounts[1].address
+        const random = new Fr(2)
+
+        const contract = await EasyPrivateVotingContract.deploy(wallets[0], accounts[0].address).send().deployed();
+        await contract.methods.delegate_vote(delegatee, random).send().wait();
+
+        const tx = await contract.withWallet(wallets[1]).methods.cast_delegated_vote(candidate).send().wait();
+        let count = await contract.methods.get_vote(candidate).simulate();
+        expect(count).toBe(1n);
+    }, 300_000)
+
+    it("It should fail when trying to both delegate and vote", async () => {
+        const candidate = new Fr(1)
+        const delegatee = accounts[1].address
+        const random = new Fr(2)
+
+        const contract = await EasyPrivateVotingContract.deploy(wallets[0], accounts[0].address).send().deployed();
+        await contract.methods.delegate_vote(delegatee, random).send().wait();
+
+        // We try voting again, but our TX is dropped due to trying to emit duplicate nullifiers
+        // first confirm that it fails simulation
+        await expect(contract.methods.cast_vote(candidate).send().wait()).rejects.toThrow(/Nullifier collision/);
+        // if we skip simulation, tx is dropped
+        await expect(
+            contract.methods.cast_vote(candidate).send({ skipPublicSimulation: true }).wait(),
+        ).rejects.toThrow('Reason: Tx dropped by P2P node.');
+    }, 300_000)
+
 });

src/test/utils.nr

@@ -8,9 +8,9 @@ use dep::aztec::{
 use crate::EasyPrivateVoting;
 
 pub fn setup() -> (&mut TestEnvironment, AztecAddress, AztecAddress) {
-    let mut env = TestEnvironment::new();
+    let mut env = unsafe { TestEnvironment::new() };
 
-    let admin = env.create_account();
+    let admin = unsafe { env.create_account() };
 
     let initializer_call_interface = EasyPrivateVoting::interface().constructor(admin);
     let voting_contract = env.deploy_self("EasyPrivateVoting").with_public_void_initializer(