From e958bab8171d10598cbc1eae447a1e48695a062e Mon Sep 17 00:00:00 2001
From: Nilesh Gupta <guptanilesh2312@gmail.com>
Date: Fri, 1 May 2026 09:50:30 +0530
Subject: [PATCH] fix: fixed unsafe modification order in MarkBallotExpired
 nand MarkBallotFinalized

---
 x/uvalidator/keeper/ballot.go | 31 +++++++++++++++++++------------
 x/uvalidator/keeper/voting.go | 23 ++++++++++-------------
 2 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/x/uvalidator/keeper/ballot.go b/x/uvalidator/keeper/ballot.go
index 4cf7b8e9..f483894d 100644
--- a/x/uvalidator/keeper/ballot.go
+++ b/x/uvalidator/keeper/ballot.go
@@ -110,7 +110,12 @@ func (k Keeper) DeleteBallot(ctx context.Context, id string) error {
 	return nil
 }
 
-// MarkBallotExpired moves a ballot from active to expired
+// MarkBallotExpired moves a ballot from active to expired.
+// Side-effect ordering: secondary indexes are updated before the canonical
+// ballot record is rewritten, so the status field is only persisted once the
+// active/expired set membership is in its final shape (defensive CEI-style
+// ordering; collections.KeySet.Remove is a no-op on absent keys, so retries
+// remain safe).
 func (k Keeper) MarkBallotExpired(ctx context.Context, id string) error {
 	ballot, err := k.Ballots.Get(ctx, id)
 	if err != nil {
@@ -122,18 +127,20 @@ func (k Keeper) MarkBallotExpired(ctx context.Context, id string) error {
 		"expiry_height", ballot.BlockHeightExpiry,
 	)
 
-	ballot.Status = types.BallotStatus_BALLOT_STATUS_EXPIRED
-	if err := k.Ballots.Set(ctx, id, ballot); err != nil {
+	if err := k.ActiveBallotIDs.Remove(ctx, id); err != nil {
 		return err
 	}
-
-	if err := k.ActiveBallotIDs.Remove(ctx, id); err != nil {
+	if err := k.ExpiredBallotIDs.Set(ctx, id); err != nil {
 		return err
 	}
-	return k.ExpiredBallotIDs.Set(ctx, id)
+
+	ballot.Status = types.BallotStatus_BALLOT_STATUS_EXPIRED
+	return k.Ballots.Set(ctx, id, ballot)
 }
 
-// MarkBallotFinalized moves a ballot from active to finalized (PASSED or REJECTED)
+// MarkBallotFinalized moves a ballot from active to finalized (PASSED or REJECTED).
+// Side-effect ordering matches MarkBallotExpired: secondary indexes are
+// updated before the canonical ballot record is rewritten with its final status.
 func (k Keeper) MarkBallotFinalized(ctx context.Context, id string, status types.BallotStatus) error {
 	if status != types.BallotStatus_BALLOT_STATUS_PASSED && status != types.BallotStatus_BALLOT_STATUS_REJECTED {
 		return fmt.Errorf("invalid finalization status: %v", status)
@@ -149,15 +156,15 @@ func (k Keeper) MarkBallotFinalized(ctx context.Context, id string, status types
 		"final_status", status.String(),
 	)
 
-	ballot.Status = status
-	if err := k.Ballots.Set(ctx, id, ballot); err != nil {
+	if err := k.ActiveBallotIDs.Remove(ctx, id); err != nil {
 		return err
 	}
-
-	if err := k.ActiveBallotIDs.Remove(ctx, id); err != nil {
+	if err := k.FinalizedBallotIDs.Set(ctx, id); err != nil {
 		return err
 	}
-	return k.FinalizedBallotIDs.Set(ctx, id)
+
+	ballot.Status = status
+	return k.Ballots.Set(ctx, id, ballot)
 }
 
 // ExpireBallotsBeforeHeight checks active ballots and marks expired ones.
diff --git a/x/uvalidator/keeper/voting.go b/x/uvalidator/keeper/voting.go
index b7733f74..44dcabfa 100644
--- a/x/uvalidator/keeper/voting.go
+++ b/x/uvalidator/keeper/voting.go
@@ -171,22 +171,14 @@ func (k Keeper) VoteOnBallot(
 	if err != nil {
 		return ballot, false, false, err
 	}
-	if isFinalizing {
-		k.Logger().Debug("ballot finalized",
-			"ballot_id", id,
-			"ballot_status", ballot.Status.String(),
-		)
-		if err := k.ActiveBallotIDs.Remove(ctx, id); err != nil {
-			return ballot, false, isNew, errors.Wrap(err, "failed removing from active ballots")
-		}
-		if err := k.FinalizedBallotIDs.Set(ctx, id); err != nil {
-			return ballot, false, isNew, errors.Wrap(err, "failed adding to finalized ballots")
-		}
-	}
 
 	return ballot, isFinalizing, isNew, nil
 }
 
+// CheckIfFinalizingVote inspects whether the just-cast vote pushes the ballot
+// over its threshold and, if so, drives the finalization through
+// MarkBallotFinalized — the single canonical write path for terminal status
+// transitions, which applies CEI-style ordering on the secondary indexes.
 func (k Keeper) CheckIfFinalizingVote(ctx context.Context, b types.Ballot) (types.Ballot, bool, error) {
 	ballot, isFinalizing := b.IsFinalizingVote()
 	if !isFinalizing {
@@ -198,8 +190,13 @@ func (k Keeper) CheckIfFinalizingVote(ctx context.Context, b types.Ballot) (type
 		"ballot_status", ballot.Status.String(),
 	)
 
-	if err := k.SetBallot(ctx, ballot); err != nil {
+	if err := k.MarkBallotFinalized(ctx, ballot.Id, ballot.Status); err != nil {
 		return ballot, false, errors.Wrap(err, "failed updating finalized ballot")
 	}
+
+	k.Logger().Debug("ballot finalized",
+		"ballot_id", ballot.Id,
+		"ballot_status", ballot.Status.String(),
+	)
 	return ballot, true, nil
 }