From c3664f9e20467e19ff98c4e39b66f7735df02c64 Mon Sep 17 00:00:00 2001
From: Tomasz Slabon <tomasz.slabon@keep.network>
Date: Wed, 15 May 2024 14:04:34 +0200
Subject: [PATCH] Added check for unstaking operators

---
 pkg/chain/ethereum/tbtc.go |  6 ------
 pkg/tbtc/chain.go          |  9 ---------
 pkg/tbtc/chain_test.go     |  5 -----
 pkg/tbtc/heartbeat.go      | 28 ++++++++++++++++++++++++++--
 4 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/pkg/chain/ethereum/tbtc.go b/pkg/chain/ethereum/tbtc.go
index 9cce628ce0..3fd3ce6172 100644
--- a/pkg/chain/ethereum/tbtc.go
+++ b/pkg/chain/ethereum/tbtc.go
@@ -2363,9 +2363,3 @@ func (tc *TbtcChain) GetRedemptionDelay(
 func (tc *TbtcChain) GetDepositMinAge() (uint32, error) {
 	return tc.walletProposalValidator.DEPOSITMINAGE()
 }
-
-func (tc *TbtcChain) IsOperatorUnstaking() (bool, error) {
-	// TODO: Implement by checking if the operator has deauthorized their entire
-	//       stake.
-	return false, nil
-}
diff --git a/pkg/tbtc/chain.go b/pkg/tbtc/chain.go
index 1b0a0c7c2a..1b579e326b 100644
--- a/pkg/tbtc/chain.go
+++ b/pkg/tbtc/chain.go
@@ -24,14 +24,6 @@ const (
 	Challenge
 )
 
-// StakingChain defines the subset of the TBTC chain interface that pertains to
-// the staking activities.
-type StakingChain interface {
-	// IsOperatorUnstaking checks if the operator is unstaking. It returns true
-	// if the operator has deauthorized their entire stake, false otherwise.
-	IsOperatorUnstaking() (bool, error)
-}
-
 // GroupSelectionChain defines the subset of the TBTC chain interface that
 // pertains to the group selection activities.
 type GroupSelectionChain interface {
@@ -539,7 +531,6 @@ type Chain interface {
 	GetBlockHashByNumber(blockNumber uint64) ([32]byte, error)
 
 	sortition.Chain
-	StakingChain
 	GroupSelectionChain
 	DistributedKeyGenerationChain
 	InactivityClaimChain
diff --git a/pkg/tbtc/chain_test.go b/pkg/tbtc/chain_test.go
index 95e250bd5e..a347c89a96 100644
--- a/pkg/tbtc/chain_test.go
+++ b/pkg/tbtc/chain_test.go
@@ -216,11 +216,6 @@ func (lc *localChain) IsEligibleForRewards() (bool, error) {
 	panic("unsupported")
 }
 
-func (lc *localChain) IsOperatorUnstaking() (bool, error) {
-	// TODO: Implement and use in unit tests.
-	return false, nil
-}
-
 func (lc *localChain) CanRestoreRewardEligibility() (bool, error) {
 	panic("unsupported")
 }
diff --git a/pkg/tbtc/heartbeat.go b/pkg/tbtc/heartbeat.go
index 65821f65a4..b38c888bda 100644
--- a/pkg/tbtc/heartbeat.go
+++ b/pkg/tbtc/heartbeat.go
@@ -123,9 +123,12 @@ func newHeartbeatAction(
 
 func (ha *heartbeatAction) execute() error {
 	// Do not execute the heartbeat action if the operator is unstaking.
-	isUnstaking, err := ha.chain.IsOperatorUnstaking()
+	isUnstaking, err := ha.isOperatorUnstaking()
 	if err != nil {
-		return fmt.Errorf("failed to check if the operator is unstaking")
+		return fmt.Errorf(
+			"failed to check if the operator is unstaking [%v]",
+			err,
+		)
 	}
 
 	if isUnstaking {
@@ -247,6 +250,27 @@ func (ha *heartbeatAction) actionType() WalletActionType {
 	return ActionHeartbeat
 }
 
+func (ha *heartbeatAction) isOperatorUnstaking() (bool, error) {
+	stakingProvider, isRegistered, err := ha.chain.OperatorToStakingProvider()
+	if err != nil {
+		return false, fmt.Errorf("failed to get staking provider for operator")
+	}
+
+	if !isRegistered {
+		return false, fmt.Errorf("staking provider not registered for operator")
+	}
+
+	// Eligible stake is defined as the currently authorized stake minus the
+	// pending authorization decrease.
+	eligibleStake, err := ha.chain.EligibleStake(stakingProvider)
+	if err != nil {
+		return false, fmt.Errorf("failed to check eligible stake for operator")
+	}
+
+	// The operator is considered unstaking if their eligible stake is `0`.
+	return eligibleStake.Cmp(big.NewInt(0)) == 0, nil
+}
+
 // heartbeatFailureCounter holds counters keeping track of consecutive
 // heartbeat failures. Each wallet has a separate counter. The key used in
 // the map is the uncompressed public key (with 04 prefix) of the wallet.