From a7793bc10425d6f01e0b00a63520166837990221 Mon Sep 17 00:00:00 2001
From: Cian Hatton <cian@interchain.io>
Date: Fri, 30 Jun 2023 09:43:15 +0100
Subject: [PATCH] Add flushing check to WriteUpgradeAckChannel (#3976)

* chore: adding check for in flight packets in WriteUpgradeAckChannel

* added test for TestWriteUpgradeAckChannel

* linter

* add client update to UpgradeAckChannel test

* mv test

* merge

* fix post-merge

* fix merge issues

* review comment

---------

Co-authored-by: Charly <charly@interchain.berlin>
Co-authored-by: Carlos Rodriguez <carlos@interchain.io>
---
 modules/core/04-channel/keeper/upgrade.go     |  4 ++
 .../core/04-channel/keeper/upgrade_test.go    | 64 +++++++++++++++++++
 modules/core/keeper/msg_server_test.go        |  1 +
 3 files changed, 69 insertions(+)

diff --git a/modules/core/04-channel/keeper/upgrade.go b/modules/core/04-channel/keeper/upgrade.go
index 92cbea537a0..f2d3dae35bf 100644
--- a/modules/core/04-channel/keeper/upgrade.go
+++ b/modules/core/04-channel/keeper/upgrade.go
@@ -294,6 +294,10 @@ func (k Keeper) WriteUpgradeAckChannel(
 	channel.State = types.ACKUPGRADE
 	channel.FlushStatus = types.FLUSHING
 
+	if !k.hasInflightPackets(ctx, portID, channelID) {
+		channel.FlushStatus = types.FLUSHCOMPLETE
+	}
+
 	k.SetChannel(ctx, portID, channelID, channel)
 
 	upgrade, found := k.GetUpgrade(ctx, portID, channelID)
diff --git a/modules/core/04-channel/keeper/upgrade_test.go b/modules/core/04-channel/keeper/upgrade_test.go
index 7b0ac07d557..15a15ed12d9 100644
--- a/modules/core/04-channel/keeper/upgrade_test.go
+++ b/modules/core/04-channel/keeper/upgrade_test.go
@@ -552,6 +552,70 @@ func (suite *KeeperTestSuite) TestChanUpgradeAck() {
 	}
 }
 
+func (suite *KeeperTestSuite) TestWriteChannelUpgradeAck() {
+	var (
+		path            *ibctesting.Path
+		proposedUpgrade types.Upgrade
+	)
+
+	testCases := []struct {
+		name                 string
+		malleate             func()
+		hasPacketCommitments bool
+	}{
+		{
+			"success with no packet commitments",
+			func() {},
+			false,
+		},
+		{
+			"success with packet commitments",
+			func() {
+				// manually set packet commitment
+				sequence, err := path.EndpointA.SendPacket(suite.chainB.GetTimeoutHeight(), 0, ibctesting.MockPacketData)
+				suite.Require().NoError(err)
+				suite.Require().Equal(uint64(1), sequence)
+			},
+			true,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		suite.Run(tc.name, func() {
+			suite.SetupTest()
+
+			path = ibctesting.NewPath(suite.chainA, suite.chainB)
+			suite.coordinator.Setup(path)
+
+			// create upgrade proposal with different version in init upgrade to see if the WriteUpgradeAck overwrites the version
+			path.EndpointA.ChannelConfig.ProposedUpgrade.Fields.Version = "original-version"
+
+			tc.malleate()
+
+			// perform the upgrade handshake.
+			suite.Require().NoError(path.EndpointA.ChanUpgradeInit())
+
+			proposedUpgrade = path.EndpointA.GetChannelUpgrade()
+			suite.Require().Equal("original-version", proposedUpgrade.Fields.Version)
+
+			suite.Require().NoError(path.EndpointB.ChanUpgradeTry())
+
+			suite.chainA.GetSimApp().IBCKeeper.ChannelKeeper.WriteUpgradeAckChannel(suite.chainA.GetContext(), path.EndpointA.ChannelConfig.PortID, path.EndpointA.ChannelID, mock.UpgradeVersion)
+
+			channel := path.EndpointA.GetChannel()
+			upgrade := path.EndpointA.GetChannelUpgrade()
+			suite.Require().Equal(mock.UpgradeVersion, upgrade.Fields.Version)
+
+			if tc.hasPacketCommitments {
+				suite.Require().Equal(types.FLUSHING, channel.FlushStatus)
+			} else {
+				suite.Require().Equal(types.FLUSHCOMPLETE, channel.FlushStatus)
+			}
+		})
+	}
+}
+
 func (suite *KeeperTestSuite) TestChanUpgradeOpen() {
 	var path *ibctesting.Path
 	testCases := []struct {
diff --git a/modules/core/keeper/msg_server_test.go b/modules/core/keeper/msg_server_test.go
index beefe6293f6..6285d8bfae2 100644
--- a/modules/core/keeper/msg_server_test.go
+++ b/modules/core/keeper/msg_server_test.go
@@ -947,6 +947,7 @@ func (suite *KeeperTestSuite) TestChannelUpgradeAck() {
 				channel := path.EndpointA.GetChannel()
 				suite.Require().Equal(channeltypes.ACKUPGRADE, channel.State)
 				suite.Require().Equal(uint64(1), channel.UpgradeSequence)
+				suite.Require().Equal(channeltypes.FLUSHCOMPLETE, channel.FlushStatus)
 			},
 		},
 		{