From 88fd29458afe10ae2ee8994e992570b7d668c946 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matija=20Petruni=C4=87?= <matija.petrunic@gmail.com>
Date: Thu, 20 Jun 2024 13:51:27 +0200
Subject: [PATCH] fix: make frost keygen retryable (#251)

---
 contracts/FROSTKeygen.sol       | 25 ++++++++++++++++---------
 test/frostKeygen/frostKeygen.js | 19 +++++++++++--------
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/contracts/FROSTKeygen.sol b/contracts/FROSTKeygen.sol
index 2776723b..8243b95c 100644
--- a/contracts/FROSTKeygen.sol
+++ b/contracts/FROSTKeygen.sol
@@ -7,20 +7,27 @@ import "@openzeppelin/contracts/access/Ownable.sol";
 
 contract FROSTKeygen is Ownable {
 
-    bool private keygenStarted; 
-    event StartedFROSTKeygen(); 
+    bool private keygenEnded = false; 
+
+    event StartedFROSTKeygen();
+    event EndedFROSTKeygen();
 
-    modifier onlyOnce(){
-        require (!keygenStarted, "FROST keygen can be called only once");
-        _; 
-        keygenStarted = true; 
-    }
-    
     /**
        @notice Emits {StartedFROSTKeygen} event
      */
-    function startFROSTKeygen() public onlyOwner onlyOnce { 
+    function startFROSTKeygen() public onlyOwner { 
+        require (!keygenEnded, "FROST keygen ended");
+
         emit StartedFROSTKeygen(); 
     }
 
+    /**
+       @notice Blocks further calls for starting keygen.
+     */
+    function endFROSTKeygen() public onlyOwner { 
+        keygenEnded = true;
+
+        emit EndedFROSTKeygen();
+    }
+
 }
\ No newline at end of file
diff --git a/test/frostKeygen/frostKeygen.js b/test/frostKeygen/frostKeygen.js
index bd6002da..41b9d0c1 100644
--- a/test/frostKeygen/frostKeygen.js
+++ b/test/frostKeygen/frostKeygen.js
@@ -12,7 +12,6 @@ contract("FROSTKeygen", (accounts) => {
     });
 
     it("should emit StartedFROSTKeygen event when startFROSTKeygen is called by the owner", async () => {
-
       const tx = await FROSTKeygenInstance.startFROSTKeygen({from: accounts[0]})
 
       TruffleAssert.eventEmitted(tx, "StartedFROSTKeygen");
@@ -20,20 +19,24 @@ contract("FROSTKeygen", (accounts) => {
     });
 
     it("should revert when startFROSTKeygen is not called by the owner", async () => {
-
       await TruffleAssert.reverts(
         FROSTKeygenInstance.startFROSTKeygen({from: accounts[1]}),
       )
 
     });
 
-    it("should revert when startFROSTKeygen is called more than once", async() => {
-
-      const tx = await FROSTKeygenInstance.startFROSTKeygen({from: accounts[0]})
+    it("should revert when keygen ended", async() => {
+      const tx = await FROSTKeygenInstance.endFROSTKeygen({from: accounts[0]})
+      TruffleAssert.eventEmitted(tx, "EndedFROSTKeygen");
 
-      TruffleAssert.eventEmitted(tx, "StartedFROSTKeygen");
+      await TruffleAssert.reverts(
+        FROSTKeygenInstance.startFROSTKeygen({from: accounts[1]}),
+      )
+    });
 
+    it("should revert when end keygen not called by owner", async() => {
       await TruffleAssert.reverts(
-        FROSTKeygenInstance.startFROSTKeygen({from: accounts[0]}))
-      })
+        FROSTKeygenInstance.endFROSTKeygen({from: accounts[1]}),
+      )
+    });
 })