From 25f27224e77dd4fb99596fe2443f711a3bdb7a9a Mon Sep 17 00:00:00 2001
From: Jayden Lee <41176085+tkxkd0159@users.noreply.github.com>
Date: Fri, 3 May 2024 10:47:56 +0900
Subject: [PATCH] add invariants for metadata

---
 x/fbridge/keeper/genesis.go | 15 -----------
 x/fbridge/module/abci.go    | 51 ++++++++++++++++++++++---------------
 x/fbridge/module/module.go  |  2 +-
 3 files changed, 32 insertions(+), 36 deletions(-)

diff --git a/x/fbridge/keeper/genesis.go b/x/fbridge/keeper/genesis.go
index 5e93b1bf83..e3ed8445b2 100644
--- a/x/fbridge/keeper/genesis.go
+++ b/x/fbridge/keeper/genesis.go
@@ -19,23 +19,8 @@ func (k Keeper) InitGenesis(ctx sdk.Context, gs *types.GenesisState) error {
 		k.setVote(ctx, vote.ProposalId, sdk.MustAccAddressFromBech32(vote.Voter), vote.Option)
 	}
 
-	expectedRoleMeta := types.RoleMetadata{}
 	for _, pair := range gs.Roles {
 		k.SetRole(ctx, pair.Role, sdk.MustAccAddressFromBech32(pair.Address))
-		switch pair.Role {
-		case types.RoleGuardian:
-			expectedRoleMeta.Guardian++
-		case types.RoleOperator:
-			expectedRoleMeta.Operator++
-		case types.RoleJudge:
-			expectedRoleMeta.Judge++
-		}
-	}
-
-	if expectedRoleMeta.Guardian != gs.RoleMetadata.Guardian ||
-		expectedRoleMeta.Operator != gs.RoleMetadata.Operator ||
-		expectedRoleMeta.Judge != gs.RoleMetadata.Judge {
-		panic("role metadata does not match the number of roles")
 	}
 	k.SetRoleMetadata(ctx, gs.RoleMetadata)
 
diff --git a/x/fbridge/module/abci.go b/x/fbridge/module/abci.go
index 0eb1368348..e5e8c37bb4 100644
--- a/x/fbridge/module/abci.go
+++ b/x/fbridge/module/abci.go
@@ -27,7 +27,6 @@ func EndBlocker(ctx sdk.Context, k keeper.Keeper) {
 
 		var total uint32 = 0
 		roleMeta := k.GetRoleMetadata(ctx)
-		previousRole := k.GetRole(ctx, sdk.MustAccAddressFromBech32(proposal.Target))
 		switch proposal.Role {
 		case types.RoleGuardian:
 			total = roleMeta.Guardian
@@ -40,31 +39,43 @@ func EndBlocker(ctx sdk.Context, k keeper.Keeper) {
 		}
 
 		if types.CheckTrustLevelThreshold(uint64(total), uint64(voteYes), guardianTrustLevel) {
-			if proposal.Role == types.RoleEmpty {
-				k.DeleteRole(ctx, sdk.MustAccAddressFromBech32(proposal.Target))
-			} else {
-				k.SetRole(ctx, proposal.Role, sdk.MustAccAddressFromBech32(proposal.Target))
+			if err := k.UpdateRole(ctx, proposal.Role, sdk.MustAccAddressFromBech32(proposal.Target)); err != nil {
+				panic(err)
 			}
 
-			switch proposal.Role {
-			case types.RoleGuardian:
-				roleMeta.Guardian++
-			case types.RoleOperator:
-				roleMeta.Operator++
-			case types.RoleJudge:
-				roleMeta.Judge++
-			}
-			switch previousRole {
+			k.DeleteRoleProposal(ctx, proposal.Id)
+		}
+	}
+}
+
+// RegisterInvariants registers all staking invariants
+func RegisterInvariants(ir sdk.InvariantRegistry, k keeper.Keeper) {
+	ir.RegisterRoute(types.ModuleName, "role-metadata", RoleMeatadataInvariant(k))
+}
+
+func RoleMeatadataInvariant(k keeper.Keeper) sdk.Invariant {
+	return func(ctx sdk.Context) (string, bool) {
+		actualRoleMeta := k.GetRoleMetadata(ctx)
+		expectedRoleMeta := types.RoleMetadata{}
+		for _, pair := range k.GetRolePairs(ctx) {
+			k.SetRole(ctx, pair.Role, sdk.MustAccAddressFromBech32(pair.Address))
+			switch pair.Role {
 			case types.RoleGuardian:
-				roleMeta.Guardian--
+				expectedRoleMeta.Guardian++
 			case types.RoleOperator:
-				roleMeta.Operator--
+				expectedRoleMeta.Operator++
 			case types.RoleJudge:
-				roleMeta.Judge--
+				expectedRoleMeta.Judge++
 			}
-			k.SetRoleMetadata(ctx, roleMeta)
-
-			k.DeleteRoleProposal(ctx, proposal.Id)
 		}
+
+		broken := expectedRoleMeta.Guardian != actualRoleMeta.Guardian ||
+			expectedRoleMeta.Operator != actualRoleMeta.Operator ||
+			expectedRoleMeta.Judge != actualRoleMeta.Judge
+
+		return sdk.FormatInvariant(types.ModuleName, "registered members and role metadata", fmt.Sprintf(
+			"Saved Role Metadata: %+v"+
+				"Calculated Role Metadata: %+v",
+			actualRoleMeta, expectedRoleMeta)), broken
 	}
 }
diff --git a/x/fbridge/module/module.go b/x/fbridge/module/module.go
index 9c654c7ed5..03bede71d9 100644
--- a/x/fbridge/module/module.go
+++ b/x/fbridge/module/module.go
@@ -131,7 +131,7 @@ func (am AppModule) EndBlock(ctx sdk.Context, _ abci.RequestEndBlock) []abci.Val
 
 // RegisterInvariants does nothing, there are no invariants to enforce
 func (am AppModule) RegisterInvariants(ir sdk.InvariantRegistry) {
-	_ = ir
+	RegisterInvariants(ir, am.keeper)
 }
 
 // Deprecated: Route does nothing.