Addressing smartnode-issue-572

This is commit 1/3 to address smartnode-issue-572. This commit includes changes to:
1) add an enable-partial-rebuild flag to the rebuild flow, 2) return error data per key from the API and 3) add logging to inform the user of the outcomes.

Update recover-keys.go

Update go.mod

Author: SolezOfScience

client/wallet.go

@@ -82,8 +82,8 @@ func (r *WalletRequester) Masquerade(address common.Address) (*types.ApiResponse
 }
 
 // Rebuild the validator keys associated with the wallet
-func (r *WalletRequester) Rebuild() (*types.ApiResponse[api.WalletRebuildData], error) {
-	return client.SendGetRequest[api.WalletRebuildData](r, "rebuild", "Rebuild", nil)
+func (r *WalletRequester) Rebuild(enablePartialRebuild bool) (*types.ApiResponse[api.WalletRebuildData], error) {
+	return client.SendGetRequest[api.WalletRebuildData](r, "rebuild", "Rebuild", map[string]string{"enable-partial-rebuild": strconv.FormatBool(enablePartialRebuild)})
 }
 
 // Recover wallet

go.mod

@@ -173,6 +173,8 @@ require (
 	rsc.io/tmplfunc v0.0.3 // indirect
 )
 
+require golang.org/x/exp v0.0.0-20240325151524-a685a6edb6d8
+
 replace github.com/wealdtech/go-merkletree v1.0.1-0.20190605192610-2bb163c2ea2a => github.com/rocket-pool/go-merkletree v1.0.1-0.20220406020931-c262d9b976dd
 
 //replace github.com/rocket-pool/node-manager-core => ../node-manager-core

rocketpool-cli/commands/wallet/commands.go

@@ -94,6 +94,9 @@ func RegisterCommands(app *cli.App, name string, aliases []string) {
 				Name:    "rebuild",
 				Aliases: []string{"b"},
 				Usage:   "Rebuild validator keystores from derived keys",
+				Flags: []cli.Flag{
+					enablePartialRebuild,
+				},
 				Action: func(c *cli.Context) error {
 					// Validate args
 					utils.ValidateArgCount(c, 0)

rocketpool-cli/commands/wallet/rebuild.go

@@ -2,11 +2,11 @@ package wallet
 
 import (
 	"fmt"
+	"github.com/rocket-pool/smartnode/v2/rocketpool-cli/utils"
 	"os"
 
 	"github.com/rocket-pool/node-manager-core/wallet"
 	"github.com/rocket-pool/smartnode/v2/rocketpool-cli/client"
-	"github.com/rocket-pool/smartnode/v2/rocketpool-cli/utils"
 	"github.com/urfave/cli/v2"
 )
 
@@ -54,41 +54,52 @@ func rebuildWallet(c *cli.Context) error {
 		}(customKeyPasswordFile)
 	}
 
+	var enablePartialRebuildValue = false
+	if enablePartialRebuild.Name != "" {
+		enablePartialRebuildValue = c.Bool(enablePartialRebuild.Name)
+	}
+
 	// Log
 	fmt.Println("Rebuilding node validator keystores...")
+	fmt.Printf("Partial rebuild enabled: %s.\n", enablePartialRebuild.Name)
 
 	// Rebuild wallet
-	response, err := rp.Api.Wallet.Rebuild()
-	if err != nil {
-		return err
+	response, _ := rp.Api.Wallet.Rebuild(enablePartialRebuildValue)
+
+	// Handle and print failure reasons with associated public keys
+	if len(response.Data.FailureReasons) > 0 {
+		fmt.Println("Failure reasons:")
+		for pubkey, reason := range response.Data.FailureReasons {
+			fmt.Printf("Public Key: %s - Failure Reason: %s\n", pubkey.Hex(), reason)
+		}
+	} else {
+		fmt.Println("No failures reported.")
 	}
 
-	// Log & return
-	fmt.Println("The node wallet was successfully rebuilt.")
-	if len(response.Data.ValidatorKeys) > 0 {
+	fmt.Println("The response for rebuilding the node wallet was successfully received.")
+	if len(response.Data.RebuiltValidatorKeys) > 0 {
 		fmt.Println("Validator keys:")
-		for _, key := range response.Data.ValidatorKeys {
+		for _, key := range response.Data.RebuiltValidatorKeys {
 			fmt.Println(key.Hex())
 		}
-		fmt.Println()
-	} else {
-		fmt.Println("No validator keys were found.")
-	}
 
-	if !utils.Confirm("Would you like to restart your Validator Client now so it can attest with the recovered keys?") {
-		fmt.Println("Please restart the Validator Client manually at your earliest convenience to load the keys.")
-		return nil
-	}
+		if !utils.Confirm("Would you like to restart your Validator Client now so it can attest with the recovered keys?") {
+			fmt.Println("Please restart the Validator Client manually at your earliest convenience to load the keys.")
+			return nil
+		}
 
-	// Restart the VC
-	fmt.Println("Restarting Validator Client...")
-	_, err = rp.Api.Service.RestartVc()
-	if err != nil {
-		fmt.Printf("Error restarting Validator Client: %s\n", err.Error())
-		fmt.Println("Please restart the Validator Client manually at your earliest convenience to load the keys.")
-		return nil
+		// Restart the VC
+		fmt.Println("Restarting Validator Client...")
+		_, err = rp.Api.Service.RestartVc()
+		if err != nil {
+			fmt.Printf("Error restarting Validator Client: %s\n", err.Error())
+			fmt.Println("Please restart the Validator Client manually at your earliest convenience to load the keys.")
+			return nil
+		}
+		fmt.Println("Validator Client restarted successfully.")
+	} else {
+		fmt.Println("No validator keys were found.")
 	}
-	fmt.Println("Validator Client restarted successfully.")
 
 	return nil
 }

rocketpool-cli/commands/wallet/utils.go

@@ -57,6 +57,11 @@ var (
 		Aliases: []string{"a"},
 		Usage:   "If you are recovering a wallet that was not generated by the Smartnode and don't know the derivation path or index of it, enter the address here. The Smartnode will search through its library of paths and indices to try to find it.",
 	}
+	enablePartialRebuild = &cli.StringSliceFlag{
+		Name:    "enable-partial-rebuild",
+		Aliases: []string{"epr"},
+		Usage:   "Allows the wallet rebuild process to partially succeed, responding with public keys for successfully rebuilt targets and errors for rebuild failures",
+	}
 )
 
 // Prompt for a new wallet password

rocketpool-daemon/api/wallet/rebuild.go

@@ -2,10 +2,11 @@ package wallet
 
 import (
 	"fmt"
-	"net/url"
-
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/gorilla/mux"
+	"github.com/rocket-pool/node-manager-core/utils/input"
+	key_recovery_manager "github.com/rocket-pool/smartnode/v2/rocketpool-daemon/common/validator/key-recovery-manager"
+	"net/url"
 
 	"github.com/rocket-pool/node-manager-core/api/server"
 	"github.com/rocket-pool/node-manager-core/api/types"
@@ -24,7 +25,9 @@ func (f *walletRebuildContextFactory) Create(args url.Values) (*walletRebuildCon
 	c := &walletRebuildContext{
 		handler: f.handler,
 	}
-	return c, nil
+	inputError := server.ValidateOptionalArg("enable-partial-rebuild", args, input.ValidateBool, &c.enablePartialRebuild, nil)
+
+	return c, inputError
 }
 
 func (f *walletRebuildContextFactory) RegisterRoute(router *mux.Router) {
@@ -38,12 +41,15 @@ func (f *walletRebuildContextFactory) RegisterRoute(router *mux.Router) {
 // ===============
 
 type walletRebuildContext struct {
-	handler *WalletHandler
+	handler              *WalletHandler
+	enablePartialRebuild bool
 }
 
 func (c *walletRebuildContext) PrepareData(data *api.WalletRebuildData, opts *bind.TransactOpts) (types.ResponseStatus, error) {
 	sp := c.handler.serviceProvider
 	vMgr := sp.GetValidatorManager()
+	partialKeyRecoveryManager := key_recovery_manager.NewPartialRecoveryManager(vMgr)
+	strictKeyRecoveryManager := key_recovery_manager.NewStrictRecoveryManager(vMgr)
 
 	// Requirements
 	err := sp.RequireWalletReady()
@@ -56,7 +62,11 @@ func (c *walletRebuildContext) PrepareData(data *api.WalletRebuildData, opts *bi
 	}
 
 	// Recover validator keys
-	data.ValidatorKeys, err = vMgr.RecoverMinipoolKeys(false)
+	if c.enablePartialRebuild {
+		data.RebuiltValidatorKeys, data.FailureReasons, err = partialKeyRecoveryManager.RecoverMinipoolKeys()
+	} else {
+		data.RebuiltValidatorKeys, data.FailureReasons, err = strictKeyRecoveryManager.RecoverMinipoolKeys()
+	}
 	if err != nil {
 		return types.ResponseStatus_Error, fmt.Errorf("error recovering minipool keys: %w", err)
 	}

rocketpool-daemon/common/validator/key-recovery-manager/dry-run-key-recovery-manager.go

@@ -0,0 +1,149 @@
+package key_recovery_manager
+
+import (
+	"fmt"
+	"github.com/rocket-pool/node-manager-core/beacon"
+	"github.com/rocket-pool/node-manager-core/utils"
+	"github.com/rocket-pool/smartnode/v2/rocketpool-daemon/common/validator"
+	"golang.org/x/exp/maps"
+	"strings"
+)
+
+type DryRunKeyRecoveryManager struct {
+	manager *validator.ValidatorManager
+}
+
+func NewDryRunKeyRecoveryManager(m *validator.ValidatorManager) *DryRunKeyRecoveryManager {
+	return &DryRunKeyRecoveryManager{
+		manager: m,
+	}
+}
+
+func (d *DryRunKeyRecoveryManager) RecoverMinipoolKeys() ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error, error) {
+	status, err := d.manager.GetWalletStatus()
+	if err != nil {
+		return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err
+	}
+
+	rpNode, mpMgr, err := d.manager.InitializeBindings(status)
+	if err != nil {
+		return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err
+	}
+
+	publicKeys, err := d.manager.GetMinipools(rpNode, mpMgr)
+	if err != nil {
+		return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err
+	}
+
+	recoveredCustomPublicKeys, unrecoverableCustomPublicKeys, _ := d.checkForAndRecoverCustomKeys(publicKeys)
+	recoveredPublicKeys, unrecoverablePublicKeys := d.recoverConventionalKeys(publicKeys)
+
+	allRecoveredPublicKeys := []beacon.ValidatorPubkey{}
+	allRecoveredPublicKeys = append(allRecoveredPublicKeys, maps.Keys(recoveredCustomPublicKeys)...)
+	allRecoveredPublicKeys = append(allRecoveredPublicKeys, recoveredPublicKeys...)
+
+	for publicKey, err := range unrecoverablePublicKeys {
+		unrecoverableCustomPublicKeys[publicKey] = err
+	}
+
+	return allRecoveredPublicKeys, unrecoverablePublicKeys, nil
+}
+
+func (d *DryRunKeyRecoveryManager) checkForAndRecoverCustomKeys(publicKeys map[beacon.ValidatorPubkey]bool) (map[beacon.ValidatorPubkey]bool, map[beacon.ValidatorPubkey]error, error) {
+
+	recoveredKeys := make(map[beacon.ValidatorPubkey]bool)
+	recoveryFailures := make(map[beacon.ValidatorPubkey]error)
+	var passwords map[string]string
+
+	keyFiles, err := d.manager.LoadFiles()
+	if err != nil {
+		return recoveredKeys, recoveryFailures, err
+	}
+
+	if len(keyFiles) > 0 {
+		passwords, err = d.manager.LoadCustomKeyPasswords()
+		if err != nil {
+			return recoveredKeys, recoveryFailures, err
+		}
+
+		for _, file := range keyFiles {
+			keystore, err := d.manager.ReadCustomKeystore(file)
+			if err != nil {
+				continue
+			}
+
+			if _, exists := publicKeys[keystore.Pubkey]; !exists {
+				err := fmt.Errorf("custom keystore for pubkey %s not found in minipool keyset", keystore.Pubkey.Hex())
+				recoveryFailures[keystore.Pubkey] = err
+				continue
+			}
+
+			formattedPublicKey := strings.ToUpper(utils.RemovePrefix(keystore.Pubkey.Hex()))
+			password, exists := passwords[formattedPublicKey]
+			if !exists {
+				err := fmt.Errorf("custom keystore for pubkey %s needs a password, but none was provided", keystore.Pubkey.Hex())
+				recoveryFailures[keystore.Pubkey] = err
+				continue
+			}
+
+			privateKey, err := d.manager.DecryptCustomKeystore(keystore, password)
+			if err != nil {
+				err := fmt.Errorf("error recreating private key for validator %s: %w", keystore.Pubkey.Hex(), err)
+				recoveryFailures[keystore.Pubkey] = err
+				continue
+			}
+
+			reconstructedPublicKey := beacon.ValidatorPubkey(privateKey.PublicKey().Marshal())
+			if reconstructedPublicKey != keystore.Pubkey {
+				err := fmt.Errorf("private keystore file %s claims to be for validator %s but it's for validator %s", file.Name(), keystore.Pubkey.Hex(), reconstructedPublicKey.Hex())
+				recoveryFailures[keystore.Pubkey] = err
+				continue
+			}
+
+			recoveredKeys[reconstructedPublicKey] = true
+		}
+	}
+
+	return recoveredKeys, recoveryFailures, nil
+}
+
+func (d *DryRunKeyRecoveryManager) recoverConventionalKeys(publicKeys map[beacon.ValidatorPubkey]bool) ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error) {
+	recoveredPublicKeys := []beacon.ValidatorPubkey{}
+	unrecoverablePublicKeys := map[beacon.ValidatorPubkey]error{}
+
+	bucketStart := uint64(0)
+	for {
+		if bucketStart >= bucketLimit {
+			break
+		}
+		bucketEnd := bucketStart + bucketSize
+		if bucketEnd > bucketLimit {
+			bucketEnd = bucketLimit
+		}
+
+		keys, err := d.manager.GetValidatorKeys(bucketStart, bucketEnd-bucketStart)
+		if err != nil {
+			continue
+		}
+
+		for _, validatorKey := range keys {
+			if exists := publicKeys[validatorKey.PublicKey]; exists {
+				delete(publicKeys, validatorKey.PublicKey)
+				recoveredPublicKeys = append(recoveredPublicKeys, validatorKey.PublicKey)
+			} else {
+				err := fmt.Errorf("keystore for pubkey %s not found in minipool keyset", validatorKey.PublicKey)
+				unrecoverablePublicKeys[validatorKey.PublicKey] = err
+				continue
+			}
+		}
+
+		if len(publicKeys) == 0 {
+			// All keys have been recovered.
+			break
+		}
+
+		bucketStart = bucketEnd
+	}
+
+	return recoveredPublicKeys, unrecoverablePublicKeys
+}

rocketpool-daemon/common/validator/key-recovery-manager/key-recovery-manager.go

@@ -0,0 +1,14 @@
+package key_recovery_manager
+
+import (
+	"github.com/rocket-pool/node-manager-core/beacon"
+)
+
+type KeyRecoveryManager interface {
+	RecoverMinipoolKeys() ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error, error)
+}
+
+const (
+	bucketSize  uint64 = 20
+	bucketLimit uint64 = 2000
+)