Skip to content

Commit

Permalink
Addressing smartnode-issue-572
Browse files Browse the repository at this point in the history
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
  • Loading branch information
SolezOfScience committed Aug 11, 2024
1 parent b3fdb9d commit 9bd5510
Show file tree
Hide file tree
Showing 13 changed files with 707 additions and 42 deletions.
4 changes: 2 additions & 2 deletions client/wallet.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
3 changes: 3 additions & 0 deletions rocketpool-cli/commands/wallet/commands.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
59 changes: 35 additions & 24 deletions rocketpool-cli/commands/wallet/rebuild.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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
}
5 changes: 5 additions & 0 deletions rocketpool-cli/commands/wallet/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
20 changes: 15 additions & 5 deletions rocketpool-daemon/api/wallet/rebuild.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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) {
Expand All @@ -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()
Expand All @@ -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)
}
Expand Down
Original file line number Diff line number Diff line change
@@ -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
}
Original file line number Diff line number Diff line change
@@ -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
)
Loading

0 comments on commit 9bd5510

Please sign in to comment.