feat: fix contract issues with key rotation, introduce get npk_m_hash_at

noir-projects/aztec-nr/aztec/src/history/public_storage.nr

@@ -6,7 +6,7 @@ use dep::std::merkle::compute_merkle_root;
 
 use crate::{context::PrivateContext, oracle::get_public_data_witness::get_public_data_witness};
 
-fn _public_storage_historical_read(storage_slot: Field, contract_address: AztecAddress, header: Header) -> Field {
+pub fn _public_storage_historical_read(storage_slot: Field, contract_address: AztecAddress, header: Header) -> Field {
     // 1) Compute the leaf slot by siloing the storage slot with the contract address
     let public_value_leaf_slot = pedersen_hash(
         [contract_address.to_field(), storage_slot],

noir-projects/aztec-nr/aztec/src/keys/getters.nr

@@ -1,5 +1,5 @@
 use dep::protocol_types::{
-    abis::key_validation_request::KeyValidationRequest, address::AztecAddress,
+    header::Header, abis::key_validation_request::KeyValidationRequest, address::AztecAddress,
     constants::CANONICAL_KEY_REGISTRY_ADDRESS, grumpkin_point::GrumpkinPoint,
     storage::map::derive_storage_slot_in_map
 };
@@ -17,8 +17,28 @@ pub fn get_npk_m(context: &mut PrivateContext, address: AztecAddress) -> Grumpki
     get_master_key(context, address, NULLIFIER_INDEX)
 }
 
+// For historical access, we move the fetching of the header to the top level to minimize redundant fetching / proving validity of the header.
+pub fn get_npk_m_at(
+    context: &mut PrivateContext,
+    address: AztecAddress,
+    block_number: u32,
+    header: Header
+) -> GrumpkinPoint {
+    assert_eq(block_number, header.global_variables.block_number as u32);
+    get_master_key_at(context, address, NULLIFIER_INDEX, header)
+}
+
 pub fn get_npk_m_hash(context: &mut PrivateContext, address: AztecAddress) -> Field {
-    get_master_key(context, address, NULLIFIER_INDEX).hash()
+    get_npk_m(context, address).hash()
+}
+
+pub fn get_npk_m_hash_at(
+    context: &mut PrivateContext,
+    address: AztecAddress,
+    block_number: u32,
+    header: Header
+) -> Field {
+    get_npk_m_at(context, address, block_number, header).hash()
 }
 
 pub fn get_ivpk_m(context: &mut PrivateContext, address: AztecAddress) -> GrumpkinPoint {
@@ -40,7 +60,23 @@ fn get_master_key(
     address: AztecAddress,
     key_index: Field
 ) -> GrumpkinPoint {
-    let key = fetch_key_from_registry(context, key_index, address);
+    let (x_coordinate_registry, y_coordinate_registry) = get_registry_with_key_index(context, key_index, address);
+    let key = fetch_key_from_registry(x_coordinate_registry, y_coordinate_registry);
+    fetch_key_from_oracle_if_zero(key, address, key_index)
+}
+
+fn get_master_key_at(
+    context: &mut PrivateContext,
+    address: AztecAddress,
+    key_index: Field,
+    header: Header
+) -> GrumpkinPoint {
+    let (x_coordinate_registry, y_coordinate_registry) = get_registry_with_key_index(context, key_index, address);
+    let key = fetch_key_from_registry_at(x_coordinate_registry, y_coordinate_registry, header);
+    fetch_key_from_oracle_if_zero(key, address, key_index)
+}
+
+fn fetch_key_from_oracle_if_zero(key: GrumpkinPoint, address: AztecAddress, key_index: Field) -> GrumpkinPoint {
     if key.is_zero() {
         // Keys were not registered in registry yet --> fetch key from PXE
         let keys = fetch_and_constrain_keys(address);
@@ -52,11 +88,11 @@ fn get_master_key(
     }
 }
 
-fn fetch_key_from_registry(
+fn get_registry_with_key_index(
     context: &mut PrivateContext,
     key_index: Field,
     address: AztecAddress
-) -> GrumpkinPoint {
+) -> (SharedMutablePrivateGetter<Field, DELAY>, SharedMutablePrivateGetter<Field, DELAY>) {
     let x_coordinate_map_slot = key_index * 2 + 1;
     let y_coordinate_map_slot = x_coordinate_map_slot + 1;
     let x_coordinate_derived_slot = derive_storage_slot_in_map(x_coordinate_map_slot, address);
@@ -72,10 +108,29 @@ fn fetch_key_from_registry(
         AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS),
         y_coordinate_derived_slot
     );
-    let x_coordinate = x_coordinate_registry.get_current_value_in_private();
-    let y_coordinate = y_coordinate_registry.get_current_value_in_private();
 
-    GrumpkinPoint::new(x_coordinate, y_coordinate)
+    (x_coordinate_registry, y_coordinate_registry)
+}
+
+fn fetch_key_from_registry(
+    x_coordinate_registry: SharedMutablePrivateGetter<Field, DELAY>,
+    y_coordinate_registry: SharedMutablePrivateGetter<Field, DELAY>
+) -> GrumpkinPoint {
+    let key_x_coordinate = x_coordinate_registry.get_current_value_in_private();
+    let key_y_coordinate = y_coordinate_registry.get_current_value_in_private();
+
+    GrumpkinPoint::new(key_x_coordinate, key_y_coordinate)
+}
+
+fn fetch_key_from_registry_at(
+    x_coordinate_registry: SharedMutablePrivateGetter<Field, DELAY>,
+    y_coordinate_registry: SharedMutablePrivateGetter<Field, DELAY>,
+    header: Header
+) -> GrumpkinPoint {
+    let key_x_coordinate = x_coordinate_registry.get_current_value_in_private_at(header.global_variables.block_number as u32, header);
+    let key_y_coordinate = y_coordinate_registry.get_current_value_in_private_at(header.global_variables.block_number as u32, header);
+
+    GrumpkinPoint::new(key_x_coordinate, key_y_coordinate)
 }
 
 // Passes only when keys were not rotated - is expected to be called only when keys were not registered yet

noir-projects/aztec-nr/aztec/src/state_vars/shared_mutable/shared_mutable_private_getter.nr

@@ -1,7 +1,7 @@
-use dep::protocol_types::{hash::pedersen_hash, traits::FromField, address::AztecAddress};
+use dep::protocol_types::{hash::pedersen_hash, traits::FromField, address::AztecAddress, header::Header};
 
+use crate::history::public_storage::{public_storage_historical_read, _public_storage_historical_read};
 use crate::context::PrivateContext;
-use crate::history::public_storage::public_storage_historical_read;
 use crate::public_storage;
 use crate::state_vars::{
     storage::Storage,
@@ -34,36 +34,50 @@ impl<T, INITIAL_DELAY> SharedMutablePrivateGetter<T, INITIAL_DELAY> {
 
     pub fn get_current_value_in_private(self) -> T where T: FromField {
         let mut context = self.context;
+        let header = context.historical_header;
 
-        let (value_change, delay_change, historical_block_number) = self.historical_read_from_public_storage(context);
+        let (value_change, delay_change, historical_block_number) = self.historical_read_from_public_storage(header);
         let effective_minimum_delay = delay_change.get_effective_minimum_delay_at(historical_block_number);
         let block_horizon = value_change.get_block_horizon(historical_block_number, effective_minimum_delay);
 
         context.set_tx_max_block_number(block_horizon);
         value_change.get_current_at(historical_block_number)
     }
 
+    pub fn get_current_value_in_private_at(
+        self,
+        block_number: u32,
+        header: Header
+    ) -> T where T: FromField {
+        assert_eq(block_number, header.global_variables.block_number as u32);
+        // We don't need the delay change because even if we execute this in the future, we are looking at the value
+        // from a specific time in the past.
+        let (value_change, _, historical_block_number) = self.historical_read_from_public_storage(header);
+
+        value_change.get_current_at(historical_block_number)
+    }
+
     fn historical_read_from_public_storage(
         self,
-        context: PrivateContext
+        header: Header
     ) -> (ScheduledValueChange<T>, ScheduledDelayChange<INITIAL_DELAY>, u32) where T: FromField {
         let value_change_slot = self.get_value_change_storage_slot();
         let mut raw_value_change_fields = [0; 3];
         for i in 0..3 {
-            raw_value_change_fields[i] = public_storage_historical_read(
-                    context,
+            raw_value_change_fields[i] = _public_storage_historical_read(
                     value_change_slot + i as Field,
-                    self.other_contract_address
+                    self.other_contract_address,
+                    header,
                 );
         }
 
         let delay_change_slot = self.get_delay_change_storage_slot();
-        let raw_delay_change_fields = [public_storage_historical_read(context, delay_change_slot, context.this_address())];
+        let raw_delay_change_fields = [_public_storage_historical_read(delay_change_slot, self.other_contract_address, header)];
 
         let value_change = ScheduledValueChange::deserialize(raw_value_change_fields);
         let delay_change = ScheduledDelayChange::deserialize(raw_delay_change_fields);
 
-        let historical_block_number = context.historical_header.global_variables.block_number as u32;
+        let historical_block_number = header.global_variables.block_number as u32;
 
         (value_change, delay_change, historical_block_number)
     }

noir-projects/aztec-nr/easy-private-state/src/easy_private_uint.nr

@@ -49,11 +49,6 @@ impl<Context> EasyPrivateUint<&mut PrivateContext> {
             if maybe_notes[i].is_some() {
                 let note = maybe_notes[i].unwrap_unchecked();
 
-                // Ensure the notes are actually owned by the owner (to prevent user from generating a valid proof while
-                // spending someone else's notes).
-                // TODO (#6312): This will break with key rotation. Fix this. Will not be able to pass this assert after rotating key
-                assert(note.npk_m_hash.eq(owner_npk_m_hash));
-
                 // Removes the note from the owner's set of notes.
                 // docs:start:remove
                 self.set.remove(note);

noir-projects/aztec-nr/value-note/src/utils.nr

@@ -55,23 +55,12 @@ pub fn decrement_by_at_most(
     let options = create_note_getter_options_for_decreasing_balance(max_amount);
     let opt_notes = balance.get_notes(options);
 
-    let owner_npk_m_hash = get_npk_m_hash(balance.context, owner);
-
     let mut decremented = 0;
     for i in 0..opt_notes.len() {
         if opt_notes[i].is_some() {
             let note = opt_notes[i].unwrap_unchecked();
 
-            // This is similar to destroy_note, except we only compute the owner_npk_m_hash once instead of doing it in
-            // each loop iteration.
-
-            // Ensure the note is actually owned by the owner (to prevent user from generating a valid proof while
-            // spending someone else's notes).
-            // TODO (#6312): This will break with key rotation. Fix this. Will not be able to pass this after rotating keys.
-            assert(note.npk_m_hash.eq(owner_npk_m_hash));
-            decremented += note.value;
-
-            balance.remove(note);
+            decremented += destroy_note(balance, note);
         }
     }
 
@@ -88,18 +77,7 @@ pub fn decrement_by_at_most(
 
 // Removes the note from the owner's set of notes.
 // Returns the value of the destroyed note.
-pub fn destroy_note(
-    balance: PrivateSet<ValueNote, &mut PrivateContext>,
-    owner: AztecAddress,
-    note: ValueNote
-) -> Field {
-    // Ensure the note is actually owned by the owner (to prevent user from generating a valid proof while
-    // spending someone else's notes).
-    let owner_npk_m_hash = get_npk_m_hash(balance.context, owner);
-
-    // TODO (#6312): This will break with key rotation. Fix this. Will not be able to pass this after rotating keys.
-    assert(note.npk_m_hash.eq(owner_npk_m_hash));
-
+pub fn destroy_note(balance: PrivateSet<ValueNote, &mut PrivateContext>, note: ValueNote) -> Field {
     balance.remove(note);
 
     note.value

noir-projects/noir-contracts/contracts/card_game_contract/src/cards.nr

@@ -118,19 +118,13 @@ impl Deck<&mut PrivateContext> {
         inserted_cards
     }
 
-    pub fn get_cards<N>(&mut self, cards: [Card; N], owner: AztecAddress) -> [CardNote; N] {
-        let owner_npk_m_hash = get_npk_m_hash(self.set.context, owner);
-
+    pub fn get_cards<N>(&mut self, cards: [Card; N]) -> [CardNote; N] {
         let options = NoteGetterOptions::with_filter(filter_cards, cards);
         let maybe_notes = self.set.get_notes(options);
         let mut found_cards = [Option::none(); N];
         for i in 0..maybe_notes.len() {
             if maybe_notes[i].is_some() {
                 let card_note = CardNote::from_note(maybe_notes[i].unwrap_unchecked());
-                // Ensure the notes are actually owned by the owner (to prevent user from generating a valid proof while
-                // spending someone else's notes).
-                // TODO (#6312): This will break with key rotation. Fix this. Will not be able to pass this assert after rotating keys.
-                assert(card_note.note.npk_m_hash.eq(owner_npk_m_hash));
 
                 for j in 0..cards.len() {
                     if found_cards[j].is_none()
@@ -150,8 +144,8 @@ impl Deck<&mut PrivateContext> {
         )
     }
 
-    pub fn remove_cards<N>(&mut self, cards: [Card; N], owner: AztecAddress) {
-        let card_notes = self.get_cards(cards, owner);
+    pub fn remove_cards<N>(&mut self, cards: [Card; N]) {
+        let card_notes = self.get_cards(cards);
         for card_note in card_notes {
             self.set.remove(card_note.note);
         }

noir-projects/noir-contracts/contracts/card_game_contract/src/main.nr

@@ -36,7 +36,7 @@ contract CardGame {
         let player = context.msg_sender();
 
         let mut collection = storage.collections.at(player);
-        collection.remove_cards(cards, player);
+        collection.remove_cards(cards);
         let mut game_deck = storage.game_decks.at(game as Field).at(player);
         let _added_to_game_deck = game_deck.add_cards(cards, player);
         let strength = compute_deck_strength(cards);
@@ -68,7 +68,7 @@ contract CardGame {
         let player = context.msg_sender();
 
         let mut game_deck = storage.game_decks.at(game as Field).at(player);
-        game_deck.remove_cards([card], player);
+        game_deck.remove_cards([card]);
 
         // docs:start:call_public_function
         CardGame::at(context.this_address()).on_card_played(game, player, card.to_field()).enqueue(&mut context);

noir-projects/noir-contracts/contracts/child_contract/src/main.nr

@@ -1,6 +1,6 @@
 // A contract used along with `Parent` contract to test nested calls.
 contract Child {
-    use dep::aztec::prelude::{AztecAddress, FunctionSelector, PublicMutable, PrivateSet, PrivateContext, Deserialize};
+    use dep::aztec::prelude::{AztecAddress, FunctionSelector, PublicMutable, PrivateSet, PrivateContext, Deserialize, Map};
 
     use dep::aztec::{
         context::gas::GasOpts,
@@ -13,7 +13,7 @@ contract Child {
     #[aztec(storage)]
     struct Storage {
         current_value: PublicMutable<Field>,
-        a_private_value: PrivateSet<ValueNote>,
+        a_map_with_private_values: Map<AztecAddress, PrivateSet<ValueNote>>,
     }
 
     // Returns a sum of the input and the chain id and version of the contract in private circuit public input's return_values.
@@ -56,22 +56,15 @@ contract Child {
         let owner_ivpk_m = get_ivpk_m(&mut context, owner);
 
         let mut note = ValueNote::new(new_value, owner_npk_m_hash);
-        storage.a_private_value.insert(&mut note, true, owner_ivpk_m);
+        storage.a_map_with_private_values.at(owner).insert(&mut note, true, owner_ivpk_m);
         new_value
     }
 
     #[aztec(private)]
     fn private_get_value(amount: Field, owner: AztecAddress) -> Field {
-        let owner_npk_m_hash = get_npk_m_hash(&mut context, owner);
-
-        // TODO (#6312): This will break with key rotation. Fix this. Will not be able to find any notes after rotating key.
         let mut options = NoteGetterOptions::new();
-        options = options.select(ValueNote::properties().value, amount, Option::none()).select(
-            ValueNote::properties().npk_m_hash,
-            owner_npk_m_hash,
-            Option::none()
-        ).set_limit(1);
-        let notes = storage.a_private_value.get_notes(options);
+        options = options.select(ValueNote::properties().value, amount, Option::none()).set_limit(1);
+        let notes = storage.a_map_with_private_values.at(owner).get_notes(options);
         notes[0].unwrap_unchecked().value
     }
 

noir-projects/noir-contracts/contracts/claim_contract/src/main.nr

@@ -30,17 +30,13 @@ contract Claim {
             target_address == proof_note.header.contract_address, "Note does not correspond to the target contract"
         );
 
-        let msg_sender_npk_m_hash = get_npk_m_hash(&mut context, context.msg_sender());
-
-        // TODO (#6312): This will break with key rotation. Fix this. Will not be able to pass this assert after key rotation.
-        assert_eq(proof_note.npk_m_hash, msg_sender_npk_m_hash, "Note does not belong to the sender");
-
         // 2) Prove that the note hash exists in the note hash tree
         prove_note_inclusion(proof_note, context);
 
         // 3) Compute and emit a nullifier which is unique to the note and this contract to ensure the reward can be
         // claimed only once with the given note.
-        // Note: The nullifier is unique to the note and THIS contract because the protocol siloes all nullifiers with
+        // Note: Only the owner of the npk_m will be able to produce the nsk_app and compute this nullifier.
+        // The nullifier is unique to the note and THIS contract because the protocol siloes all nullifiers with
         // the address of a contract it was emitted from.
         context.push_new_nullifier(proof_note.compute_nullifier(&mut context), 0);