fixed tests #2

rarimo · Jun 11, 2024 · 064da3b · 064da3b
1 parent f143496
commit 064da3b
Show file tree

Hide file tree

Showing 8 changed files with 182 additions and 167 deletions.
diff --git a/contracts/registration/Registration.sol b/contracts/registration/Registration.sol
@@ -200,17 +200,6 @@ contract Registration is Initializable, TSSUpgradeable {
         stateKeeper.reissueBondIdentity(bytes32(passportKey_), bytes32(identityKey_), dgCommit_);
     }
 
-    /**
-     * @notice Change the Rarimo TSS signer via Rarimo TSS
-     * @param newSignerPubKey_ the new signer public key
-     * @param signature_ the Rarimo TSS signature
-     */
-    function changeSigner(bytes memory newSignerPubKey_, bytes memory signature_) external {
-        _checkSignature(keccak256(newSignerPubKey_), signature_);
-
-        signer = _convertPubKeyToAddress(newSignerPubKey_);
-    }
-
     /**
      * @notice Adds or removes a dispatcher via Rarimo TSS
      * @param methodId_ the method id (AddDispatcher or RemoveDispatcher)

diff --git a/contracts/state/PoseidonSMT.sol b/contracts/state/PoseidonSMT.sol
@@ -55,17 +55,6 @@ contract PoseidonSMT is Initializable, TSSUpgradeable {
         stateKeeper = stateKeeper_;
     }
 
-    /**
-     * @notice Change the Rarimo TSS signer via Rarimo TSS
-     * @param newSignerPubKey_ the new signer public key
-     * @param signature_ the Rarimo TSS signature
-     */
-    function changeSigner(bytes memory newSignerPubKey_, bytes memory signature_) external {
-        _checkSignature(keccak256(newSignerPubKey_), signature_);
-
-        signer = _convertPubKeyToAddress(newSignerPubKey_);
-    }
-
     /**
      * @notice Adds the new element to the tree.
      */

diff --git a/contracts/state/StateKeeper.sol b/contracts/state/StateKeeper.sol
@@ -222,17 +222,6 @@ contract StateKeeper is Initializable, TSSUpgradeable {
         icaoMasterTreeMerkleRoot = newRoot_;
     }
 
-    /**
-     * @notice Change the Rarimo TSS signer via Rarimo TSS
-     * @param newSignerPubKey_ the new signer public key
-     * @param signature_ the Rarimo TSS signature
-     */
-    function changeSigner(bytes memory newSignerPubKey_, bytes memory signature_) external {
-        _checkSignature(keccak256(newSignerPubKey_), signature_);
-
-        signer = _convertPubKeyToAddress(newSignerPubKey_);
-    }
-
     /**
      * @notice Add or Remove registrations via Rarimo TSS
      * @param methodId_ the method id (AddRegistrations or RemoveRegistrations)
@@ -324,6 +313,10 @@ contract StateKeeper is Initializable, TSSUpgradeable {
         return _registrations[key_];
     }
 
+    function isRegistration(address registration_) external view returns (bool) {
+        return _registrationExists[registration_];
+    }
+
     function _onlyRegistration() internal view {
         require(_registrationExists[msg.sender], "StateKeeper: not a registration");
     }

diff --git a/contracts/state/TSSSigner.sol b/contracts/state/TSSSigner.sol
@@ -21,6 +21,17 @@ abstract contract TSSSigner {
         chainName = chainName_;
     }
 
+    /**
+     * @notice Change the Rarimo TSS signer via Rarimo TSS
+     * @param newSignerPubKey_ the new signer public key
+     * @param signature_ the Rarimo TSS signature
+     */
+    function changeSigner(bytes memory newSignerPubKey_, bytes memory signature_) external {
+        _checkSignature(keccak256(newSignerPubKey_), signature_);
+
+        signer = _convertPubKeyToAddress(newSignerPubKey_);
+    }
+
     function getNonce(uint8 methodId_) external view returns (uint256) {
         return _nonces[methodId_];
     }

