feat: aztec nr lib constraining nullifier key is fresh (#5939)

resolves #5688

---------

Co-authored-by: Jan Beneš <[email protected]>

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

@@ -1 +1,4 @@
-mod point_to_symmetric_key;
+mod getters;
+mod point_to_symmetric_key;
+
+use crate::keys::getters::get_fresh_nullifier_public_key_hash;

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

@@ -0,0 +1,79 @@
+use dep::protocol_types::{
+  address::{
+    AztecAddress,
+    PartialAddress
+  }, 
+  constants::{
+    GENERATOR_INDEX__PUBLIC_KEYS_HASH,
+    GENERATOR_INDEX__CONTRACT_ADDRESS_V1,
+    CANONICAL_KEY_REGISTRY_ADDRESS
+  },
+  grumpkin_point::GrumpkinPoint,
+};
+
+use crate::context::PrivateContext;
+use crate::hash::{
+  pedersen_hash, 
+  poseidon2_hash,
+};
+use crate::oracle::keys::get_public_keys_and_partial_address;
+use crate::state_vars::{
+  map::derive_storage_slot_in_map,
+  shared_mutable::shared_mutable_private_getter::SharedMutablePrivateGetter,
+};
+
+struct PublicKeyTypeEnum {
+  NULLIFIER: u8,
+}
+
+global PublicKeyType = PublicKeyTypeEnum {
+  NULLIFIER: 0,
+};
+
+pub fn get_fresh_nullifier_public_key_hash(
+  context: &mut PrivateContext,
+  address: AztecAddress, 
+) -> Field {
+  // This is the storage slot of the nullifier_public_key inside the key registry contract
+  // TODO: (#6133) We should have this be directly imported from the other contract if possible, or at least this should not be this brittle
+  let storage_slot_of_nullifier_public_key = 1;
+
+  let derived_slot = derive_storage_slot_in_map(storage_slot_of_nullifier_public_key, address);
+
+  // We read from the canonical Key Registry
+  // TODO: (#6134) It's a bit wonky because we need to know the delay for get_current_value_in_private to work correctly.
+  // We should allow for this usecase without needing to hard code it here.
+  let registry_private_getter: SharedMutablePrivateGetter<Field, 5> = SharedMutablePrivateGetter::new(*context, AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS), derived_slot);
+  let nullifier_public_key_hash_in_registry = registry_private_getter.get_current_value_in_private();
+
+  let nullifier_public_key_hash = if nullifier_public_key_hash_in_registry == 0 {
+    let keys = get_original_public_keys_internal(address);
+    poseidon2_hash(keys[PublicKeyType.NULLIFIER].serialize())
+  } else {
+    nullifier_public_key_hash_in_registry
+  };
+
+  nullifier_public_key_hash
+}
+
+// This constraint only works on keys that have not been rotated, otherwise this call will fail as the public keys are not constrained
+fn get_original_public_keys_internal(address: AztecAddress) -> [GrumpkinPoint; 4] {
+  let (public_keys, partial_address) = get_public_keys_and_partial_address(address);
+
+  let nullifier_pub_key = public_keys[0];
+  let incoming_pub_key = public_keys[1];
+  let outgoing_pub_key = public_keys[2];
+  let tagging_pub_key = public_keys[3];
+
+  let computed_address = AztecAddress::compute_from_public_keys_and_partial_address(
+    nullifier_pub_key,
+    incoming_pub_key,
+    outgoing_pub_key,
+    tagging_pub_key,
+    partial_address,
+  );
+
+  assert(computed_address.eq(address));
+
+  [nullifier_pub_key, incoming_pub_key, outgoing_pub_key, tagging_pub_key]
+}

noir-projects/aztec-nr/aztec/src/oracle.nr

@@ -11,6 +11,7 @@ mod get_nullifier_membership_witness;
 mod get_public_data_witness;
 mod get_membership_witness;
 mod get_public_key;
+mod keys;
 mod nullifier_key;
 mod get_sibling_path;
 mod unsafe_rand;

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

