diff --git a/noir-projects/noir-contracts/contracts/token_contract/src/main.nr b/noir-projects/noir-contracts/contracts/token_contract/src/main.nr index 253a39a4b4e..2c1e8b917cb 100644 --- a/noir-projects/noir-contracts/contracts/token_contract/src/main.nr +++ b/noir-projects/noir-contracts/contracts/token_contract/src/main.nr @@ -225,7 +225,8 @@ contract Token { #[aztec(private)] fn privately_mint_private_note(amount: Field) { let caller = context.msg_sender(); - storage.balances.add(caller, U128::from_integer(amount)).emit(encode_and_encrypt_note(&mut context, caller, caller)); + let caller_npk_m = get_current_public_keys(&mut context, caller).npk_m; + storage.balances.add(caller, caller_npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note(&mut context, caller, caller)); Token::at(context.this_address()).assert_minter_and_mint(context.msg_sender(), amount).enqueue(&mut context); } @@ -318,7 +319,8 @@ contract Token { // Note: Using context.msg_sender() as a sender below makes this incompatible with escrows because we send // outgoing logs to that address and to send outgoing logs you need to get a hold of ovsk_m. let from = context.msg_sender(); - storage.balances.add(to, U128::from_integer(amount)).emit(encode_and_encrypt_note(&mut context, from, to)); + let to_npk_m = get_current_public_keys(&mut context, to).npk_m; + storage.balances.add(to, to_npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note(&mut context, from, to)); } // docs:end:redeem_shield @@ -331,7 +333,8 @@ contract Token { assert(nonce == 0, "invalid nonce"); } - storage.balances.sub(from, U128::from_integer(amount)).emit(encode_and_encrypt_note(&mut context, from, from)); + let from_npk_m = get_current_public_keys(&mut context, from).npk_m; + storage.balances.sub(from, from_npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note(&mut context, from, from)); Token::at(context.this_address())._increase_public_balance(to, amount).enqueue(&mut context); } @@ -361,11 +364,11 @@ contract Token { INITIAL_TRANSFER_CALL_MAX_NOTES ); - storage.balances.add(from, change).emit( + storage.balances.add(from, from_keys.npk_m, change).emit( encode_and_encrypt_note_with_keys_unconstrained(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from) ); - storage.balances.add(to, amount).emit( + storage.balances.add(to, to_keys.npk_m, amount).emit( encode_and_encrypt_note_with_keys_unconstrained(&mut context, from_keys.ovpk_m, to_keys.ivpk_m, to) ); @@ -462,10 +465,10 @@ contract Token { let amount = U128::from_integer(amount); // docs:start:increase_private_balance // docs:start:encrypted - storage.balances.sub(from, amount).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from)); + storage.balances.sub(from, from_keys.npk_m, amount).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from)); // docs:end:encrypted // docs:end:increase_private_balance - storage.balances.add(to, amount).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, to_keys.ivpk_m, to)); + storage.balances.add(to, to_keys.npk_m, amount).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, to_keys.ivpk_m, to)); } // docs:end:transfer_from @@ -478,7 +481,8 @@ contract Token { assert(nonce == 0, "invalid nonce"); } - storage.balances.sub(from, U128::from_integer(amount)).emit(encode_and_encrypt_note(&mut context, from, from)); + let from_npk_m = get_current_public_keys(&mut context, from).npk_m; + storage.balances.sub(from, from_npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note(&mut context, from, from)); Token::at(context.this_address())._reduce_total_supply(amount).enqueue(&mut context); } @@ -528,7 +532,7 @@ contract Token { U128::from_integer(funded_amount), INITIAL_TRANSFER_CALL_MAX_NOTES ); - storage.balances.add(user, change).emit( + storage.balances.add(user, user_keys.npk_m, change).emit( encode_and_encrypt_note_with_keys_unconstrained(&mut context, user_keys.ovpk_m, user_keys.ivpk_m, user) ); diff --git a/noir-projects/noir-contracts/contracts/token_contract/src/types/balances_map.nr b/noir-projects/noir-contracts/contracts/token_contract/src/types/balances_map.nr index 0fece696437..af1c9affa9c 100644 --- a/noir-projects/noir-contracts/contracts/token_contract/src/types/balances_map.nr +++ b/noir-projects/noir-contracts/contracts/token_contract/src/types/balances_map.nr @@ -1,4 +1,4 @@ -use dep::aztec::prelude::{AztecAddress, NoteGetterOptions, NoteViewerOptions, NoteHeader, NoteInterface, PrivateSet, Map}; +use dep::aztec::prelude::{AztecAddress, NoteGetterOptions, NoteViewerOptions, NoteHeader, NoteInterface, PrivateSet, Map, Point}; use dep::aztec::{ context::{PrivateContext, UnconstrainedContext}, protocol_types::constants::MAX_NOTE_HASH_READ_REQUESTS_PER_CALL, @@ -62,6 +62,7 @@ impl BalancesMap { pub fn add( self: Self, owner: AztecAddress, + owner_npk_m: Point, addend: U128 ) -> OuterNoteEmission where T: NoteInterface + OwnedNote + Eq { if addend == U128::from_integer(0) { @@ -70,8 +71,7 @@ impl BalancesMap { let context = self.map.context; // We fetch the nullifier public key hash from the registry / from our PXE - let owner_npk_m_hash = get_current_public_keys(context, owner).npk_m.hash(); - let mut addend_note = T::new(addend, owner_npk_m_hash); + let mut addend_note = T::new(addend, owner_npk_m.hash()); // docs:start:insert OuterNoteEmission::new(Option::some(self.map.at(owner).insert(&mut addend_note))) @@ -82,6 +82,7 @@ impl BalancesMap { pub fn sub( self: Self, owner: AztecAddress, + owner_npk_m: Point, amount: U128 ) -> OuterNoteEmission where T: NoteInterface + OwnedNote + Eq { let subtracted = self.try_sub(owner, amount, MAX_NOTE_HASH_READ_REQUESTS_PER_CALL); @@ -89,7 +90,7 @@ impl BalancesMap { // try_sub may have substracted more or less than amount. We must ensure that we subtracted at least as much as // we needed, and then create a new note for the owner for the change (if any). assert(subtracted >= amount, "Balance too low"); - self.add(owner, subtracted - amount) + self.add(owner, owner_npk_m, subtracted - amount) } // Attempts to remove 'target_amount' from the owner's balance. try_sub returns how much was actually subtracted