From c663ba08f2f2f47f7d33f0a04b3b5a543a25968c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?=
 <jedrzej.stuczynski@gmail.com>
Date: Thu, 15 Feb 2024 21:01:46 +0000
Subject: [PATCH] fixed dkg incorrectly setting state deadlines

---
 .../coconut-dkg/src/types.rs                  | 49 ++++++++++++-------
 .../transactions/advance_epoch_state.rs       |  3 +-
 2 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/common/cosmwasm-smart-contracts/coconut-dkg/src/types.rs b/common/cosmwasm-smart-contracts/coconut-dkg/src/types.rs
index a4d0afb077..0e527dee71 100644
--- a/common/cosmwasm-smart-contracts/coconut-dkg/src/types.rs
+++ b/common/cosmwasm-smart-contracts/coconut-dkg/src/types.rs
@@ -34,6 +34,26 @@ pub struct TimeConfiguration {
     pub in_progress_time_secs: u64,
 }
 
+impl TimeConfiguration {
+    pub fn state_duration(&self, state: EpochState) -> Option<u64> {
+        match state {
+            EpochState::WaitingInitialisation => None,
+            EpochState::PublicKeySubmission { .. } => Some(self.public_key_submission_time_secs),
+            EpochState::DealingExchange { .. } => Some(self.dealing_exchange_time_secs),
+            EpochState::VerificationKeySubmission { .. } => {
+                Some(self.verification_key_submission_time_secs)
+            }
+            EpochState::VerificationKeyValidation { .. } => {
+                Some(self.verification_key_validation_time_secs)
+            }
+            EpochState::VerificationKeyFinalization { .. } => {
+                Some(self.verification_key_finalization_time_secs)
+            }
+            EpochState::InProgress => Some(self.in_progress_time_secs),
+        }
+    }
+}
+
 impl FromStr for TimeConfiguration {
     type Err = String;
 
@@ -116,25 +136,8 @@ impl Epoch {
         time_configuration: TimeConfiguration,
         current_timestamp: Timestamp,
     ) -> Self {
-        let duration = match state {
-            EpochState::WaitingInitialisation => None,
-            EpochState::PublicKeySubmission { .. } => {
-                Some(time_configuration.public_key_submission_time_secs)
-            }
-            EpochState::DealingExchange { .. } => {
-                Some(time_configuration.dealing_exchange_time_secs)
-            }
-            EpochState::VerificationKeySubmission { .. } => {
-                Some(time_configuration.verification_key_submission_time_secs)
-            }
-            EpochState::VerificationKeyValidation { .. } => {
-                Some(time_configuration.verification_key_validation_time_secs)
-            }
-            EpochState::VerificationKeyFinalization { .. } => {
-                Some(time_configuration.verification_key_finalization_time_secs)
-            }
-            EpochState::InProgress => Some(time_configuration.in_progress_time_secs),
-        };
+        let duration = time_configuration.state_duration(state);
+
         Epoch {
             state,
             epoch_id,
@@ -144,6 +147,14 @@ impl Epoch {
         }
     }
 
+    pub fn update(mut self, next_state: EpochState, current_timestamp: Timestamp) -> Self {
+        self.state = next_state;
+        let duration = self.time_configuration.state_duration(next_state);
+        self.deadline = duration.map(|d| current_timestamp.plus_seconds(d));
+
+        self
+    }
+
     pub fn final_timestamp_secs(&self) -> Option<u64> {
         let mut finish = self.deadline?.seconds();
         let time_configuration = self.time_configuration;
diff --git a/contracts/coconut-dkg/src/epoch_state/transactions/advance_epoch_state.rs b/contracts/coconut-dkg/src/epoch_state/transactions/advance_epoch_state.rs
index 57681ed3cf..789eeff601 100644
--- a/contracts/coconut-dkg/src/epoch_state/transactions/advance_epoch_state.rs
+++ b/contracts/coconut-dkg/src/epoch_state/transactions/advance_epoch_state.rs
@@ -73,8 +73,7 @@ pub fn try_advance_epoch_state(deps: DepsMut<'_>, env: Env) -> Result<Response,
     }
 
     // update the epoch state
-    let mut next_epoch = current_epoch;
-    next_epoch.state = next_state;
+    let next_epoch = current_epoch.update(next_state, env.block.time);
     CURRENT_EPOCH.save(deps.storage, &next_epoch)?;
 
     Ok(Response::new())