feat!: alternative key registry contract (#7523)

This is a tentative redesign of the key registry contract, which lets us
read all 4 keys in a single merkle inclusion proof, while also
performing fewer public storage writes. It relies on `PublicMutable`
with custom deadlines instead of `SharedMutable` as the key registry
must be lax during reads to prevent abuse from accounts performing
frequent rotation.

This is meant to be a proof of concept and reference for discussion, so
there are no tests yet. If we were to adopt this contract it should be
relatively straightforward to replace the current getters with
`get_current_public_keys`.

---

Update: we've now decided to move forward with this design, even if
whether we actually do have full key rotation is up for debate. This is
a good middle ground in terms of the code being performant and easy to
either add or remove rotation in the future.

What this PR does is introduce the new registry and switch all old
contracts to use it. They'll be using historical mode, i.e. reading keys
at a specific block in the past instead of reading the current keys,
resulting in no max block number constraints. This is a departure from
the current behavior, but is only temporary until we switch over from
the old API to the new one
(AztecProtocol/aztec-packages#7953) so that we
can then benefit from reduce gate counts
(AztecProtocol/aztec-packages#7954). I've also
left the original contract, libraries, etc., unmodified so as to later
delete them cleanly instead of making this PR too messy
(AztecProtocol/aztec-packages#7955).

---------

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

aztec/src/context/unconstrained_context.nr

@@ -31,6 +31,12 @@ impl UnconstrainedContext {
         Self { block_number, contract_address, version, chain_id }
     }
 
+    unconstrained fn at_historical(contract_address: AztecAddress, block_number: u32) -> Self {
+        let chain_id = get_chain_id();
+        let version = get_version();
+        Self { block_number, contract_address, version, chain_id }
+    }
+
     fn block_number(self) -> u32 {
         self.block_number
     }

aztec/src/history/public_storage.nr

@@ -10,7 +10,7 @@ trait PublicStorageHistoricalRead {
     fn public_storage_historical_read(header: Header, storage_slot: Field, contract_address: AztecAddress) -> Field;
 }
 
-impl PublicStorageHistoricalRead for Header { 
+impl PublicStorageHistoricalRead for Header {
     fn public_storage_historical_read(self, storage_slot: Field, contract_address: AztecAddress) -> Field {
         // 1) Compute the leaf slot by siloing the storage slot with the contract address
         let public_data_tree_index = poseidon2_hash_with_separator(

aztec/src/keys/getters/mod.nr

@@ -0,0 +1,154 @@
+use dep::protocol_types::{
+    header::Header, abis::validation_requests::KeyValidationRequest, address::AztecAddress,
+    constants::CANONICAL_KEY_REGISTRY_ADDRESS, point::Point, storage::map::derive_storage_slot_in_map,
+    traits::is_empty
+};
+use crate::{
+    context::{PrivateContext, UnconstrainedContext},
+    oracle::{keys::get_public_keys_and_partial_address, key_validation_request::get_key_validation_request},
+    keys::{
+    public_keys::{PublicKeys, PUBLIC_KEYS_LENGTH}, stored_keys::StoredKeys,
+    constants::{NULLIFIER_INDEX, INCOMING_INDEX, OUTGOING_INDEX, TAGGING_INDEX}
+},
+    state_vars::{
+    shared_mutable::shared_mutable_private_getter::SharedMutablePrivateGetter,
+    public_mutable::PublicMutable, map::Map
+}
+};
+
+global DELAY = 5;
+
+mod test;
+
+// docs:start:key-getters
+trait KeyGetters {
+  fn get_npk_m(header: Header, context: &mut PrivateContext, address: AztecAddress) -> Point;
+  fn get_ivpk_m(header: Header, context: &mut PrivateContext, address: AztecAddress) -> Point;
+  fn get_ovpk_m(header: Header, context: &mut PrivateContext, address: AztecAddress) -> Point;
+  fn get_tpk_m(header: Header, context: &mut PrivateContext, address: AztecAddress) -> Point;
+  fn get_npk_m_hash(header: Header, context: &mut PrivateContext, address: AztecAddress) -> Field;
+}
+
+impl KeyGetters for Header {
+    fn get_npk_m(self, _context: &mut PrivateContext, address: AztecAddress) -> Point {
+        get_historical_public_keys(self, address).npk_m
+    }
+
+    fn get_ivpk_m(self, _context: &mut PrivateContext, address: AztecAddress) -> Point {
+        get_historical_public_keys(self, address).ivpk_m
+    }
+
+    fn get_ovpk_m(self, _context: &mut PrivateContext, address: AztecAddress) -> Point {
+        get_historical_public_keys(self, address).ovpk_m
+    }
+
+    fn get_tpk_m(self, _context: &mut PrivateContext, address: AztecAddress) -> Point {
+        get_historical_public_keys(self, address).tpk_m
+    }
+
+    fn get_npk_m_hash(self, context: &mut PrivateContext, address: AztecAddress) -> Field {
+        self.get_npk_m(context, address).hash()
+    }
+}
+// docs:end:key-getters
+
+// A helper function since requesting nsk_app is very common
+// TODO(#6543)
+pub fn get_nsk_app(npk_m_hash: Field) -> Field {
+    get_key_validation_request(npk_m_hash, NULLIFIER_INDEX).sk_app
+}
+
+// This is the number of blocks that must pass after a key rotation event until the old keys are fully phased out and
+// become invalid.
+global KEY_REGISTRY_UPDATE_BLOCKS = 5;
+
+global KEY_REGISTRY_STORAGE_SLOT = 1;
+
+// Returns all current public keys for a given account, applying proper constraints to the context. We read all
+// keys at once since the constraints for reading them all are actually fewer than if we read them one at a time - any
+// read keys that are not required by the caller can simply be discarded.
+pub fn get_current_public_keys(context: &mut PrivateContext, account: AztecAddress) -> PublicKeys {
+    // We're going to perform historical reads from public storage, and so need to constrain the caller so that they
+    // cannot use very old blocks when constructing proofs, and hence e.g. read very old keys. We are lax and allow
+    // _any_ recent block number to be used, regardless of whether there may have been a recent key rotation. This means
+    // that multiple sets of keys are valid for a while immediately after rotation, until the old keys become phased
+    // out. We *must* be lax to prevent denial of service and transaction fingerprinting attacks by accounts that rotate
+    // their keys frequently.
+    // Note that we constrain the max block number even if the registry ends up being empty: this ensures that proof of
+    // an empty registry is also fresh.
+    let current_header = context.get_header();
+    context.set_tx_max_block_number(current_header.global_variables.block_number as u32 + KEY_REGISTRY_UPDATE_BLOCKS);
+
+    get_historical_public_keys(current_header, account)
+}
+
+// Returns historical public keys for a given account at some block determined by a block header. We read all keys at
+// once since the constraints for reading them all are actually fewer than if we read them one at a time - any read keys
+// that are not required by the caller can simply be discarded.
+// WARNING: if called with a historical header created from a fixed block this function will explicitly ignore key
+// rotation! This means that callers of this may force a user to use old keys, potentially leaking privacy (e.g. if the
+// old keys were leaked). Only call this function with a header from a fixed block if you understand the implications of
+// breaking key rotation very well.
+pub fn get_historical_public_keys(historical_header: Header, account: AztecAddress) -> PublicKeys {
+    // TODO: improve this so that we always hint the correct set of keys (either registry or canonical) and hash them
+    // once instead of having two different hints and twice as many constraints due to the double hashing.
+
+    // The key registry is the primary source of information for keys, as that's where accounts store their new keys
+    // when they perform rotation. The key registry conveniently stores a hash of each user's keys, so we can read that
+    // single field and then prove that we know its preimage (i.e. the current set of keys).
+    let key_registry_hash = key_registry_hash_public_historical_read(historical_header, account);
+    if key_registry_hash != 0 {
+        let hinted_registry_public_keys = key_registry_get_stored_keys_hint(
+            account,
+            historical_header.global_variables.block_number as u32
+        );
+        assert_eq(hinted_registry_public_keys.hash().to_field(), key_registry_hash);
+
+        hinted_registry_public_keys
+    } else {
+        // If nothing was written to the registry, we may still be able to produce the correct keys if we happen to know
+        // the canonical set (i.e. the ones that are part of the account's preimage).
+        let (hinted_canonical_public_keys, partial_address) = get_public_keys_and_partial_address(account);
+        assert_eq(
+            account, AztecAddress::compute(hinted_canonical_public_keys.hash(), partial_address), "Invalid public keys hint for address"
+        );
+
+        hinted_canonical_public_keys
+    }
+}
+
+fn key_registry_hash_public_historical_read(historical_header: Header, account: AztecAddress) -> Field {
+    // The keys are stored in a Map that is keyed with the address of each account, so we first derive the corresponding
+    // slot for this account.
+    let keys_storage_slot = derive_storage_slot_in_map(KEY_REGISTRY_STORAGE_SLOT, account);
+
+    // The keys are stored as [ ...serialized_keys, hash ], and since arrays get allocated sequential storage slots
+    // (prior to siloing!), we simply add the length to the base slot to get the last element.
+    let hash_storage_slot = keys_storage_slot + PUBLIC_KEYS_LENGTH as Field;
+
+    historical_header.public_storage_historical_read(
+        hash_storage_slot,
+        AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS)
+    )
+}
+
+unconstrained fn key_registry_get_stored_keys_hint(account: AztecAddress, block_number: u32) -> PublicKeys {
+    // This is equivalent to the key registry contract having an unconstrained getter that we call from an oracle, but
+    // PXE does not yet support that functionality so we do this manually instad. Note that this would be a *historical*
+    // call!
+
+    // TODO (#7524): call the unconstrained KeyRegistry.get_current_keys() function instead
+
+    let context = UnconstrainedContext::at_historical(
+        AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS),
+        block_number
+    );
+    let keys_storage = Map::new(
+        context,
+        KEY_REGISTRY_STORAGE_SLOT,
+        |context, slot| { PublicMutable::new(context, slot) }
+    );
+
+    let stored_keys: StoredKeys = keys_storage.at(account).read();
+    stored_keys.public_keys
+}

aztec/src/keys/getters/test.nr

@@ -0,0 +1,58 @@
+use crate::keys::getters::{get_current_public_keys, get_historical_public_keys, KEY_REGISTRY_UPDATE_BLOCKS};
+use crate::context::PrivateContext;
+use dep::protocol_types::address::AztecAddress;
+
+use crate::test::helpers::{cheatcodes, test_environment::TestEnvironment, utils::TestAccount};
+use dep::std::test::OracleMock;
+
+global KEY_ORACLE_RESPONSE_LENGTH = 13; // 12 fields for the keys, one field for the partial address
+
+fn setup() -> (TestEnvironment, PrivateContext, TestAccount) {
+    let mut env = TestEnvironment::new();
+    let account = cheatcodes::create_account();
+
+    let historical_block_number = env.block_number();
+    let context = env.private_at(historical_block_number);
+
+    (env, context, account)
+}
+
+#[test(should_fail_with="Invalid public keys hint for address")]
+fn test_get_current_keys_unknown_unregistered() {
+    let (_, context, account) = setup();
+
+    let _ = OracleMock::mock("getPublicKeysAndPartialAddress").returns([0; KEY_ORACLE_RESPONSE_LENGTH]).times(1);
+    let _ = get_current_public_keys(&mut context, account.address);
+}
+
+#[test(should_fail_with="Invalid public keys hint for address")]
+fn test_get_historical_keys_unknown_unregistered() {
+    let (_, context, account) = setup();
+    let historical_header = context.get_header();
+
+    let _ = OracleMock::mock("getPublicKeysAndPartialAddress").returns([0; KEY_ORACLE_RESPONSE_LENGTH]).times(1);
+    let _ = get_historical_public_keys(historical_header, account.address);
+}
+
+#[test]
+fn test_get_current_keys_known_unregistered() {
+    let (_, mut context, account) = setup();
+
+    let current_keys = get_current_public_keys(&mut context, account.address);
+
+    assert_eq(current_keys, account.keys);
+    assert_eq(
+        context.max_block_number.unwrap(), context.historical_header.global_variables.block_number as u32 + KEY_REGISTRY_UPDATE_BLOCKS
+    );
+}
+
+#[test]
+fn test_get_historical_keys_known_unregistered() {
+    let (_, context, account) = setup();
+
+    let historical_header = context.get_header();
+
+    let historical_keys = get_historical_public_keys(historical_header, account.address);
+    assert_eq(historical_keys, account.keys);
+    assert(context.max_block_number.is_none());
+}

aztec/src/keys/mod.nr

@@ -2,5 +2,6 @@ mod constants;
 mod getters;
 mod point_to_symmetric_key;
 mod public_keys;
+mod stored_keys;
 
 use crate::keys::public_keys::{PublicKeys, PUBLIC_KEYS_LENGTH};

aztec/src/keys/public_keys.nr

@@ -4,7 +4,7 @@ use dep::protocol_types::{
 };
 use crate::keys::constants::{NUM_KEY_TYPES, NULLIFIER_INDEX, INCOMING_INDEX, OUTGOING_INDEX};
 
-global PUBLIC_KEYS_LENGTH = 12;
+global PUBLIC_KEYS_LENGTH: u32 = 12;
 
 struct PublicKeys {
     npk_m: Point,