@@ -0,0 +1,28 @@
+use dep::protocol_types::{
+  address::{
+    AztecAddress,
+    PartialAddress,
+  }, 
+  grumpkin_point::GrumpkinPoint,
+};
+
+use crate::hash::poseidon2_hash;
+
+#[oracle(getPublicKeysAndPartialAddress)]
+fn get_public_keys_and_partial_address_oracle(_address: AztecAddress) -> [Field; 9] {}
+
+unconstrained fn get_public_keys_and_partial_address_oracle_wrapper(address: AztecAddress) -> [Field; 9] {
+  get_public_keys_and_partial_address_oracle(address)
+}
+
+fn get_public_keys_and_partial_address(address: AztecAddress) -> ([GrumpkinPoint; 4], PartialAddress) {
+  let result = get_public_keys_and_partial_address_oracle_wrapper(address);
+
+  let nullifier_pub_key = GrumpkinPoint::new(result[0], result[1]);
+  let incoming_pub_key = GrumpkinPoint::new(result[2], result[3]);
+  let outgoing_pub_key = GrumpkinPoint::new(result[4], result[5]);
+  let tagging_pub_key = GrumpkinPoint::new(result[6], result[7]);
+  let partial_address = PartialAddress::from_field(result[8]);
+
+  ([nullifier_pub_key, incoming_pub_key, outgoing_pub_key, tagging_pub_key], partial_address)
+}

noir-projects/aztec-nr/aztec/src/state_vars/map.nr

@@ -27,10 +27,14 @@ impl<K, V> Map<K, V> {
     // docs:start:at
     pub fn at(self, key: K) -> V where K: ToField {
         // TODO(#1204): use a generator index for the storage slot
-        let derived_storage_slot = pedersen_hash([self.storage_slot, key.to_field()], 0);
+        let derived_storage_slot = derive_storage_slot_in_map(self.storage_slot, key);
 
         let state_var_constructor = self.state_var_constructor;
         state_var_constructor(self.context, derived_storage_slot)
     }
     // docs:end:at
 }
+
+pub fn derive_storage_slot_in_map<K>(storage_slot: Field, key: K) -> Field where K: ToField {
+    pedersen_hash([storage_slot, key.to_field()], 0)
+}

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

