From 9175dd6f19ebf492a027fd6ccc16ee4b90642e2b Mon Sep 17 00:00:00 2001
From: Lazar <12626340+Lazar955@users.noreply.github.com>
Date: Fri, 8 Nov 2024 15:47:09 +0100
Subject: [PATCH] chore(swe): prevent stale reads (#96)

references
[issue](https://github.com/babylonlabs-io/vigilante/issues/86)
---
 CHANGELOG.md                                      |  1 +
 .../stakingeventwatcher/stakingeventwatcher.go    |  7 +------
 .../stakingeventwatcher/tracked_delegations.go    | 15 +++++++++++++++
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 54ce6384..24faed89 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -38,6 +38,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## Unreleased
 
 * [#94](https://github.com/babylonlabs-io/vigilante/pull/94) adds gosec and fixes gosec issues
+* [#96](https://github.com/babylonlabs-io/vigilante/pull/96) fixes potential stale data read
 
 ## v0.15.0
 
diff --git a/btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go b/btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go
index 6eb90dee..95c131be 100644
--- a/btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go
+++ b/btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go
@@ -580,12 +580,7 @@ func (sew *StakingEventWatcher) handlerVerifiedDelegations() {
 // checkBtcForStakingTx gets a snapshot of current Delegations in cache
 // checks if staking tx is in BTC, generates a proof and invokes sending of MsgAddBTCDelegationInclusionProof
 func (sew *StakingEventWatcher) checkBtcForStakingTx() error {
-	delegations := sew.pendingTracker.GetDelegations()
-	if delegations == nil {
-		return nil
-	}
-
-	for _, del := range delegations {
+	for del := range sew.pendingTracker.DelegationsIter() {
 		if del.ActivationInProgress {
 			continue
 		}
diff --git a/btcstaking-tracker/stakingeventwatcher/tracked_delegations.go b/btcstaking-tracker/stakingeventwatcher/tracked_delegations.go
index 7bbcecf3..5d2ddce6 100644
--- a/btcstaking-tracker/stakingeventwatcher/tracked_delegations.go
+++ b/btcstaking-tracker/stakingeventwatcher/tracked_delegations.go
@@ -2,6 +2,7 @@ package stakingeventwatcher
 
 import (
 	"fmt"
+	"iter"
 	"sync"
 
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
@@ -58,6 +59,20 @@ func (td *TrackedDelegations) GetDelegations() []*TrackedDelegation {
 	return delegations
 }
 
+func (td *TrackedDelegations) DelegationsIter() iter.Seq[*TrackedDelegation] {
+	return func(yield func(*TrackedDelegation) bool) {
+		td.mu.RLock()
+		defer td.mu.RUnlock()
+
+		// we lock for the entirety of the iteration
+		for _, v := range td.mapping {
+			if !yield(v) {
+				return
+			}
+		}
+	}
+}
+
 func (td *TrackedDelegations) AddDelegation(
 	StakingTx *wire.MsgTx,
 	StakingOutputIdx uint32,