feat: making keys getters complete

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

@@ -1,4 +1,8 @@
 mod getters;
 mod point_to_symmetric_key;
 
-use crate::keys::getters::get_fresh_nullifier_public_key_hash;
+use crate::keys::getters::{get_npk_m, get_ivpk_m,
+// Commented out as it's currently not enabled in key registry
+// get_ovpk_m,
+// get_tpk_m
+};

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

@@ -1,79 +1,99 @@
-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 dep::protocol_types::{address::AztecAddress, constants::CANONICAL_KEY_REGISTRY_ADDRESS, grumpkin_point::GrumpkinPoint};
+use crate::{
+    context::PrivateContext, oracle::keys::get_public_keys_and_partial_address,
+    state_vars::{
+    map::derive_storage_slot_in_map,
+    shared_mutable::shared_mutable_private_getter::SharedMutablePrivateGetter
+}
 };
 
-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,
-};
+// Note: In fetch_key_from_registry we expect that the shared mutable slot is index * 2 + 1 for the x coordinate and
+// index * 2 + 2 for the y coordinate. For example, the npk_m x coordinates will be stored in a map at storage slot
+// 0 * 2 + 1 = 1, and the npk_m y coordinates at 2 * 2 + 2 = 6. If this changes the function will need to be
+// refactored.
+global NULLIFIER_INDEX = 0;
+global INCOMING_INDEX = 1;
+global OUTGOING_INDEX = 2;
+global TAGGING_INDEX = 3;
+
+global DELAY = 5;
 
-struct PublicKeyTypeEnum {
-  NULLIFIER: u8,
+pub fn get_npk_m(context: &mut PrivateContext, address: AztecAddress) -> GrumpkinPoint {
+    get_master_key(context, address, NULLIFIER_INDEX)
 }
 
-global PublicKeyType = PublicKeyTypeEnum {
-  NULLIFIER: 0,
-};
+pub fn get_ivpk_m(context: &mut PrivateContext, address: AztecAddress) -> GrumpkinPoint {
+    get_master_key(context, address, INCOMING_INDEX)
+}
 
-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;
+// Commented out as it's currently not enabled in key registry
+// pub fn get_ovpk_m(context: &mut PrivateContext, address: AztecAddress) -> GrumpkinPoint {
+//     get_master_key(context, address, OUTGOING_INDEX)
+// }
+// 
+// pub fn get_tpk_m(context: &mut PrivateContext, address: AztecAddress) -> GrumpkinPoint {
+//     get_master_key(context, address, TAGGING_INDEX)
+// }
 
-  let derived_slot = derive_storage_slot_in_map(storage_slot_of_nullifier_public_key, address);
+fn get_master_key(
+    context: &mut PrivateContext,
+    address: AztecAddress,
+    key_index: Field
+) -> GrumpkinPoint {
+    let key = fetch_key_from_registry(context, key_index, address);
+    if key.is_zero() {
+        // Keys were not registered in registry yet --> fetch key from PXE
+        fetch_and_constrain_keys(address)[key_index]
+    } else {
+        // Keys were registered --> return the key
+        key
+    }
+}
 
-  // 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();
+fn fetch_key_from_registry(
+    context: &mut PrivateContext,
+    key_index: Field,
+    address: AztecAddress
+) -> GrumpkinPoint {
+    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);
+    let y_coordinate_derived_slot = derive_storage_slot_in_map(y_coordinate_map_slot, address);
 
-  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
-  };
+    let x_coordinate_registry: SharedMutablePrivateGetter<Field, DELAY> = SharedMutablePrivateGetter::new(
+        *context,
+        AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS),
+        x_coordinate_derived_slot
+    );
+    let y_coordinate_registry: SharedMutablePrivateGetter<Field, DELAY> = SharedMutablePrivateGetter::new(
+        *context,
+        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();
 
-  nullifier_public_key_hash
+    GrumpkinPoint::new(x_coordinate, y_coordinate)
 }
 
-// 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);
+// Passes only when keys were not rotated - is expected to be called only when keys were not registered yet
+fn fetch_and_constrain_keys(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 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,
-  );
+    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));
+    assert(computed_address.eq(address));
 
-  [nullifier_pub_key, incoming_pub_key, outgoing_pub_key, tagging_pub_key]
+    [nullifier_pub_key, incoming_pub_key, outgoing_pub_key, tagging_pub_key]
 }

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

@@ -1,104 +1,88 @@
 contract KeyRegistry {
-  use dep::authwit::auth::assert_current_call_valid_authwit_public;
+    use dep::authwit::auth::assert_current_call_valid_authwit_public;
 
-  use dep::aztec::{
-    state_vars::{
-      SharedMutable,
-      Map
-    }, 
-    protocol_types::{
-      grumpkin_point::GrumpkinPoint,
-      address::{
-        AztecAddress,
-        PublicKeysHash,
-        PartialAddress,
-      },
-      constants::{
-        GENERATOR_INDEX__CONTRACT_ADDRESS_V1,
-        GENERATOR_INDEX__PUBLIC_KEYS_HASH
-      },
-      hash::poseidon2_hash,
-    },
-  };
+    use dep::aztec::{
+        state_vars::{SharedMutable, Map},
+        protocol_types::{
+        grumpkin_point::GrumpkinPoint, address::{AztecAddress, PartialAddress},
+    }
+    };
 
-  global KEY_ROTATION_DELAY = 5;
+    global KEY_ROTATION_DELAY = 5;
 
-  #[aztec(storage)]
+    #[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>>,
+    // The following stores a hash of individual master public keys
+    // If you change slots of vars below, you must update the slots in `SharedMutablePrivateGetter` in aztec-nr/keys.
+    // We store x and y coordinates in individual shared mutables as shared mutable currently supports only 1 field
+    npk_m_x_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
+    npk_m_y_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
-    // 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>>,
+    ivpk_m_x_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
+    ivpk_m_y_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
+
+    ovpk_m_x_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
+    ovpk_m_y_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
+
+    tpk_m_x_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
+    tpk_m_y_registry: Map<AztecAddress, SharedMutable<Field, KEY_ROTATION_DELAY>>,
   }
 
