From accb4b1190d4eb5e751a3070bf582874b1e6ccd3 Mon Sep 17 00:00:00 2001
From: Tomasz Slabon <tomasz.slabon@keep.network>
Date: Thu, 6 Jun 2024 16:42:42 +0200
Subject: [PATCH] Added wallet archiving on wallet closed event

---
 pkg/tbtc/chain.go      |  8 +++++++
 pkg/tbtc/chain_test.go |  6 +++++
 pkg/tbtc/node.go       | 28 +++++++++++++++++++++--
 pkg/tbtc/registry.go   | 52 ++++++++++++++++++++++++++++++++++++++++++
 pkg/tbtc/tbtc.go       |  9 +++++++-
 5 files changed, 100 insertions(+), 3 deletions(-)

diff --git a/pkg/tbtc/chain.go b/pkg/tbtc/chain.go
index 2a851d4264..9eaa645b86 100644
--- a/pkg/tbtc/chain.go
+++ b/pkg/tbtc/chain.go
@@ -262,6 +262,14 @@ type BridgeChain interface {
 	// according to the on-chain Bridge rules.
 	ComputeMainUtxoHash(mainUtxo *bitcoin.UnspentTransactionOutput) [32]byte
 
+	// PastNewWalletRegisteredEvents fetches past new wallet registered events
+	// according to the provided filter or unfiltered if the filter is nil.
+	// Returned events are sorted by the block number in the ascending order,
+	// i.e. the latest event is at the end of the slice.
+	PastNewWalletRegisteredEvents(
+		filter *NewWalletRegisteredEventFilter,
+	) ([]*NewWalletRegisteredEvent, error)
+
 	// PastDepositRevealedEvents fetches past deposit reveal events according
 	// to the provided filter or unfiltered if the filter is nil. Returned
 	// events are sorted by the block number in the ascending order, i.e. the
diff --git a/pkg/tbtc/chain_test.go b/pkg/tbtc/chain_test.go
index 2f7f57e5dc..64effd2b3b 100644
--- a/pkg/tbtc/chain_test.go
+++ b/pkg/tbtc/chain_test.go
@@ -697,6 +697,12 @@ func (lc *localChain) GetInactivityClaimNonce(walletID [32]byte) (*big.Int, erro
 	return big.NewInt(int64(nonce)), nil
 }
 
+func (lc *localChain) PastNewWalletRegisteredEvents(
+	filter *NewWalletRegisteredEventFilter,
+) ([]*NewWalletRegisteredEvent, error) {
+	panic("unsupported")
+}
+
 func (lc *localChain) PastDepositRevealedEvents(
 	filter *DepositRevealedEventFilter,
 ) ([]*DepositRevealedEvent, error) {
diff --git a/pkg/tbtc/node.go b/pkg/tbtc/node.go
index 66574d83bc..1234449f9d 100644
--- a/pkg/tbtc/node.go
+++ b/pkg/tbtc/node.go
@@ -1125,8 +1125,32 @@ func (n *node) handleWalletClosure(walletID [32]byte) error {
 		return fmt.Errorf("wallet closure not confirmed")
 	}
 
-	// TODO: Continue with wallet closure handling: save key material and remove
-	//       the wallet from coordination mechanism.
+	events, err := n.chain.PastNewWalletRegisteredEvents(
+		&NewWalletRegisteredEventFilter{
+			EcdsaWalletID: [][32]byte{walletID},
+		},
+	)
+	if err != nil {
+		return fmt.Errorf("could not get past new wallet registered events")
+	}
+
+	// There should be only one event returned and the ECDSA wallet ID should
+	// match the requested wallet ID. These errors should never happen, but
+	// check just in case.
+	if len(events) != 1 {
+		return fmt.Errorf("wrong number of past new wallet registered events")
+	}
+
+	if events[0].EcdsaWalletID != walletID {
+		return fmt.Errorf("wrong past new wallet registered event returned")
+	}
+
+	walletPublicKeyHash := events[0].WalletPublicKeyHash
+
+	err = n.walletRegistry.archiveWallet(walletPublicKeyHash)
+	if err != nil {
+		return fmt.Errorf("failed to archive the wallet: [%v]", err)
+	}
 
 	return nil
 }
diff --git a/pkg/tbtc/registry.go b/pkg/tbtc/registry.go
index 80334f7475..6f469862de 100644
--- a/pkg/tbtc/registry.go
+++ b/pkg/tbtc/registry.go
@@ -156,6 +156,45 @@ func (wr *walletRegistry) getWalletByPublicKeyHash(
 	return wallet{}, false
 }
 
+func (wr *walletRegistry) archiveWallet(
+	walletPublicKeyHash [20]byte,
+) error {
+	wr.mutex.Lock()
+	defer wr.mutex.Unlock()
+
+	var walletPublicKey *ecdsa.PublicKey
+
+	for _, value := range wr.walletCache {
+		if value.walletPublicKeyHash == walletPublicKeyHash {
+			// All signers belong to one wallet. Take the wallet public key from
+			//  the first signer.
+			walletPublicKey = value.signers[0].wallet.publicKey
+		}
+	}
+
+	if walletPublicKey == nil {
+		logger.Infof(
+			"node does not control wallet with public key hash [0x%x]; "+
+				"quitting wallet archiving",
+			walletPublicKeyHash,
+		)
+		return nil
+	}
+
+	walletStorageKey := getWalletStorageKey(walletPublicKey)
+
+	// Archive the entire wallet storage.
+	err := wr.walletStorage.archiveWallet(walletStorageKey)
+	if err != nil {
+		return fmt.Errorf("could not archive wallet: [%v]", err)
+	}
+
+	// Remove the wallet from the wallet cache.
+	delete(wr.walletCache, walletStorageKey)
+
+	return nil
+}
+
 // walletStorage is the component that persists data of the wallets managed
 // by the given node using the underlying persistence layer. It should be
 // used directly only by the walletRegistry.
@@ -194,6 +233,19 @@ func (ws *walletStorage) saveSigner(signer *signer) error {
 	return nil
 }
 
+func (ws *walletStorage) archiveWallet(walletStoragePath string) error {
+	err := ws.persistence.Archive(walletStoragePath)
+	if err != nil {
+		return fmt.Errorf(
+			"could not archive wallet storage using the "+
+				"underlying persistence layer: [%w]",
+			err,
+		)
+	}
+
+	return nil
+}
+
 // loadSigners loads all signers stored using the underlying persistence layer.
 // This function should not be called from any other place than walletRegistry.
 func (ws *walletStorage) loadSigners() map[string][]*signer {
diff --git a/pkg/tbtc/tbtc.go b/pkg/tbtc/tbtc.go
index 17fb285a99..937851e98a 100644
--- a/pkg/tbtc/tbtc.go
+++ b/pkg/tbtc/tbtc.go
@@ -274,9 +274,16 @@ func Initialize(
 				event.BlockNumber,
 			)
 
-			node.handleWalletClosure(
+			err := node.handleWalletClosure(
 				event.WalletID,
 			)
+			if err != nil {
+				logger.Errorf(
+					"Failure while handling wallet with ID [0x%x]: [%v]",
+					event.WalletID,
+					err,
+				)
+			}
 		}()
 	})