From 062c06ac2e58d88334e59448df09146f1e0f2716 Mon Sep 17 00:00:00 2001
From: rymnc <43716372+rymnc@users.noreply.github.com>
Date: Tue, 23 Jul 2024 17:31:46 +0530
Subject: [PATCH] fix: add user address instead of defaulting to msg.sender

---
 contracts/src/IScKeystore.sol   |  6 +++++-
 contracts/src/ScKeystore.sol    | 12 ++++++------
 contracts/test/ScKeystore.t.sol |  4 ++--
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/contracts/src/IScKeystore.sol b/contracts/src/IScKeystore.sol
index ccc1449..5e2e4f0 100644
--- a/contracts/src/IScKeystore.sol
+++ b/contracts/src/IScKeystore.sol
@@ -13,8 +13,12 @@ struct UserInfo {
 
 interface IScKeystore {
     function userExists(address user) external view returns (bool);
-    function addUser(bytes calldata signaturePubKey, KeyPackage calldata keyPackage) external;
+
+    function addUser(address user, bytes calldata signaturePubKey, KeyPackage calldata keyPackage) external;
+
     function getUser(address user) external view returns (UserInfo memory);
+
     function addKeyPackage(KeyPackage calldata) external;
+
     function getAvailableKeyPackage(address user) external view returns (KeyPackage memory);
 }
diff --git a/contracts/src/ScKeystore.sol b/contracts/src/ScKeystore.sol
index 96baa5a..5f451c2 100644
--- a/contracts/src/ScKeystore.sol
+++ b/contracts/src/ScKeystore.sol
@@ -22,19 +22,19 @@ contract ScKeystore is Ownable, IScKeystore {
         return users[user].signaturePubKey.length > 0;
     }
 
-    function addUser(bytes calldata signaturePubKey, KeyPackage calldata keyPackage) external onlyOwner {
+    function addUser(address user, bytes calldata signaturePubKey, KeyPackage calldata keyPackage) external onlyOwner {
         if (signaturePubKey.length == 0) revert MalformedUserInfo();
         if (keyPackage.data.length == 0) revert MalformedKeyPackage();
-        if (userExists(msg.sender)) revert UserAlreadyExists();
+        if (userExists(user)) revert UserAlreadyExists();
 
         keyPackages.push(keyPackage);
         uint256 keyPackageIndex = keyPackages.length - 1;
 
-        users[msg.sender] = UserInfo(new uint256[](0), signaturePubKey);
-        users[msg.sender].signaturePubKey = signaturePubKey;
-        users[msg.sender].keyPackageIndices.push(keyPackageIndex);
+        users[user] = UserInfo(new uint256[](0), signaturePubKey);
+        users[user].signaturePubKey = signaturePubKey;
+        users[user].keyPackageIndices.push(keyPackageIndex);
 
-        emit UserAdded(msg.sender, signaturePubKey);
+        emit UserAdded(user, signaturePubKey);
     }
 
     function getUser(address user) external view returns (UserInfo memory) {
diff --git a/contracts/test/ScKeystore.t.sol b/contracts/test/ScKeystore.t.sol
index 8980a76..d79c57b 100644
--- a/contracts/test/ScKeystore.t.sol
+++ b/contracts/test/ScKeystore.t.sol
@@ -19,7 +19,7 @@ contract ScKeystoreTest is Test {
 
     function addUser() internal {
         KeyPackage memory keyPackage = KeyPackage({ data: new bytes[](1) });
-        s.addUser("0x", keyPackage);
+        s.addUser(address(this), "0x", keyPackage);
     }
 
     function test__owner() public view {
@@ -32,7 +32,7 @@ contract ScKeystoreTest is Test {
 
     function test__addUser__reverts__whenUserInfoIsMalformed() public {
         vm.expectRevert(MalformedUserInfo.selector);
-        s.addUser("", KeyPackage({ data: new bytes[](0) }));
+        s.addUser(address(this), "", KeyPackage({ data: new bytes[](0) }));
     }
 
     function test__addUser__reverts__whenUserAlreadyExists() public {