-  #[aztec(public)]
+    #[aztec(public)]
   fn rotate_nullifier_public_key(
-    address: AztecAddress,
-    new_nullifier_public_key: GrumpkinPoint,
-    nonce: Field,
-  ) {
-    assert(
-      !new_nullifier_public_key.is_zero(),
-      "New nullifier public key must be non-zero"
-    );
+        address: AztecAddress,
+        new_nullifier_public_key: GrumpkinPoint,
+        nonce: Field
+    ) {
+        // TODO: (#6137)
+        if (!address.eq(context.msg_sender())) {
+            assert_current_call_valid_authwit_public(&mut context, address);
+        } else {
+            assert(nonce == 0, "invalid nonce");
+        }
 
-    // TODO: (#6137)
-    if (!address.eq(context.msg_sender())) {
-      assert_current_call_valid_authwit_public(&mut context, address);
-    } else {
-      assert(nonce == 0, "invalid nonce");
+        let npk_m_x_registry = storage.npk_m_x_registry.at(address);
+        let npk_m_y_registry = storage.npk_m_y_registry.at(address);
+        npk_m_x_registry.schedule_value_change(new_nullifier_public_key.x);
+        npk_m_y_registry.schedule_value_change(new_nullifier_public_key.y);
     }
 
-    let nullifier_key_registry = storage.nullifier_public_key_hash_registry.at(address);
-
-    nullifier_key_registry.schedule_value_change(poseidon2_hash(new_nullifier_public_key.serialize()));
-  }
-
-  #[aztec(public)]
+    #[aztec(public)]
   fn register(
-    address: AztecAddress,
-    partial_address: PartialAddress,
-    nullifier_public_key: GrumpkinPoint,
-    incoming_public_key: GrumpkinPoint,
-    outgoing_public_key: GrumpkinPoint,
-    tagging_public_key: GrumpkinPoint,
-  ) {
-    assert(
-      !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"
-    );
+        address: AztecAddress,
+        partial_address: PartialAddress,
+        nullifier_public_key: GrumpkinPoint,
+        incoming_public_key: GrumpkinPoint,
+        outgoing_public_key: GrumpkinPoint,
+        tagging_public_key: GrumpkinPoint
+    ) {
+        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
+        );
 
-    // 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 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");
 
-    assert(computed_address.eq(address), "Computed address does not match supplied address");
+        let npk_m_x_registry = storage.npk_m_x_registry.at(address);
+        let npk_m_y_registry = storage.npk_m_y_registry.at(address);
+        let ivpk_m_x_registry = storage.ivpk_m_x_registry.at(address);
+        let ivpk_m_y_registry = storage.ivpk_m_y_registry.at(address);
+        // let ovpk_m_x_registry = storage.ovpk_m_x_registry.at(address);
+        // let ovpk_m_y_registry = storage.ovpk_m_y_registry.at(address);
+        // let tpk_m_x_registry = storage.tpk_m_x_registry.at(address);
+        // let tpk_m_y_registry = storage.tpk_m_y_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_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);
-    // tagging_key_registry.schedule_value_change(new_tagging_public_key);
-  }
+        npk_m_x_registry.schedule_value_change(nullifier_public_key.x);
+        npk_m_y_registry.schedule_value_change(nullifier_public_key.y);
+        ivpk_m_x_registry.schedule_value_change(incoming_public_key.x);
+        ivpk_m_y_registry.schedule_value_change(incoming_public_key.y);
+        // Commented out as we hit the max enqueued public calls limit when not done so
+        // ovpk_m_x_registry.schedule_value_change(outgoing_public_key.x);
+        // ovpk_m_y_registry.schedule_value_change(outgoing_public_key.y);
+        // tpk_m_x_registry.schedule_value_change(tagging_public_key.x);
+        // tpk_m_y_registry.schedule_value_change(tagging_public_key.y);
+    }
 }

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

@@ -19,7 +19,7 @@ contract Test {
     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,
+        keys::getters::get_npk_m,
         context::{Context, inputs::private_context_inputs::PrivateContextInputs},
         hash::{pedersen_hash, poseidon2_hash, compute_secret_hash, ArgsHasher},
         note::{
@@ -425,23 +425,15 @@ contract Test {
 
         // 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<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();
-
-        nullifier_public_key
+        registry_private_getter.get_current_value_in_private()
     }
 
     #[aztec(private)]
     fn test_nullifier_key_freshness(
         address: AztecAddress,
         public_nullifying_key: GrumpkinPoint,
     ) {
-        assert_eq(get_fresh_nullifier_public_key_hash(&mut context, address), poseidon2_hash(public_nullifying_key.serialize()));
-    }
-
-    #[aztec(public)]
-    fn delay() {
-        // We use this as a util function to "mine a block"
-        context.emit_unencrypted_log("dummy");
+        assert_eq(get_npk_m(&mut context, address), public_nullifying_key);
     }
 
     // Purely exists for testing