From 4824ffd89dea00fe48d0d00d3d1900d6edd12066 Mon Sep 17 00:00:00 2001
From: Damian Nolan <damiannolan@gmail.com>
Date: Wed, 16 Mar 2022 18:39:49 +0100
Subject: [PATCH] chore: 06-solomachine adding CheckForMisbehaviour and
 UpdateStateOnMisbehaviour (#1111)

* adding VerifyClientMessage and tests

* splitting tests for verify header and misbehaviour

* rename update to UpdateState

* adding CheckForMisbehaviour and UpdateStateOnMisbehaviour

* adding misbehaviour checks

* updating CheckForMisbehaviour codestyle and adding basic test cases
---
 .../06-solomachine/types/update.go            | 23 ++++++
 .../06-solomachine/types/update_test.go       | 82 +++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/modules/light-clients/06-solomachine/types/update.go b/modules/light-clients/06-solomachine/types/update.go
index eb3941b89d2..d04d93c36af 100644
--- a/modules/light-clients/06-solomachine/types/update.go
+++ b/modules/light-clients/06-solomachine/types/update.go
@@ -23,6 +23,11 @@ func (cs ClientState) CheckHeaderAndUpdateState(
 		return nil, nil, err
 	}
 
+	foundMisbehaviour := cs.CheckForMisbehaviour(ctx, cdc, clientStore, msg)
+	if foundMisbehaviour {
+		return cs.UpdateStateOnMisbehaviour(ctx, cdc, clientStore)
+	}
+
 	return cs.UpdateState(ctx, cdc, clientStore, msg)
 }
 
@@ -109,3 +114,21 @@ func (cs ClientState) UpdateState(ctx sdk.Context, cdc codec.BinaryCodec, client
 	cs.ConsensusState = consensusState
 	return &cs, consensusState, nil
 }
+
+// CheckForMisbehaviour returns true for type Misbehaviour (passed VerifyClientMessage check), otherwise returns false
+func (cs ClientState) CheckForMisbehaviour(_ sdk.Context, _ codec.BinaryCodec, _ sdk.KVStore, clientMsg exported.ClientMessage) bool {
+	if _, ok := clientMsg.(*Misbehaviour); ok {
+		return true
+	}
+
+	return false
+}
+
+// UpdateStateOnMisbehaviour updates state upon misbehaviour. This method should only be called on misbehaviour
+// as it does not perform any misbehaviour checks.
+func (cs ClientState) UpdateStateOnMisbehaviour(
+	_ sdk.Context, _ codec.BinaryCodec, _ sdk.KVStore, // prematurely include args for self storage of consensus state
+) (*ClientState, exported.ConsensusState, error) {
+	cs.IsFrozen = true
+	return &cs, cs.ConsensusState, nil
+}
diff --git a/modules/light-clients/06-solomachine/types/update_test.go b/modules/light-clients/06-solomachine/types/update_test.go
index ddd239ab9a4..56f620b15aa 100644
--- a/modules/light-clients/06-solomachine/types/update_test.go
+++ b/modules/light-clients/06-solomachine/types/update_test.go
@@ -624,3 +624,85 @@ func (suite *SoloMachineTestSuite) TestUpdateState() {
 		}
 	}
 }
+
+func (suite *SoloMachineTestSuite) TestCheckForMisbehaviour() {
+	var (
+		clientMsg exported.ClientMessage
+	)
+
+	// test singlesig and multisig public keys
+	for _, solomachine := range []*ibctesting.Solomachine{suite.solomachine, suite.solomachineMulti} {
+		testCases := []struct {
+			name     string
+			malleate func()
+			expPass  bool
+		}{
+			{
+				"success",
+				func() {
+					clientMsg = solomachine.CreateMisbehaviour()
+				},
+				true,
+			},
+			{
+				"normal header returns false",
+				func() {
+					clientMsg = solomachine.CreateHeader()
+				},
+				false,
+			},
+		}
+
+		for _, tc := range testCases {
+			tc := tc
+
+			suite.Run(tc.name, func() {
+				clientState := solomachine.ClientState()
+
+				tc.malleate()
+
+				foundMisbehaviour := clientState.CheckForMisbehaviour(suite.chainA.GetContext(), suite.chainA.Codec, suite.store, clientMsg)
+
+				if tc.expPass {
+					suite.Require().True(foundMisbehaviour)
+				} else {
+					suite.Require().False(foundMisbehaviour)
+				}
+
+			})
+		}
+	}
+}
+
+func (suite *SoloMachineTestSuite) TestUpdateStateOnMisbehaviour() {
+	// test singlesig and multisig public keys
+	for _, solomachine := range []*ibctesting.Solomachine{suite.solomachine, suite.solomachineMulti} {
+		testCases := []struct {
+			name     string
+			malleate func()
+			expPass  bool
+		}{
+			{
+				"success",
+				func() {},
+				true,
+			},
+		}
+
+		for _, tc := range testCases {
+			tc := tc
+
+			suite.Run(tc.name, func() {
+				clientState := solomachine.ClientState()
+
+				tc.malleate()
+
+				cs, _, _ := clientState.UpdateStateOnMisbehaviour(suite.chainA.GetContext(), suite.chainA.Codec, suite.store)
+
+				if tc.expPass {
+					suite.Require().True(cs.IsFrozen)
+				}
+			})
+		}
+	}
+}