From 6da7bd00383ed21f5af8858e0011e6cf09dba1e6 Mon Sep 17 00:00:00 2001
From: rymnc <43716372+rymnc@users.noreply.github.com>
Date: Thu, 3 Aug 2023 18:25:13 +0530
Subject: [PATCH] fix: tests

---
 contracts/WakuRlnRegistry.sol | 24 +++++++++++----
 docs/index.md                 | 18 +++++++++--
 test/WakuRln.t.sol            |  2 +-
 test/WakuRlnRegistry.t.sol    | 56 +++++++++++++++++++++++++++++++++++
 4 files changed, 90 insertions(+), 10 deletions(-)
 create mode 100644 test/WakuRlnRegistry.t.sol

diff --git a/contracts/WakuRlnRegistry.sol b/contracts/WakuRlnRegistry.sol
index cb8688a..a4f1094 100644
--- a/contracts/WakuRlnRegistry.sol
+++ b/contracts/WakuRlnRegistry.sol
@@ -4,10 +4,12 @@ pragma solidity 0.8.15;
 import {WakuRln} from "./WakuRln.sol";
 import {IPoseidonHasher} from "rln-contract/PoseidonHasher.sol";
 import {Ownable} from "openzeppelin-contracts/contracts/access/Ownable.sol";
+import "forge-std/console.sol";
 
 error StorageAlreadyExists(address storageAddress);
 error NoStorageContractAvailable();
-error FailedToRegister(string reason);
+error IncompatibleStorage();
+error IncompatibleStorageIndex();
 
 contract WakuRlnRegistry is Ownable {
     uint16 public nextStorageIndex;
@@ -31,6 +33,9 @@ contract WakuRlnRegistry is Ownable {
 
     function registerStorage(address storageAddress) external onlyOwner {
         if (storages[nextStorageIndex] != address(0)) revert StorageAlreadyExists(storageAddress);
+        WakuRln wakuRln = WakuRln(storageAddress);
+        if (wakuRln.poseidonHasher() != poseidonHasher) revert IncompatibleStorage();
+        if (wakuRln.contractIndex() != nextStorageIndex) revert IncompatibleStorageIndex();
         _insertIntoStorageMap(storageAddress);
     }
 
@@ -39,19 +44,26 @@ contract WakuRlnRegistry is Ownable {
         _insertIntoStorageMap(address(newStorageContract));
     }
 
-    function register(uint256 commitment) external payable {
+    function register(uint256[] calldata commitments) external payable {
         if (usingStorageIndex >= nextStorageIndex) revert NoStorageContractAvailable();
 
         // iteratively check if the storage contract is full, and increment the usingStorageIndex if it is
         while (true) {
-            try WakuRln(storages[usingStorageIndex]).register{value: msg.value}(commitment) {
+            try WakuRln(storages[usingStorageIndex]).register(commitments) {
                 break;
-            } catch Error(string memory reason) {
-                if (keccak256(abi.encodePacked(reason)) != keccak256(abi.encodePacked("FullTree()"))) {
-                    revert FailedToRegister(reason);
+            } catch (bytes memory err) {
+                if (keccak256(err) != keccak256(abi.encodeWithSignature("FullTree()"))) {
+                    assembly {
+                        revert(add(32, err), mload(err))
+                    }
                 }
                 usingStorageIndex += 1;
             }
         }
     }
+
+    function forceProgress() external onlyOwner {
+        if (usingStorageIndex >= nextStorageIndex) revert NoStorageContractAvailable();
+        usingStorageIndex += 1;
+    }
 }
diff --git a/docs/index.md b/docs/index.md
index d7931ec..0b17306 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -104,10 +104,16 @@ error StorageAlreadyExists(address storageAddress)
 error NoStorageContractAvailable()
 ```
 
-## FailedToRegister
+## IncompatibleStorage
 
 ```solidity
-error FailedToRegister(string reason)
+error IncompatibleStorage()
+```
+
+## IncompatibleStorageIndex
+
+```solidity
+error IncompatibleStorageIndex()
 ```
 
 ## WakuRlnRegistry
@@ -169,5 +175,11 @@ function newStorage() external
 ### register
 
 ```solidity
-function register(uint256 commitment) external payable
+function register(uint256[] commitments) external payable
+```
+
+### forceProgress
+
+```solidity
+function forceProgress() external
 ```
diff --git a/test/WakuRln.t.sol b/test/WakuRln.t.sol
index 4ac0ffb..12cdd5b 100644
--- a/test/WakuRln.t.sol
+++ b/test/WakuRln.t.sol
@@ -1,4 +1,4 @@
-// SPDX-License-Identifier: Unlicense
+// SPDX-License-Identifier: MIT
 pragma solidity ^0.8.15;
 
 import {PoseidonHasher} from "rln-contract/PoseidonHasher.sol";
diff --git a/test/WakuRlnRegistry.t.sol b/test/WakuRlnRegistry.t.sol
new file mode 100644
index 0000000..60f73d9
--- /dev/null
+++ b/test/WakuRlnRegistry.t.sol
@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.15;
+
+import {PoseidonHasher} from "rln-contract/PoseidonHasher.sol";
+import "../contracts/WakuRlnRegistry.sol";
+import {noDuplicate} from "./WakuRln.t.sol";
+import {DuplicateIdCommitment} from "rln-contract/RlnBase.sol";
+import "forge-std/Test.sol";
+import "forge-std/StdCheats.sol";
+
+contract WakuRlnRegistryTest is Test {
+    WakuRlnRegistry public wakuRlnRegistry;
+    PoseidonHasher public poseidonHasher;
+
+    function setUp() public {
+        poseidonHasher = new PoseidonHasher();
+        wakuRlnRegistry = new WakuRlnRegistry(address(poseidonHasher));
+    }
+
+    function test__NewStorage() public {
+        wakuRlnRegistry.newStorage();
+    }
+
+    function test__RegisterStorage_BadIndex() public {
+        wakuRlnRegistry.registerStorage(address(new WakuRln(address(poseidonHasher), 0)));
+        address newStorage = address(new WakuRln(address(poseidonHasher), 0));
+        vm.expectRevert(IncompatibleStorageIndex.selector);
+        wakuRlnRegistry.registerStorage(newStorage);
+    }
+
+    function test__RegisterStorage_BadImpl() public {
+        address newStorage = address(new WakuRln(address(new PoseidonHasher()), 0));
+        vm.expectRevert(IncompatibleStorage.selector);
+        wakuRlnRegistry.registerStorage(newStorage);
+    }
+
+    function test__Register(uint256[] calldata commitments) public {
+        vm.assume(noDuplicate(commitments));
+        wakuRlnRegistry.newStorage();
+        wakuRlnRegistry.register(commitments);
+    }
+
+    function test__BadRegister(uint256[] calldata commitments) public {
+        vm.assume(!noDuplicate(commitments));
+        wakuRlnRegistry.newStorage();
+        vm.expectRevert(DuplicateIdCommitment.selector);
+        wakuRlnRegistry.register(commitments);
+    }
+
+    function test__forceProgression() public {
+        wakuRlnRegistry.newStorage();
+        wakuRlnRegistry.newStorage();
+        wakuRlnRegistry.forceProgress();
+        require(wakuRlnRegistry.usingStorageIndex() == 1);
+    }
+}