From 4e19c4d2ff41df670d33f2ec8353f0e7adea2598 Mon Sep 17 00:00:00 2001
From: Tomasz Slabon <tomasz.slabon@keep.network>
Date: Tue, 4 Jun 2024 11:14:55 +0200
Subject: [PATCH] Added waiting for wallet closure confirmation

---
 pkg/chain/ethereum/tbtc.go | 15 ++++++++++
 pkg/tbtc/chain.go          |  2 ++
 pkg/tbtc/chain_test.go     |  4 +++
 pkg/tbtc/node.go           | 56 ++++++++++++++++++++++++++++++++++++++
 pkg/tbtc/tbtc.go           | 16 +++++++++++
 pkg/tbtcpg/chain_test.go   |  4 +++
 6 files changed, 97 insertions(+)

diff --git a/pkg/chain/ethereum/tbtc.go b/pkg/chain/ethereum/tbtc.go
index 58c6058c4a..a970bbff6e 100644
--- a/pkg/chain/ethereum/tbtc.go
+++ b/pkg/chain/ethereum/tbtc.go
@@ -1415,6 +1415,21 @@ func (tc *TbtcChain) PastNewWalletRegisteredEvents(
 	return convertedEvents, err
 }
 
+func (tc *TbtcChain) IsWalletRegistered(EcdsaWalletID [32]byte) (bool, error) {
+	isWalletRegistered, err := tc.walletRegistry.IsWalletRegistered(
+		EcdsaWalletID,
+	)
+	if err != nil {
+		return false, fmt.Errorf(
+			"cannot check if wallet with ECDSA ID [0x%x] is registered: [%v]",
+			EcdsaWalletID,
+			err,
+		)
+	}
+
+	return isWalletRegistered, nil
+}
+
 func (tc *TbtcChain) GetWallet(
 	walletPublicKeyHash [20]byte,
 ) (*tbtc.WalletChainData, error) {
diff --git a/pkg/tbtc/chain.go b/pkg/tbtc/chain.go
index 2c4d25506c..2a851d4264 100644
--- a/pkg/tbtc/chain.go
+++ b/pkg/tbtc/chain.go
@@ -245,6 +245,8 @@ type WalletClosedEvent struct {
 // BridgeChain defines the subset of the TBTC chain interface that pertains
 // specifically to the tBTC Bridge operations.
 type BridgeChain interface {
+	IsWalletRegistered(EcdsaWalletID [32]byte) (bool, error)
+
 	// GetWallet gets the on-chain data for the given wallet. Returns an error
 	// if the wallet was not found.
 	GetWallet(walletPublicKeyHash [20]byte) (*WalletChainData, error)
diff --git a/pkg/tbtc/chain_test.go b/pkg/tbtc/chain_test.go
index 32e2998f54..2f7f57e5dc 100644
--- a/pkg/tbtc/chain_test.go
+++ b/pkg/tbtc/chain_test.go
@@ -843,6 +843,10 @@ func buildDepositRequestKey(
 	return sha256.Sum256(append(fundingTxHash[:], buffer...))
 }
 
+func (lc *localChain) IsWalletRegistered(EcdsaWalletID [32]byte) (bool, error) {
+	panic("unsupported")
+}
+
 func (lc *localChain) GetWallet(walletPublicKeyHash [20]byte) (
 	*WalletChainData,
 	error,
diff --git a/pkg/tbtc/node.go b/pkg/tbtc/node.go
index af0c37293f..66574d83bc 100644
--- a/pkg/tbtc/node.go
+++ b/pkg/tbtc/node.go
@@ -8,6 +8,7 @@ import (
 	"math/big"
 	"sync"
 
+	"github.com/keep-network/keep-common/pkg/chain/ethereum"
 	"github.com/keep-network/keep-core/pkg/bitcoin"
 	"github.com/keep-network/keep-core/pkg/chain"
 
@@ -38,6 +39,12 @@ const (
 	// Moreover, the signature must be produced in the reasonable time.
 	// That being said, the value `5` seems to be reasonable trade-off.
 	signingAttemptsLimit = 5
+
+	// walletClosureConfirmationBlocks determines the period used when waiting
+	// for the wallet closure confirmation. This period ensures the wallet has
+	// been definitely closed and the closing transaction will not be removed by
+	// a chain reorganization.
+	walletClosureConfirmationBlocks = 32
 )
 
 // TODO: Unit tests for `node.go`.
@@ -1075,6 +1082,55 @@ func processCoordinationResult(node *node, result *coordinationResult) {
 	}
 }
 
+// handleWalletClosure handles the wallet termination or closing process.
+func (n *node) handleWalletClosure(walletID [32]byte) error {
+	blockCounter, err := n.chain.BlockCounter()
+	if err != nil {
+		return fmt.Errorf("error getting block counter [%w]", err)
+	}
+
+	currentBlock, err := blockCounter.CurrentBlock()
+	if err != nil {
+		return fmt.Errorf("error getting current block [%w]", err)
+	}
+
+	// To verify there was no chain reorg and the wallet is really closed check
+	// if it is registered. Both terminated and closed wallets are removed
+	// from the ECDSA registry.
+	stateCheck := func() (bool, error) {
+		isRegistered, err := n.chain.IsWalletRegistered(walletID)
+		if err != nil {
+			return false, err
+		}
+
+		return !isRegistered, nil
+	}
+
+	// Wait a significant number of blocks to make sure the transaction has not
+	// been reverted for some reason, e.g. due to a chain reorganization.
+	result, err := ethereum.WaitForBlockConfirmations(
+		blockCounter,
+		currentBlock,
+		walletClosureConfirmationBlocks,
+		stateCheck,
+	)
+	if err != nil {
+		return fmt.Errorf(
+			"error while waiting for wallet closure confirmation [%w]",
+			err,
+		)
+	}
+
+	if !result {
+		return fmt.Errorf("wallet closure not confirmed")
+	}
+
+	// TODO: Continue with wallet closure handling: save key material and remove
+	//       the wallet from coordination mechanism.
+
+	return nil
+}
+
 // waitForBlockFn represents a function blocking the execution until the given
 // block height.
 type waitForBlockFn func(context.Context, uint64) error
diff --git a/pkg/tbtc/tbtc.go b/pkg/tbtc/tbtc.go
index e90356b2e9..17fb285a99 100644
--- a/pkg/tbtc/tbtc.go
+++ b/pkg/tbtc/tbtc.go
@@ -264,6 +264,22 @@ func Initialize(
 		}()
 	})
 
+	_ = chain.OnWalletClosed(func(event *WalletClosedEvent) {
+		go func() {
+			// TODO: Most likely event deduplication is needed.
+
+			logger.Infof(
+				"Wallet with ID [0x%x] has been closed at block [%v]",
+				event.WalletID,
+				event.BlockNumber,
+			)
+
+			node.handleWalletClosure(
+				event.WalletID,
+			)
+		}()
+	})
+
 	return nil
 }
 
diff --git a/pkg/tbtcpg/chain_test.go b/pkg/tbtcpg/chain_test.go
index cc0ddc48e2..a87ba98f56 100644
--- a/pkg/tbtcpg/chain_test.go
+++ b/pkg/tbtcpg/chain_test.go
@@ -1020,6 +1020,10 @@ func (lc *LocalChain) SetAverageBlockTime(averageBlockTime time.Duration) {
 	lc.averageBlockTime = averageBlockTime
 }
 
+func (lc *LocalChain) IsWalletRegistered(EcdsaWalletID [32]byte) (bool, error) {
+	panic("unsupported")
+}
+
 func (lc *LocalChain) GetWallet(walletPublicKeyHash [20]byte) (
 	*tbtc.WalletChainData,
 	error,