Validate new public key in Account (#989)

* feat: add logic and tests * feat: update CHANGELOG * feat: use ExTrait * feat: update doc example * Update src/tests/account/test_account.cairo Co-authored-by: Andrew Fleming <fleming.andrew@protonmail.com> * Update src/tests/account/test_account.cairo Co-authored-by: Andrew Fleming <fleming.andrew@protonmail.com> * Update docs/modules/ROOT/pages/api/account.adoc Co-authored-by: Andrew Fleming <fleming.andrew@protonmail.com> * Update docs/modules/ROOT/pages/api/account.adoc Co-authored-by: Andrew Fleming <fleming.andrew@protonmail.com> * Update docs/modules/ROOT/pages/api/account.adoc Co-authored-by: Andrew Fleming <fleming.andrew@protonmail.com> * Update src/account/account.cairo Co-authored-by: Andrew Fleming <fleming.andrew@protonmail.com> * Update src/account/account.cairo Co-authored-by: Andrew Fleming <fleming.andrew@protonmail.com> * feat: move NEW_PUBKEY to constants * feat: format files --------- Co-authored-by: Andrew Fleming <fleming.andrew@protonmail.com>
OpenZeppelin · May 16, 2024 · 852abcf · 852abcf
1 parent ad6b858
commit 852abcf
Show file tree

Hide file tree

Showing 11 changed files with 251 additions and 59 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -17,8 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed (Breaking)
 
 - ERC721Component internal implementation to support transfer, mint, and burn flows going through an `_update` function (#978)
-- ERC721Component implementations now require an ERC721HooksTrait implementation in scope.
-- ERC1155Component implementations now require an ERC1155HooksTrait implementation in scope.
+- ERC721Component implementations now require an ERC721HooksTrait implementation in scope (#978)
+- ERC1155Component implementations now require an ERC1155HooksTrait implementation in scope (#982)
+- AccountComponent, preset, and dispatcher now require a `signature` param in the public-key-setter functions (#989)
 
 ## 0.12.0 (2024-04-21)
 

diff --git a/docs/modules/ROOT/pages/api/account.adoc b/docs/modules/ROOT/pages/api/account.adoc
@@ -117,7 +117,7 @@ NOTE: {src5-component-required-note}
 .PublicKeyImpl
 
 * xref:#AccountComponent-get_public_key[`++get_public_key(self)++`]
-* xref:#AccountComponent-set_public_key[`++set_public_key(self, new_public_key)++`]
+* xref:#AccountComponent-set_public_key[`++set_public_key(self, new_public_key, signature)++`]
 
 [.sub-index#AccountComponent-Embeddable-Impls-SRC6CamelOnlyImpl]
 .SRC6CamelOnlyImpl
@@ -128,7 +128,7 @@ NOTE: {src5-component-required-note}
 .PublicKeyCamelImpl
 
 * xref:#AccountComponent-getPublicKey[`++getPublicKey(self)++`]
-* xref:#AccountComponent-setPublicKey[`++setPublicKey(self, newPublicKey)++`]
+* xref:#AccountComponent-setPublicKey[`++setPublicKey(self, newPublicKey, signature)++`]
 
 .SRC5Impl
 * xref:api/introspection.adoc#ISRC5-supports_interface[`supports_interface(self, interface_id: felt252)`]
@@ -141,6 +141,7 @@ NOTE: {src5-component-required-note}
 
 * xref:#AccountComponent-initializer[`++initializer(self, public_key)++`]
 * xref:#AccountComponent-assert_only_self[`++assert_only_self(self)++`]
+* xref:#AccountComponent-assert_valid_new_owner[`++assert_valid_new_owner(self, current_owner, new_owner, signature)++`]
 * xref:#AccountComponent-validate_transaction[`++validate_transaction(self)++`]
 * xref:#AccountComponent-_set_public_key[`++_set_public_key(self, new_public_key)++`]
 * xref:#AccountComponent-_is_valid_signature[`++_is_valid_signature(self, hash, signature)++`]
@@ -199,12 +200,31 @@ Returns the current public key of the account.
 
 [.contract-item]
 [[AccountComponent-set_public_key]]
-==== `[.contract-item-name]#++set_public_key++#++(ref self: ContractState, new_public_key: felt252)++` [.item-kind]#external#
+==== `[.contract-item-name]#++set_public_key++#++(ref self: ContractState, new_public_key: felt252, signature: Span<felt252>)++` [.item-kind]#external#
 
 Sets a new public key for the account. Only accessible by the account calling itself through `\\__execute__`.
 
+Requirements:
+
+- The caller must be the contract itself.
+- The signature must be valid for the new owner.
+
 Emits both an {OwnerRemoved} and an {OwnerAdded} event.
 
+[NOTE]
+====
+The message to be signed is computed in Cairo as follows:
+
+```javascript
+let message_hash = PoseidonTrait::new()
+        .update_with('StarkNet Message')
+        .update_with('accept_ownership')
+        .update_with(get_contract_address())
+        .update_with(current_owner)
+        .finalize();
+```
+====
+
 [.contract-item]
 [[AccountComponent-isValidSignature]]
 ==== `[.contract-item-name]#++isValidSignature++#++(self: @ContractState, hash: felt252, signature: Array<felt252>) → felt252++` [.item-kind]#external#
@@ -219,7 +239,7 @@ See xref:AccountComponent-get_public_key[get_public_key].
 
 [.contract-item]
 [[AccountComponent-setPublicKey]]
-==== `[.contract-item-name]#++setPublicKey++#++(ref self: ContractState, newPublicKey: felt252)++` [.item-kind]#external#
+==== `[.contract-item-name]#++setPublicKey++#++(ref self: ContractState, newPublicKey: felt252, signature: Span<felt252>)++` [.item-kind]#external#
 
 See xref:AccountComponent-set_public_key[set_public_key].
 
@@ -240,6 +260,19 @@ Emits an {OwnerAdded} event.
 
 Validates that the caller is the account itself. Otherwise it reverts.
 
+[.contract-item]
+[[AccountComponent-assert_valid_new_owner]]
+==== `[.contract-item-name]#++assert_valid_new_owner++#++(self: @ComponentState, current_owner: felt252, new_owner: felt252, signature: Span<felt252>)++` [.item-kind]#internal#
+
+Validates that `new_owner` accepted the ownership of the contract through a signature.
+
+Requirements:
+
+- `signature` must be valid for the new owner.
+
+WARNING: This function assumes that `current_owner` is the current owner of the contract, and
+does not validate this assumption.
+
 [.contract-item]
 [[AccountComponent-validate_transaction]]
 ==== `[.contract-item-name]#++validate_transaction++#++(self: @ComponentState)++ → felt252` [.item-kind]#internal#

diff --git a/src/account/account.cairo b/src/account/account.cairo
@@ -6,12 +6,14 @@
 /// The Account component enables contracts to behave as accounts.
 #[starknet::component]
 mod AccountComponent {
+    use core::hash::{HashStateExTrait, HashStateTrait};
     use openzeppelin::account::interface;
     use openzeppelin::account::utils::{MIN_TRANSACTION_VERSION, QUERY_VERSION, QUERY_OFFSET};
     use openzeppelin::account::utils::{execute_calls, is_valid_stark_signature};
     use openzeppelin::introspection::src5::SRC5Component::InternalTrait as SRC5InternalTrait;
     use openzeppelin::introspection::src5::SRC5Component::SRC5;
     use openzeppelin::introspection::src5::SRC5Component;
+    use poseidon::PoseidonTrait;
     use starknet::account::Call;
     use starknet::get_caller_address;
     use starknet::get_contract_address;
@@ -155,11 +157,20 @@ mod AccountComponent {
         /// Requirements:
         ///
         /// - The caller must be the contract itself.
+        /// - The signature must be valid for the new owner.
         ///
-        /// Emits an `OwnerRemoved` event.
-        fn set_public_key(ref self: ComponentState<TContractState>, new_public_key: felt252) {
+        /// Emits both an `OwnerRemoved` and an `OwnerAdded` event.
+        fn set_public_key(
+            ref self: ComponentState<TContractState>,
+            new_public_key: felt252,
+            signature: Span<felt252>
+        ) {
             self.assert_only_self();
-            self.emit(OwnerRemoved { removed_owner_guid: self.Account_public_key.read() });
+
+            let current_owner = self.Account_public_key.read();
+            self.assert_valid_new_owner(current_owner, new_public_key, signature);
+
+            self.emit(OwnerRemoved { removed_owner_guid: current_owner });
             self._set_public_key(new_public_key);
         }
     }
@@ -191,8 +202,12 @@ mod AccountComponent {
             self.Account_public_key.read()
         }
 
-        fn setPublicKey(ref self: ComponentState<TContractState>, newPublicKey: felt252) {
-            PublicKey::set_public_key(ref self, newPublicKey);
+        fn setPublicKey(
+            ref self: ComponentState<TContractState>,
+            newPublicKey: felt252,
+            signature: Span<felt252>
+        ) {
+            PublicKey::set_public_key(ref self, newPublicKey, signature);
         }
     }
 
@@ -218,6 +233,31 @@ mod AccountComponent {
             assert(self == caller, Errors::UNAUTHORIZED);
         }
 
+        /// Validates that `new_owner` accepted the ownership of the contract.
+        ///
+        /// WARNING: This function assumes that `current_owner` is the current owner of the contract, and
+        /// does not validate this assumption.
+        ///
+        /// Requirements:
+        ///
+        /// - The signature must be valid for the new owner.
+        fn assert_valid_new_owner(
+            self: @ComponentState<TContractState>,
+            current_owner: felt252,
+            new_owner: felt252,
+            signature: Span<felt252>
+        ) {
+            let message_hash = PoseidonTrait::new()
+                .update_with('StarkNet Message')
+                .update_with('accept_ownership')
+                .update_with(get_contract_address())
+                .update_with(current_owner)
+                .finalize();
+
+            let is_valid = is_valid_stark_signature(message_hash, new_owner, signature);
+            assert(is_valid, Errors::INVALID_SIGNATURE);
+        }
+
         /// Validates the signature for the current transaction.
         /// Returns the short string `VALID` if valid, otherwise it reverts.
         fn validate_transaction(self: @ComponentState<TContractState>) -> felt252 {
@@ -300,17 +340,25 @@ mod AccountComponent {
             PublicKey::get_public_key(self)
         }
 
-        fn set_public_key(ref self: ComponentState<TContractState>, new_public_key: felt252) {
-            PublicKey::set_public_key(ref self, new_public_key);
+        fn set_public_key(
+            ref self: ComponentState<TContractState>,
+            new_public_key: felt252,
+            signature: Span<felt252>
+        ) {
+            PublicKey::set_public_key(ref self, new_public_key, signature);
         }
 
         // IPublicKeyCamel
         fn getPublicKey(self: @ComponentState<TContractState>) -> felt252 {
             PublicKeyCamel::getPublicKey(self)
         }
 
-        fn setPublicKey(ref self: ComponentState<TContractState>, newPublicKey: felt252) {
-            PublicKeyCamel::setPublicKey(ref self, newPublicKey);
+        fn setPublicKey(
+            ref self: ComponentState<TContractState>,
+            newPublicKey: felt252,
+            signature: Span<felt252>
+        ) {
+            PublicKeyCamel::setPublicKey(ref self, newPublicKey, signature);
         }
 
         // ISRC5

diff --git a/src/account/dual_account.cairo b/src/account/dual_account.cairo
@@ -15,7 +15,7 @@ struct DualCaseAccount {
 }
 
 trait DualCaseAccountABI {
-    fn set_public_key(self: @DualCaseAccount, new_public_key: felt252);
+    fn set_public_key(self: @DualCaseAccount, new_public_key: felt252, signature: Span<felt252>);
     fn get_public_key(self: @DualCaseAccount) -> felt252;
     fn is_valid_signature(
         self: @DualCaseAccount, hash: felt252, signature: Array<felt252>
@@ -24,8 +24,9 @@ trait DualCaseAccountABI {
 }
 
 impl DualCaseAccountImpl of DualCaseAccountABI {
-    fn set_public_key(self: @DualCaseAccount, new_public_key: felt252) {
-        let args = array![new_public_key];
+    fn set_public_key(self: @DualCaseAccount, new_public_key: felt252, signature: Span<felt252>) {
+        let mut args = array![new_public_key];
+        args.append_serde(signature);
 
         try_selector_with_fallback(
             *self.contract_address, selectors::set_public_key, selectors::setPublicKey, args.span()

diff --git a/src/account/interface.cairo b/src/account/interface.cairo
@@ -35,7 +35,7 @@ trait IDeployable<TState> {
 #[starknet::interface]
 trait IPublicKey<TState> {
     fn get_public_key(self: @TState) -> felt252;
-    fn set_public_key(ref self: TState, new_public_key: felt252);
+    fn set_public_key(ref self: TState, new_public_key: felt252, signature: Span<felt252>);
 }
 
 #[starknet::interface]
@@ -46,7 +46,7 @@ trait ISRC6CamelOnly<TState> {
 #[starknet::interface]
 trait IPublicKeyCamel<TState> {
     fn getPublicKey(self: @TState) -> felt252;
-    fn setPublicKey(ref self: TState, newPublicKey: felt252);
+    fn setPublicKey(ref self: TState, newPublicKey: felt252, signature: Span<felt252>);
 }
 
 //
@@ -73,14 +73,14 @@ trait AccountABI<TState> {
 
     // IPublicKey
     fn get_public_key(self: @TState) -> felt252;
-    fn set_public_key(ref self: TState, new_public_key: felt252);
+    fn set_public_key(ref self: TState, new_public_key: felt252, signature: Span<felt252>);
 
     // ISRC6CamelOnly
     fn isValidSignature(self: @TState, hash: felt252, signature: Array<felt252>) -> felt252;
 
     // IPublicKeyCamel
     fn getPublicKey(self: @TState) -> felt252;
-    fn setPublicKey(ref self: TState, newPublicKey: felt252);
+    fn setPublicKey(ref self: TState, newPublicKey: felt252, signature: Span<felt252>);
 }
 
 //

diff --git a/src/presets/interfaces/account.cairo b/src/presets/interfaces/account.cairo
@@ -21,14 +21,14 @@ trait AccountUpgradeableABI<TState> {
 
     // IPublicKey
     fn get_public_key(self: @TState) -> felt252;
-    fn set_public_key(ref self: TState, new_public_key: felt252);
+    fn set_public_key(ref self: TState, new_public_key: felt252, signature: Span<felt252>);
 
     // ISRC6CamelOnly
     fn isValidSignature(self: @TState, hash: felt252, signature: Array<felt252>) -> felt252;
 
     // IPublicKeyCamel
     fn getPublicKey(self: @TState) -> felt252;
-    fn setPublicKey(ref self: TState, newPublicKey: felt252);
+    fn setPublicKey(ref self: TState, newPublicKey: felt252, signature: Span<felt252>);
 
     // IUpgradeable
     fn upgrade(ref self: TState, new_class_hash: ClassHash);