@@ -7,6 +7,7 @@ contract KeyRegistry {
       Map
     }, 
     protocol_types::{
+      grumpkin_point::GrumpkinPoint,
       address::{
         AztecAddress,
         PublicKeysHash,
@@ -24,9 +25,12 @@ contract KeyRegistry {
 
   #[aztec(storage)]
   struct Storage {
+    //! This should stay at storage slot 1. If you change this, make sure you change the hardcoded value in keys/assert_public_key_freshness.
+    //! We use this hardcoded storage slot with derive_storage_slot_in_map and the SharedMutablePrivateGetter to directly read the value at an address in this contract.
+    nullifier_public_key_hash_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
+
     // We are not supporting rotating / changing keys other than the nullifier public in the registry at the moment, but will in the future.
     // Uncomment lines below to enable that functionality
-    nullifier_public_key_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
     // incoming_public_key_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
     // outgoing_public_key_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
     // tagging_public_key_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
@@ -35,71 +39,63 @@ contract KeyRegistry {
   #[aztec(public)]
   fn rotate_nullifier_public_key(
     address: AztecAddress,
-    new_nullifier_public_key: Field,
+    new_nullifier_public_key: GrumpkinPoint,
+    nonce: Field,
   ) {
     assert(
-      new_nullifier_public_key != 0,
+      !new_nullifier_public_key.is_zero(),
       "New nullifier public key must be non-zero"
     );
 
+    // TODO: (#6137)
     if (!address.eq(context.msg_sender())) {
       assert_current_call_valid_authwit_public(&mut context, address);
+    } else {
+      assert(nonce == 0, "invalid nonce");
     }
 
-    let nullifier_key_registry = storage.nullifier_public_key_registry.at(address);
+    let nullifier_key_registry = storage.nullifier_public_key_hash_registry.at(address);
 
-    nullifier_key_registry.schedule_value_change(new_nullifier_public_key);
+    nullifier_key_registry.schedule_value_change(poseidon2_hash(new_nullifier_public_key.serialize()));
   }
 
   #[aztec(public)]
   fn register(
     address: AztecAddress,
     partial_address: PartialAddress,
-    nullifier_public_key: Field,
-    incoming_public_key: Field,
-    outgoing_public_key: Field,
-    tagging_public_key: Field,
+    nullifier_public_key: GrumpkinPoint,
+    incoming_public_key: GrumpkinPoint,
+    outgoing_public_key: GrumpkinPoint,
+    tagging_public_key: GrumpkinPoint,
   ) {
     assert(
-      (nullifier_public_key != 0) &
-      (incoming_public_key != 0) &
-      (outgoing_public_key != 0) &
-      (tagging_public_key != 0),
+      !partial_address.is_zero() &
+      !nullifier_public_key.is_zero() &
+      !incoming_public_key.is_zero() &
+      !outgoing_public_key.is_zero() &
+      !tagging_public_key.is_zero(),
        "All public keys must be non-zero"
     );
 
-    // TODO (ek): Do it below after refactoring all public_keys_hash_elemtns
-    // let public_keys_hash = PublicKeysHash::compute(nullifier_public_key, tagging_public_key, incoming_public_key, outgoing_public_key);
-    // let address = AztecAddress::compute(public_keys_hash, partial_address);
     // We could also pass in original_public_keys_hash instead of computing it here, if all we need the original one is for being able to prove ownership of address
-    let public_keys_hash = poseidon2_hash([
-        nullifier_public_key,
-        incoming_public_key,
-        outgoing_public_key,
-        tagging_public_key,
-        GENERATOR_INDEX__PUBLIC_KEYS_HASH,
-      ]
-    );
-
-    let computed_address = AztecAddress::from_field(
-      poseidon2_hash([
-          partial_address.to_field(),
-          public_keys_hash.to_field(),
-          GENERATOR_INDEX__CONTRACT_ADDRESS_V1 as Field,
-        ]
-      )
+    let computed_address = AztecAddress::compute_from_public_keys_and_partial_address(
+      nullifier_public_key,
+      incoming_public_key,
+      outgoing_public_key,
+      tagging_public_key,
+      partial_address,
     );
 
     assert(computed_address.eq(address), "Computed address does not match supplied address");
 
-    let nullifier_key_registry = storage.nullifier_public_key_registry.at(address);
+    let nullifier_key_hash_registry = storage.nullifier_public_key_hash_registry.at(address);
     // We are not supporting rotating / changing keys other than the nullifier public in the registry at the moment, but will in the future.
     // Uncomment lines below to enable that functionality
     // let incoming_key_registry = storage.incoming_public_key_registry.at(address);
     // let outgoing_key_registry = storage.outgoing_public_key_registry.at(address);
     // let tagging_key_registry = storage.taggin_public_key_registry.at(address);
 
-    nullifier_key_registry.schedule_value_change(nullifier_public_key);
+    nullifier_key_hash_registry.schedule_value_change(poseidon2_hash(nullifier_public_key.serialize()));
     // We are not supporting rotating / changing keys other than the nullifier public in the registry at the moment, but will in the future.
     // Uncomment lines below to enable that functionality    // incoming_key_registry.schedule_value_change(new_incoming_public_key);
     // outgoing_key_registry.schedule_value_change(new_outgoing_public_key);

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

@@ -10,16 +10,18 @@ contract Test {
 
     use dep::aztec::protocol_types::{
         abis::private_circuit_public_inputs::PrivateCircuitPublicInputs,
-        constants::MAX_NOTE_HASH_READ_REQUESTS_PER_CALL, traits::Serialize
+        constants::{MAX_NOTE_HASH_READ_REQUESTS_PER_CALL, CANONICAL_KEY_REGISTRY_ADDRESS}, traits::{Serialize, ToField, FromField},
+        grumpkin_point::GrumpkinPoint
     };
 
     use dep::aztec::note::constants::MAX_NOTES_PER_PAGE;
 
-    use dep::aztec::state_vars::shared_mutable::SharedMutablePrivateGetter;
+    use dep::aztec::state_vars::{shared_mutable::SharedMutablePrivateGetter, map::derive_storage_slot_in_map};
 
     use dep::aztec::{
+        keys::getters::get_fresh_nullifier_public_key_hash,
         context::{Context, inputs::private_context_inputs::PrivateContextInputs},
-        hash::{pedersen_hash, compute_secret_hash, ArgsHasher},
+        hash::{pedersen_hash, poseidon2_hash, compute_secret_hash, ArgsHasher},
         note::{
         lifecycle::{create_note, destroy_note}, note_getter::{get_notes, view_notes},
         note_getter_options::NoteStatus
@@ -375,38 +377,47 @@ contract Test {
         constant.value
     }
 
+    // This function is used in the e2e_state_vars to test the SharedMutablePrivateGetter in isolation
     #[aztec(private)]
-    fn test_shared_mutable_private_getter_for_registry_contract(
+    fn test_shared_mutable_private_getter<T>(
         contract_address_to_read: AztecAddress,
+        storage_slot_of_shared_mutable: Field
+    ) -> Field where T: FromField, T: ToField {
+        // It's a bit wonky because we need to know the delay for get_current_value_in_private to work correctly
+        let test: SharedMutablePrivateGetter<T, 5> = SharedMutablePrivateGetter::new(
+            context,
+            contract_address_to_read,
+            storage_slot_of_shared_mutable
+        );
+
+        let ret = test.get_current_value_in_private();
+
+        ret.to_field()
+    }
+
+    // This function is used for testing the registry contract and fresh public key getters. If nothing exists in the registry, but we have added public
+    // keys to the pxe, this function will return nothing, but the public key getters will return the correct value
+    #[aztec(private)]
+    fn test_shared_mutable_private_getter_for_registry_contract(
         storage_slot_of_shared_mutable: Field,
         address_to_get_in_registry: AztecAddress
-    ) {
+    ) -> Field {
         // We have to derive this slot to get the location of the shared mutable inside the Map
-        let derived_slot = dep::aztec::hash::pedersen_hash(
-            [storage_slot_of_shared_mutable, address_to_get_in_registry.to_field()],
-            0
-        );
+        let derived_slot = derive_storage_slot_in_map(storage_slot_of_shared_mutable, address_to_get_in_registry);
+
         // It's a bit wonky because we need to know the delay for get_current_value_in_private to work correctly
-        let registry_private_getter: SharedMutablePrivateGetter<AztecAddress, 5> = SharedMutablePrivateGetter::new(context, contract_address_to_read, derived_slot);
+        let registry_private_getter: SharedMutablePrivateGetter<Field, 5> = SharedMutablePrivateGetter::new(context, AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS), derived_slot);
         let nullifier_public_key = registry_private_getter.get_current_value_in_private();
 
-        context.emit_unencrypted_log(nullifier_public_key);
+        nullifier_public_key
     }
 
     #[aztec(private)]
-    fn test_shared_mutable_private_getter(
-        contract_address_to_read: AztecAddress,
-        storage_slot_of_shared_mutable: Field
+    fn test_nullifier_key_freshness(
+        address: AztecAddress,
+        public_nullifying_key: GrumpkinPoint,
     ) {
-        // It's a bit wonky because we need to know the delay for get_current_value_in_private to work correctly
-        let test: SharedMutablePrivateGetter<AztecAddress, 5> = SharedMutablePrivateGetter::new(
-            context,
-            contract_address_to_read,
-            storage_slot_of_shared_mutable
-        );
-        let authorized = test.get_current_value_in_private();
-
-        context.emit_unencrypted_log(authorized);
+        assert_eq(get_fresh_nullifier_public_key_hash(&mut context, address), poseidon2_hash(public_nullifying_key.serialize()));
     }
 
     #[aztec(public)]

noir-projects/noir-protocol-circuits/crates/types/src/address/aztec_address.nr

@@ -61,6 +61,25 @@ impl AztecAddress {
         )
     }
 
+    pub fn compute_from_public_keys_and_partial_address(
+        nullifier_public_key: GrumpkinPoint,
+        incoming_public_key: GrumpkinPoint,
+        outgoing_public_key: GrumpkinPoint,
+        tagging_public_key: GrumpkinPoint,
+        partial_address: PartialAddress,
+    ) -> AztecAddress {
+        let public_keys_hash = PublicKeysHash::compute_new(
+            nullifier_public_key,
+            incoming_public_key,
+            outgoing_public_key,
+            tagging_public_key,
+        );
+
+        let computed_address = AztecAddress::compute(public_keys_hash, partial_address);
+
+        computed_address
+    }
+
     pub fn is_zero(self) -> bool {
         self.inner == 0
     }

noir-projects/noir-protocol-circuits/crates/types/src/address/partial_address.nr

@@ -68,6 +68,10 @@ impl PartialAddress {
         self.inner
     }
 
+    pub fn is_zero(self) -> bool {
+        self.to_field() == 0
+    }
+
     pub fn assert_is_zero(self) {
         assert(self.to_field() == 0);
     }