From deac22dcce61c488379e8a584dd323ca337d61d8 Mon Sep 17 00:00:00 2001
From: Jun Kimura <jun.kimura@datachain.jp>
Date: Sun, 24 Nov 2024 18:39:36 +0900
Subject: [PATCH] add `counterpartyUpgradeTimeout` to the storage

Signed-off-by: Jun Kimura <jun.kimura@datachain.jp>
---
 contracts/core/04-channel/IBCChannelUpgrade.sol | 14 +++++++++-----
 contracts/core/24-host/IBCStore.sol             |  3 ++-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/contracts/core/04-channel/IBCChannelUpgrade.sol b/contracts/core/04-channel/IBCChannelUpgrade.sol
index 3eb95730..d4a0409c 100644
--- a/contracts/core/04-channel/IBCChannelUpgrade.sol
+++ b/contracts/core/04-channel/IBCChannelUpgrade.sol
@@ -35,9 +35,10 @@ abstract contract IBCChannelUpgradeBase is IBCModuleManager, IIBCChannelUpgradeB
         channel.state = Channel.State.STATE_OPEN;
 
         delete channelStorage.upgrade;
-        deleteUpgradeCommitment(portId, channelId);
-        revertCounterpartyUpgrade(channelStorage.recvStartSequence);
+        revertCounterpartyUpgrade(channelStorage);
+        delete channelStorage.counterpartyUpgradeTimeout;
 
+        deleteUpgradeCommitment(portId, channelId);
         updateChannelCommitment(portId, channelId, channel);
         writeErrorReceipt(portId, channelId, channel.upgrade_sequence, err);
     }
@@ -81,13 +82,15 @@ abstract contract IBCChannelUpgradeBase is IBCModuleManager, IIBCChannelUpgradeB
         getCommitments()[IBCCommitment.channelUpgradeCommitmentKey(portId, channelId)] = commitment;
     }
 
-    function revertCounterpartyUpgrade(RecvStartSequence storage recvStartSequence) internal {
+    function revertCounterpartyUpgrade(ChannelStorage storage channelStorage) internal {
+        RecvStartSequence storage recvStartSequence = channelStorage.recvStartSequence;
         uint64 prevRecvStartSequence = recvStartSequence.prevSequence;
         if (prevRecvStartSequence == 0) {
             return;
         }
         recvStartSequence.prevSequence = 0;
         recvStartSequence.sequence = prevRecvStartSequence;
+        delete channelStorage.counterpartyUpgradeTimeout;
     }
 
     function toString(UpgradeHandshakeError err) internal pure returns (string memory) {
@@ -135,10 +138,11 @@ abstract contract IBCChannelUpgradeCommon is IBCChannelUpgradeBase {
     function setCounterpartyUpgrade(ChannelStorage storage channelStorage, Upgrade.Data calldata upgrade) internal {
         RecvStartSequence storage recvStartSequence = channelStorage.recvStartSequence;
         if (recvStartSequence.prevSequence != 0) {
-            revertCounterpartyUpgrade(recvStartSequence);
+            revertCounterpartyUpgrade(channelStorage);
         }
         recvStartSequence.prevSequence = recvStartSequence.sequence;
         recvStartSequence.sequence = upgrade.next_sequence_send;
+        channelStorage.counterpartyUpgradeTimeout = upgrade.timeout;
     }
 
     function verifyMembership(
@@ -230,7 +234,7 @@ contract IBCChannelUpgradeInitTryAck is
         Upgrade.Data storage upgrade = channelStorage.upgrade;
         if (upgrade.fields.ordering != Channel.Order.ORDER_NONE_UNSPECIFIED) {
             delete channelStorage.upgrade;
-            revertCounterpartyUpgrade(channelStorage.recvStartSequence);
+            revertCounterpartyUpgrade(channelStorage);
             // NOTE: we do not delete the upgrade commitment here since the new upgrade will overwrite the old one
             writeErrorReceipt(msg_.portId, msg_.channelId, channel.upgrade_sequence, UpgradeHandshakeError.Overwritten);
         }
diff --git a/contracts/core/24-host/IBCStore.sol b/contracts/core/24-host/IBCStore.sol
index b6bc178c..a3e177a4 100644
--- a/contracts/core/24-host/IBCStore.sol
+++ b/contracts/core/24-host/IBCStore.sol
@@ -2,7 +2,7 @@
 pragma solidity ^0.8.20;
 
 import {ConnectionEnd} from "../../proto/Connection.sol";
-import {Channel, Upgrade} from "../../proto/Channel.sol";
+import {Channel, Upgrade, Timeout} from "../../proto/Channel.sol";
 
 abstract contract IBCStore {
     // keccak256(abi.encode(uint256(keccak256("ibc.commitment")) - 1)) & ~bytes32(uint256(0xff))
@@ -63,6 +63,7 @@ abstract contract IBCStore {
         uint64 nextSequenceRecv;
         uint64 nextSequenceAck;
         Upgrade.Data upgrade;
+        Timeout.Data counterpartyUpgradeTimeout;
         uint64 latestErrorReceiptSequence;
         RecvStartSequence recvStartSequence;
         uint64 ackStartSequence;