diff --git a/test/helpers/TSSMerkleTree.ts b/test/helpers/TSSMerkleTree.ts
@@ -62,15 +62,15 @@ export class TSSMerkleTree {
     operationType: RegistrationMethodId.AddPassportDispatcher | RegistrationMethodId.AddCertificateDispatcher,
     dispatcherType: string,
     dispatcher: string,
-    chaneName: string,
+    chainName: string,
     nonce: BigNumberish,
     contractAddress: string,
     anotherSigner: HDNodeWallet | undefined = undefined,
   ): TSSOperation {
     const encoder = new ethers.AbiCoder();
     const data = encoder.encode(["bytes32", "address"], [dispatcherType, dispatcher]);
 
-    const hash = this.getArbitraryDataSignHash(operationType, data, chaneName, nonce, contractAddress);
+    const hash = this.getArbitraryDataSignHash(operationType, data, chainName, nonce, contractAddress);
 
     return {
       data,
@@ -81,15 +81,15 @@ export class TSSMerkleTree {
   public removeDispatcherOperation(
     operationType: RegistrationMethodId.RemovePassportDispatcher | RegistrationMethodId.RemoveCertificateDispatcher,
     dispatcherType: string,
-    chaneName: string,
+    chainName: string,
     nonce: BigNumberish,
     contractAddress: string,
     anotherSigner: HDNodeWallet | undefined = undefined,
   ): TSSOperation {
     const encoder = new ethers.AbiCoder();
     const data = encoder.encode(["bytes32"], [dispatcherType]);
 
-    const hash = this.getArbitraryDataSignHash(operationType, data, chaneName, nonce, contractAddress);
+    const hash = this.getArbitraryDataSignHash(operationType, data, chainName, nonce, contractAddress);
 
     return {
       data,
@@ -98,19 +98,20 @@ export class TSSMerkleTree {
   }
 
   public addRegistrationsOperation(
+    registrationKeys: string[],
     registrations: string[],
-    chaneName: string,
+    chainName: string,
     nonce: BigNumberish,
     contractAddress: string,
     anotherSigner: HDNodeWallet | undefined = undefined,
   ): TSSOperation {
     const encoder = new ethers.AbiCoder();
-    const data = encoder.encode(["address[]"], [registrations]);
+    const data = encoder.encode(["string[]", "address[]"], [registrationKeys, registrations]);
 
     const hash = this.getArbitraryDataSignHash(
       StateKeeperMethodId.AddRegistrations,
       data,
-      chaneName,
+      chainName,
       nonce,
       contractAddress,
     );
@@ -122,19 +123,19 @@ export class TSSMerkleTree {
   }
 
   public removeRegistrationsOperation(
-    registrations: string[],
-    chaneName: string,
+    registrationKeys: string[],
+    chainName: string,
     nonce: BigNumberish,
     contractAddress: string,
     anotherSigner: HDNodeWallet | undefined = undefined,
   ): TSSOperation {
     const encoder = new ethers.AbiCoder();
-    const data = encoder.encode(["address[]"], [registrations]);
+    const data = encoder.encode(["string[]"], [registrationKeys]);
 
     const hash = this.getArbitraryDataSignHash(
       StateKeeperMethodId.RemoveRegistrations,
       data,
-      chaneName,
+      chainName,
       nonce,
       contractAddress,
     );
@@ -147,14 +148,14 @@ export class TSSMerkleTree {
 
   public authorizeUpgradeOperation(
     newImplementation: string,
-    chaneName: string,
+    chainName: string,
     nonce: BigNumberish,
     contractAddress: string,
     anotherSigner: HDNodeWallet | undefined = undefined,
   ): string {
     const hash = ethers.solidityPackedKeccak256(
       ["address", "uint8", "address", "string", "uint256"],
-      [contractAddress, TSSUpgradeableId.MAGIC_ID, newImplementation, chaneName, nonce],
+      [contractAddress, TSSUpgradeableId.MAGIC_ID, newImplementation, chainName, nonce],
     );
 
     return this.getProof(hash, true, anotherSigner);
@@ -163,13 +164,13 @@ export class TSSMerkleTree {
   public getArbitraryDataSignHash(
     methodId: number,
     data: string,
-    chaneName: string,
+    chainName: string,
     nonce: BigNumberish,
     contractAddress: string,
   ): string {
     return ethers.solidityPackedKeccak256(
       ["address", "uint8", "bytes", "string", "uint256"],
-      [contractAddress, methodId, data, chaneName, nonce],
+      [contractAddress, methodId, data, chainName, nonce],
     );
   }
 }
diff --git a/test/registration/Registration.test.ts b/test/registration/Registration.test.ts
@@ -165,6 +165,7 @@ describe("Registration", () => {
       },
     });
     const Registration = await ethers.getContractFactory("RegistrationMock");
+    const Proxy = await ethers.getContractFactory("ERC1967Proxy");
 
     registrationSmt = await PoseidonSMT.deploy();
     certificatesSmt = await PoseidonSMT.deploy();
@@ -175,21 +176,21 @@ describe("Registration", () => {
     await deployPRSASHA1Dispatcher();
     await deployPECDSASHA1Dispatcher();
 
-    const Proxy = await ethers.getContractFactory("ERC1967Proxy");
     let proxy = await Proxy.deploy(await stateKeeper.getAddress(), "0x");
     stateKeeper = stateKeeper.attach(await proxy.getAddress()) as StateKeeperMock;
 
     proxy = await Proxy.deploy(await registrationSmt.getAddress(), "0x");
     registrationSmt = registrationSmt.attach(await proxy.getAddress()) as PoseidonSMTMock;
-    await registrationSmt.__PoseidonSMT_init(SIGNER.address, chainName, await stateKeeper.getAddress(), treeSize);
 
     proxy = await Proxy.deploy(await certificatesSmt.getAddress(), "0x");
     certificatesSmt = certificatesSmt.attach(await proxy.getAddress()) as PoseidonSMTMock;
-    await certificatesSmt.__PoseidonSMT_init(SIGNER.address, chainName, await stateKeeper.getAddress(), treeSize);
 
     proxy = await Proxy.deploy(await registration.getAddress(), "0x");
     registration = registration.attach(await proxy.getAddress()) as RegistrationMock;
 
+    await registrationSmt.__PoseidonSMT_init(SIGNER.address, chainName, await stateKeeper.getAddress(), treeSize);
+    await certificatesSmt.__PoseidonSMT_init(SIGNER.address, chainName, await stateKeeper.getAddress(), treeSize);
+
     await stateKeeper.__StateKeeper_init(
       SIGNER.address,
       chainName,

diff --git a/test/state/PoseidonSMT.test.ts b/test/state/PoseidonSMT.test.ts
@@ -7,7 +7,6 @@ import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";
 import { ERC1967Proxy__factory, PoseidonSMT } from "@ethers-v6";
 
 import { getPoseidon, Reverter, TSSMerkleTree, TSSSigner } from "@/test/helpers";
-import { TSSUpgradeableId } from "@/test/helpers/constants";
 
 const treeSize = 80;
 const chainName = "Tests";
@@ -61,79 +60,4 @@ describe("PoseidonSMT", () => {
       });
     });
   });
-
-  describe("$TSS flow", () => {
-    describe("#changeSigner", () => {
-      const newSigner = ethers.Wallet.createRandom();
-      const tssPublicKey = "0x" + newSigner.signingKey.publicKey.slice(4);
-
-      it("should change signer if signature and new public key are valid", async () => {
-        expect(await tree.getFunction("signer").staticCall()).to.eq(SIGNER.address);
-
-        const signature = signHelper.signChangeSigner(tssPublicKey);
-
-        await tree.changeSigner(tssPublicKey, signature);
-
-        expect(await tree.getFunction("signer").staticCall()).to.eq(newSigner.address);
-      });
-    });
-  });
-
-  describe("$upgrade flow", () => {
-    describe("#upgrade", () => {
-      it("should upgrade the contract", async () => {
-        const PoseidonSMT = await ethers.getContractFactory("PoseidonSMT", {
-          libraries: {
-            PoseidonUnit2L: await (await getPoseidon(2)).getAddress(),
-            PoseidonUnit3L: await (await getPoseidon(3)).getAddress(),
-          },
-        });
-        const newTree = await PoseidonSMT.deploy();
-
-        const signature = merkleTree.authorizeUpgradeOperation(
-          await newTree.getAddress(),
-          chainName,
-          await tree.getNonce(TSSUpgradeableId.MAGIC_ID),
-          await tree.getAddress(),
-        );
-
-        await tree.upgradeToWithProof(await newTree.getAddress(), signature);
-
-        expect(await tree.implementation()).to.be.eq(await newTree.getAddress());
-      });
-
-      it("should revert if trying to upgrade to zero address", async () => {
-        const signature = merkleTree.authorizeUpgradeOperation(
-          ethers.ZeroAddress,
-          chainName,
-          await tree.getNonce(TSSUpgradeableId.MAGIC_ID),
-          await tree.getAddress(),
-        );
-
-        await expect(tree.upgradeToWithProof(ethers.ZeroAddress, signature)).to.be.rejectedWith(
-          "Upgradeable: Zero address",
-        );
-      });
-
-      it("should revert if operation was signed by the invalid signer", async () => {
-        const ANOTHER_SIGNER = ethers.Wallet.createRandom();
-
-        const signature = merkleTree.authorizeUpgradeOperation(
-          await tree.getAddress(),
-          chainName,
-          await tree.getNonce(TSSUpgradeableId.MAGIC_ID),
-          await tree.getAddress(),
-          ANOTHER_SIGNER,
-        );
-
-        await expect(tree.upgradeToWithProof(await tree.getAddress(), signature)).to.be.rejectedWith(
-          "TSSSigner: invalid signature",
-        );
-      });
-    });
-
-    it("should revert if trying to use default `upgradeTo` method", async () => {
-      await expect(tree.upgradeTo(ethers.ZeroAddress)).to.be.rejectedWith("Upgradeable: This upgrade method is off");
-    });
-  });
 });