fix: change init and single claim nullif to incl owner address, add testing utilities (#24892)

Reimplementation of
#24837 - same
nullifier fix, this time extending `TestEnvironment` so that we can test
the fix works as intended.

Author: nventuro

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

@@ -12,8 +12,10 @@ use crate::{
 };
 
 use crate::protocol::{
-    address::AztecAddress, constants::DOM_SEP__INITIALIZATION_NULLIFIER, hash::poseidon2_hash_with_separator,
-    traits::Packable,
+    address::AztecAddress,
+    constants::DOM_SEP__INITIALIZATION_NULLIFIER,
+    hash::poseidon2_hash_with_separator,
+    traits::{Packable, ToField},
 };
 
 mod test;
@@ -67,7 +69,7 @@ impl<Note, Context> PrivateImmutable<Note, Context> {
     /// Computes the initialization nullifier using the provided secret.
     fn compute_initialization_nullifier(self, secret: Field) -> Field {
         poseidon2_hash_with_separator(
-            [self.storage_slot, secret],
+            [self.storage_slot, self.owner.to_field(), secret],
             DOM_SEP__INITIALIZATION_NULLIFIER,
         )
     }

noir-projects/aztec-nr/aztec/src/state_vars/private_immutable/test.nr

@@ -1,7 +1,8 @@
 use crate::{
     context::{PrivateContext, UtilityContext},
+    keys::getters::get_public_keys,
     state_vars::{OwnedStateVariable, private_immutable::PrivateImmutable},
-    test::{helpers::test_environment::TestEnvironment, mocks::mock_note::MockNote},
+    test::{helpers::test_environment::{CreateAccountOptions, TestEnvironment}, mocks::mock_note::MockNote},
 };
 use crate::protocol::address::AztecAddress;
 
@@ -197,6 +198,22 @@ unconstrained fn initialize_twice_same_tx() {
     });
 }
 
+#[test]
+unconstrained fn shared_keys_distinct_addresses_have_distinct_initialization_nullifiers() {
+    let mut env = TestEnvironment::new();
+
+    let owner_a = env.create_light_account_opts(CreateAccountOptions::new().with_secret(42));
+    let owner_b = env.create_light_account_opts(CreateAccountOptions::new().with_secret(42));
+
+    assert(owner_a != owner_b);
+    env.utility_context(|_| { assert_eq(get_public_keys(owner_a).npk_m_hash, get_public_keys(owner_b).npk_m_hash); });
+
+    let nullifier_a = env.private_context(|context| in_private(context, owner_a).get_initialization_nullifier());
+    let nullifier_b = env.private_context(|context| in_private(context, owner_b).get_initialization_nullifier());
+
+    assert(nullifier_a != nullifier_b);
+}
+
 #[test(should_fail_with = "already present")]
 unconstrained fn initialize_twice_other_tx() {
     let mut env = TestEnvironment::new();

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

@@ -12,8 +12,10 @@ use crate::{
 };
 
 use crate::protocol::{
-    address::AztecAddress, constants::DOM_SEP__INITIALIZATION_NULLIFIER, hash::poseidon2_hash_with_separator,
-    traits::Packable,
+    address::AztecAddress,
+    constants::DOM_SEP__INITIALIZATION_NULLIFIER,
+    hash::poseidon2_hash_with_separator,
+    traits::{Packable, ToField},
 };
 
 mod test;
@@ -86,7 +88,7 @@ impl<Note, Context> PrivateMutable<Note, Context> {
     /// Computes the initialization nullifier using the provided secret.
     fn compute_initialization_nullifier(self, secret: Field) -> Field {
         poseidon2_hash_with_separator(
-            [self.storage_slot, secret],
+            [self.storage_slot, self.owner.to_field(), secret],
             DOM_SEP__INITIALIZATION_NULLIFIER,
         )
     }

noir-projects/aztec-nr/aztec/src/state_vars/private_mutable/test.nr

@@ -1,7 +1,8 @@
 use crate::{
     context::{PrivateContext, UtilityContext},
+    keys::getters::get_public_keys,
     state_vars::{OwnedStateVariable, private_mutable::PrivateMutable},
-    test::{helpers::test_environment::TestEnvironment, mocks::mock_note::MockNote},
+    test::{helpers::test_environment::{CreateAccountOptions, TestEnvironment}, mocks::mock_note::MockNote},
 };
 use crate::protocol::address::AztecAddress;
 
@@ -339,6 +340,22 @@ unconstrained fn initialize_or_replace_initialized_settled() {
     });
 }
 
+#[test]
+unconstrained fn shared_keys_distinct_addresses_have_distinct_initialization_nullifiers() {
+    let mut env = TestEnvironment::new();
+
+    let owner_a = env.create_light_account_opts(CreateAccountOptions::new().with_secret(42));
+    let owner_b = env.create_light_account_opts(CreateAccountOptions::new().with_secret(42));
+
+    assert(owner_a != owner_b);
+    env.utility_context(|_| { assert_eq(get_public_keys(owner_a).npk_m_hash, get_public_keys(owner_b).npk_m_hash); });
+
+    let nullifier_a = env.private_context(|context| in_private(context, owner_a).get_initialization_nullifier());
+    let nullifier_b = env.private_context(|context| in_private(context, owner_b).get_initialization_nullifier());
+
+    assert(nullifier_a != nullifier_b);
+}
+
 #[test]
 unconstrained fn get_replacement_note_has_unique_randomness() {
     let mut env = TestEnvironment::new();

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

@@ -1,5 +1,6 @@
 use crate::protocol::{
     address::AztecAddress, constants::DOM_SEP__SINGLE_USE_CLAIM_NULLIFIER, hash::poseidon2_hash_with_separator,
+    traits::ToField,
 };
 
 use crate::{
@@ -82,7 +83,7 @@ impl<Context> SingleUseClaim<Context> {
     /// [`SingleUseClaim::assert_claimed`] and [`SingleUseClaim::has_claimed`] to coherently write and read state.
     fn compute_nullifier(self, owner_nhk_app: Field) -> Field {
         poseidon2_hash_with_separator(
-            [owner_nhk_app, self.storage_slot],
+            [owner_nhk_app, self.storage_slot, self.owner.to_field()],
             DOM_SEP__SINGLE_USE_CLAIM_NULLIFIER,
         )
     }

noir-projects/aztec-nr/aztec/src/state_vars/single_use_claim/test.nr

@@ -1,10 +1,11 @@
 use crate::{
     context::{PrivateContext, UtilityContext},
+    keys::getters::get_public_keys,
     oracle::random::random,
     state_vars::{OwnedStateVariable, SingleUseClaim},
 };
 use crate::protocol::{address::AztecAddress, traits::FromField};
-use crate::test::helpers::test_environment::TestEnvironment;
+use crate::test::helpers::test_environment::{CreateAccountOptions, TestEnvironment};
 
 global STORAGE_SLOT: Field = 17;
 
@@ -134,6 +135,24 @@ unconstrained fn claim_two_claim_ids_different_state_vars() {
     });
 }
 
+#[test]
+unconstrained fn claim_by_owner_does_not_claim_for_shared_key_distinct_address_owner() {
+    let mut env = TestEnvironment::new();
+
+    let owner_a = env.create_light_account_opts(CreateAccountOptions::new().with_secret(42));
+    let owner_b = env.create_light_account_opts(CreateAccountOptions::new().with_secret(42));
+
+    assert(owner_a != owner_b);
+    env.utility_context(|_| { assert_eq(get_public_keys(owner_a).npk_m_hash, get_public_keys(owner_b).npk_m_hash); });
+
+    env.private_context(|context| { in_private(context, owner_a).claim(); });
+
+    env.utility_context(|context| {
+        assert(in_utility(context, STORAGE_SLOT, owner_a).has_claimed());
+        assert(!in_utility(context, STORAGE_SLOT, owner_b).has_claimed());
+    });
+}
+
 #[test(should_fail_with = "Public keys not registered for account")]
 unconstrained fn claim_owner_unknown_public_keys() {
     let env = TestEnvironment::new();

noir-projects/aztec-nr/aztec/src/test/helpers/test_environment.nr

@@ -80,6 +80,9 @@ pub struct TestEnvironment {
     // results in TXE being able to maximize cache usage and not have to recompute account addresses and contract
     // artifacts, which are relatively expensive operations.
     light_account_secret: Counter,
+    // The partial address used for light account creation. Kept separate from the secret (which derives the keys) so
+    // that light accounts with the same keys but different addresses can be created via `create_light_account_opts`.
+    light_account_partial_address: Counter,
     contract_account_secret: Counter,
     contract_deployment_secret: Counter,
     contract_deployment_salt: Counter,
@@ -522,6 +525,44 @@ impl DeployOptions {
     }
 }
 
+/// Configuration values for [`TestEnvironment::create_light_account_opts`]. Meant to be used by calling `new` and then
+/// chaining methods setting each value, e.g.:
+/// ```noir
+/// env.create_light_account_opts(CreateAccountOptions::new().with_secret(42).with_partial_address(1));
+/// ```
+///
+/// The account address is derived from both the public keys (which come from the secret) and the partial address.
+/// Setting the same secret with different partial addresses produces accounts that share master keys but have distinct
+/// addresses.
+pub struct CreateAccountOptions {
+    secret: Option<Field>,
+    partial_address: Option<Field>,
+}
+
+impl CreateAccountOptions {
+    /// Creates a new `CreateAccountOptions` with default values, i.e. the same as if using the `create_light_account`
+    /// method instead of `create_light_account_opts`. Use the `with_secret` and `with_partial_address` methods to set
+    /// the desired configuration values.
+    pub fn new() -> Self {
+        Self { secret: Option::none(), partial_address: Option::none() }
+    }
+
+    /// Sets the secret used for key derivation. The secret affects the account's public keys and therefore its address:
+    /// accounts created with different secrets will have different keys and addresses.
+    pub fn with_secret(&mut self, secret: Field) -> Self {
+        self.secret = Option::some(secret);
+        *self
+    }
+
+    /// Sets the partial address used for address derivation. The partial address affects the resulting account address
+    /// but not its keys: accounts created with the same secret but different partial addresses share master keys while
+    /// having different addresses.
+    pub fn with_partial_address(&mut self, partial_address: Field) -> Self {
+        self.partial_address = Option::some(partial_address);
+        *self
+    }
+}
+
 /// Configuration values for [`TestEnvironment::call_public_opts`]. Meant to be used by calling `new` and then chaining
 /// methods setting each value, e.g.:
 /// ```noir
@@ -669,6 +710,10 @@ impl TestEnvironment {
             // secrets (e.g., Token1 with deployment secret=1 and liquidity_provider with account
             // secret=1), causing key collisions since keys are derived from the secret alone.
             light_account_secret: Counter::new(),
+            // The partial address counter does not need an offset for the same reason the salt counter doesn't: keys
+            // are derived from the secret alone, so a collision between a partial address and a secret value has no
+            // effect on key derivation.
+            light_account_partial_address: Counter::new(),
             contract_account_secret: Counter::new_with_offset(1_000),
             contract_deployment_secret: Counter::new_with_offset(2_000),
             // The salt counter does not need an offset because keys (derived from secret) and salt are unrelated: a
@@ -988,8 +1033,24 @@ impl TestEnvironment {
     /// Each call to `create_light_account` will return a different address, and so it can be called repeatedly to
     /// generate multiple addresses. These addresses are also different from the ones that `create_contract_account`
     /// returns, and so these two functions can be mixed and match to create a set of unique accounts.
+    ///
+    /// See [`create_light_account_opts`](TestEnvironment::create_light_account_opts) for a variant that allows
+    /// controlling the account's secret and partial address, e.g. to create accounts that share keys but have
+    /// different addresses.
     pub unconstrained fn create_light_account(&mut self) -> AztecAddress {
-        let test_account = txe_oracles::create_account(self.light_account_secret.next());
+        self.create_light_account_opts(CreateAccountOptions::new())
+    }
+
+    /// Variant of `create_light_account` which allows specifying configuration values via `CreateAccountOptions`, such
+    /// as a custom secret or partial address. If not specified, the secret and partial address default to values from
+    /// internal counters.
+    ///
+    /// Setting the same secret with different partial addresses produces accounts that share master keys but have
+    /// distinct addresses.
+    pub unconstrained fn create_light_account_opts(&mut self, opts: CreateAccountOptions) -> AztecAddress {
+        let secret = opts.secret.unwrap_or_else(|| self.light_account_secret.next());
+        let partial_address = opts.partial_address.unwrap_or_else(|| self.light_account_partial_address.next());
+        let test_account = txe_oracles::create_account(secret, partial_address);
         test_account.address
     }
 

noir-projects/aztec-nr/aztec/src/test/helpers/test_environment/test/accounts.nr

@@ -2,7 +2,7 @@ use crate::{
     authwit::auth::{compute_authwit_message_hash, IS_VALID_SELECTOR},
     context::calls::PrivateStaticCall,
     keys::getters::get_public_keys,
-    test::helpers::{test_environment::TestEnvironment, txe_oracles},
+    test::helpers::{test_environment::{CreateAccountOptions, TestEnvironment}, txe_oracles},
 };
 use crate::protocol::abis::function_selector::FunctionSelector;
 
@@ -16,6 +16,45 @@ unconstrained fn create_light_account_does_not_repeat_accounts() {
     assert(first_account != second_account);
 }
 
+#[test]
+unconstrained fn create_light_account_opts_same_secret_and_partial_address_repeats_address() {
+    let mut env = TestEnvironment::new();
+
+    let first_account =
+        env.create_light_account_opts(CreateAccountOptions::new().with_secret(42).with_partial_address(1));
+    let second_account =
+        env.create_light_account_opts(CreateAccountOptions::new().with_secret(42).with_partial_address(1));
+
+    assert_eq(first_account, second_account);
+}
+
+#[test]
+unconstrained fn create_light_account_opts_shared_secret_distinct_partial_address_shares_keys_not_address() {
+    let mut env = TestEnvironment::new();
+
+    let first_account =
+        env.create_light_account_opts(CreateAccountOptions::new().with_secret(42).with_partial_address(1));
+    let second_account =
+        env.create_light_account_opts(CreateAccountOptions::new().with_secret(42).with_partial_address(2));
+
+    assert(first_account != second_account);
+    env.utility_context(|_| {
+        assert_eq(get_public_keys(first_account).npk_m_hash, get_public_keys(second_account).npk_m_hash);
+    });
+}
+
+#[test]
+unconstrained fn create_light_account_opts_distinct_secret_shared_partial_address_differ() {
+    let mut env = TestEnvironment::new();
+
+    let first_account =
+        env.create_light_account_opts(CreateAccountOptions::new().with_secret(42).with_partial_address(1));
+    let second_account =
+        env.create_light_account_opts(CreateAccountOptions::new().with_secret(43).with_partial_address(1));
+
+    assert(first_account != second_account);
+}
+
 #[test]
 unconstrained fn create_contract_account_does_not_repeat_accounts() {
     let mut env = TestEnvironment::new();

noir-projects/aztec-nr/aztec/src/test/helpers/txe_oracles.nr

@@ -21,8 +21,8 @@ use crate::protocol::{
 /// [`oracle::version`](crate::oracle::version), which covers oracles used during contract execution by PXE.
 ///
 /// The TypeScript counterparts are in `yarn-project/txe/src/txe_oracle_version.ts`.
-pub global TXE_ORACLE_VERSION_MAJOR: Field = 3;
-pub global TXE_ORACLE_VERSION_MINOR: Field = 1;
+pub global TXE_ORACLE_VERSION_MAJOR: Field = 4;
+pub global TXE_ORACLE_VERSION_MINOR: Field = 0;
 
 /// Asserts that the TXE oracle interface version is compatible.
 pub unconstrained fn assert_compatible_txe_oracle_version() {
@@ -249,7 +249,7 @@ pub unconstrained fn deploy_oracle<let N: u32, let P: u32>(
 ) -> [Field; CONTRACT_INSTANCE_LENGTH] {}
 
 #[oracle(aztec_txe_createAccount)]
-pub unconstrained fn create_account(secret: Field) -> TestAccount {}
+pub unconstrained fn create_account(secret: Field, partial_address: Field) -> TestAccount {}
 
 #[oracle(aztec_txe_addAccount)]
 pub unconstrained fn add_account(secret: Field) -> TestAccount {}

noir-projects/noir-protocol-circuits/crates/types/src/constants_tests.nr

@@ -13,9 +13,9 @@ use crate::{
         DOM_SEP__CONSTRAINED_MSG_NULLIFIER, DOM_SEP__CONSTRAINED_MSG_SENDER_SECRET,
         DOM_SEP__CONTRACT_ADDRESS_V2, DOM_SEP__CONTRACT_CLASS_ID, DOM_SEP__ECDH_FIELD_MASK,
         DOM_SEP__ECDH_SUBKEY, DOM_SEP__EVENT_COMMITMENT, DOM_SEP__EVENT_LOG_TAG, DOM_SEP__FBSK_M,
-        DOM_SEP__FUNCTION_ARGS, DOM_SEP__INITIALIZATION_NULLIFIER, DOM_SEP__INITIALIZER,
-        DOM_SEP__HANDSHAKE_FORGERY_PROTECTION, DOM_SEP__INTERACTIVE_HANDSHAKE_SIGNATURE,
-        DOM_SEP__IVSK_M, DOM_SEP__MERKLE_HASH,
+        DOM_SEP__FUNCTION_ARGS, DOM_SEP__HANDSHAKE_FORGERY_PROTECTION,
+        DOM_SEP__INITIALIZATION_NULLIFIER, DOM_SEP__INITIALIZER,
+        DOM_SEP__INTERACTIVE_HANDSHAKE_SIGNATURE, DOM_SEP__IVSK_M, DOM_SEP__MERKLE_HASH,
         DOM_SEP__MESSAGE_NULLIFIER, DOM_SEP__MSSK_M, DOM_SEP__NHK_M,
         DOM_SEP__NON_INTERACTIVE_HANDSHAKE_LOG_TAG, DOM_SEP__NOTE_COMPLETION_LOG_TAG,
         DOM_SEP__NOTE_HASH, DOM_SEP__NOTE_HASH_NONCE, DOM_SEP__NOTE_NULLIFIER,