From 9bd551093d9de7eddb5cc7f985989bdf18c8b585 Mon Sep 17 00:00:00 2001 From: SolezOfScience Date: Sat, 10 Aug 2024 10:34:29 -0700 Subject: [PATCH 1/2] 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 --- client/wallet.go | 4 +- go.mod | 2 + rocketpool-cli/commands/wallet/commands.go | 3 + rocketpool-cli/commands/wallet/rebuild.go | 59 ++++--- rocketpool-cli/commands/wallet/utils.go | 5 + rocketpool-daemon/api/wallet/rebuild.go | 20 ++- .../dry-run-key-recovery-manager.go | 149 ++++++++++++++++ .../key-recovery-manager.go | 14 ++ .../partial-key-recovery-manager.go | 160 +++++++++++++++++ .../strict-key-recovery-manager.go | 166 ++++++++++++++++++ .../common/validator/recover-keys.go | 6 - .../common/validator/validator-manager.go | 156 +++++++++++++++- shared/types/api/wallet.go | 5 +- 13 files changed, 707 insertions(+), 42 deletions(-) create mode 100644 rocketpool-daemon/common/validator/key-recovery-manager/dry-run-key-recovery-manager.go create mode 100644 rocketpool-daemon/common/validator/key-recovery-manager/key-recovery-manager.go create mode 100644 rocketpool-daemon/common/validator/key-recovery-manager/partial-key-recovery-manager.go create mode 100644 rocketpool-daemon/common/validator/key-recovery-manager/strict-key-recovery-manager.go diff --git a/client/wallet.go b/client/wallet.go index ef271095f..674d02297 100644 --- a/client/wallet.go +++ b/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 diff --git a/go.mod b/go.mod index 5fc7a5713..e2698c6e5 100644 --- a/go.mod +++ b/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 diff --git a/rocketpool-cli/commands/wallet/commands.go b/rocketpool-cli/commands/wallet/commands.go index 18429c7ec..18f8e8ae4 100644 --- a/rocketpool-cli/commands/wallet/commands.go +++ b/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) diff --git a/rocketpool-cli/commands/wallet/rebuild.go b/rocketpool-cli/commands/wallet/rebuild.go index 5304ed167..14627f689 100644 --- a/rocketpool-cli/commands/wallet/rebuild.go +++ b/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 } diff --git a/rocketpool-cli/commands/wallet/utils.go b/rocketpool-cli/commands/wallet/utils.go index 7c399334e..de0025999 100644 --- a/rocketpool-cli/commands/wallet/utils.go +++ b/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 diff --git a/rocketpool-daemon/api/wallet/rebuild.go b/rocketpool-daemon/api/wallet/rebuild.go index a4ac22ea1..8eca93ede 100644 --- a/rocketpool-daemon/api/wallet/rebuild.go +++ b/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) } diff --git a/rocketpool-daemon/common/validator/key-recovery-manager/dry-run-key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager/dry-run-key-recovery-manager.go new file mode 100644 index 000000000..fd854232e --- /dev/null +++ b/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 +} diff --git a/rocketpool-daemon/common/validator/key-recovery-manager/key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager/key-recovery-manager.go new file mode 100644 index 000000000..dab48bdd5 --- /dev/null +++ b/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 +) diff --git a/rocketpool-daemon/common/validator/key-recovery-manager/partial-key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager/partial-key-recovery-manager.go new file mode 100644 index 000000000..c96a93c9a --- /dev/null +++ b/rocketpool-daemon/common/validator/key-recovery-manager/partial-key-recovery-manager.go @@ -0,0 +1,160 @@ +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 PartialRecoveryManager struct { + manager *validator.ValidatorManager +} + +func NewPartialRecoveryManager(m *validator.ValidatorManager) *PartialRecoveryManager { + return &PartialRecoveryManager{ + manager: m, + } +} + +func (p PartialRecoveryManager) RecoverMinipoolKeys() ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error, error) { + status, err := p.manager.GetWalletStatus() + if err != nil { + return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err + } + + rpNode, mpMgr, err := p.manager.InitializeBindings(status) + if err != nil { + return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err + } + + publicKeys, err := p.manager.GetMinipools(rpNode, mpMgr) + if err != nil { + return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err + } + + recoveredCustomPublicKeys, unrecoverableCustomPublicKeys, _ := p.checkForAndRecoverCustomKeys(publicKeys) + recoveredPublicKeys, unrecoverablePublicKeys := p.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 (p PartialRecoveryManager) 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 := p.manager.LoadFiles() + if err != nil { + return recoveredKeys, recoveryFailures, err + } + + if len(keyFiles) > 0 { + passwords, err = p.manager.LoadCustomKeyPasswords() + if err != nil { + return recoveredKeys, recoveryFailures, err + } + + for _, file := range keyFiles { + keystore, err := p.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 := p.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 + } + + if err := p.manager.StoreValidatorKey(&privateKey, keystore.Path); err != nil { + recoveryFailures[reconstructedPublicKey] = fmt.Errorf("error storing private keystore for %s: %w", reconstructedPublicKey.Hex(), err) + } else { + recoveredKeys[reconstructedPublicKey] = true + } + + delete(publicKeys, keystore.Pubkey) + } + } + + return recoveredKeys, recoveryFailures, nil +} + +func (p *PartialRecoveryManager) 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 := p.manager.GetValidatorKeys(bucketStart, bucketEnd-bucketStart) + if err != nil { + continue + } + + for _, validatorKey := range keys { + delete(publicKeys, validatorKey.PublicKey) + if exists := publicKeys[validatorKey.PublicKey]; exists { + if err := p.manager.SaveValidatorKey(validatorKey); err != nil { + unrecoverablePublicKeys[validatorKey.PublicKey] = err + } else { + 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 +} diff --git a/rocketpool-daemon/common/validator/key-recovery-manager/strict-key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager/strict-key-recovery-manager.go new file mode 100644 index 000000000..ff5079c70 --- /dev/null +++ b/rocketpool-daemon/common/validator/key-recovery-manager/strict-key-recovery-manager.go @@ -0,0 +1,166 @@ +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 StrictRecoveryManager struct { + manager *validator.ValidatorManager +} + +func NewStrictRecoveryManager(m *validator.ValidatorManager) *StrictRecoveryManager { + return &StrictRecoveryManager{ + manager: m, + } +} + +func (s *StrictRecoveryManager) RecoverMinipoolKeys() ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error, error) { + status, err := s.manager.GetWalletStatus() + if err != nil { + return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err + } + + rpNode, mpMgr, err := s.manager.InitializeBindings(status) + if err != nil { + return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err + } + + publicKeys, err := s.manager.GetMinipools(rpNode, mpMgr) + if err != nil { + return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err + } + + recoveredCustomPublicKeys, unrecoverableCustomPublicKeys, err := s.checkForAndRecoverCustomKeys(publicKeys) + if err != nil { + return maps.Keys(recoveredCustomPublicKeys), unrecoverableCustomPublicKeys, err + } + + recoveredPublicKeys, unrecoverablePublicKeys := s.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 (s *StrictRecoveryManager) 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 := s.manager.LoadFiles() + if err != nil { + return recoveredKeys, recoveryFailures, err + } + + if len(keyFiles) > 0 { + passwords, err = s.manager.LoadCustomKeyPasswords() + if err != nil { + return recoveredKeys, recoveryFailures, err + } + + for _, file := range keyFiles { + keystore, err := s.manager.ReadCustomKeystore(file) + if err != nil { + return recoveredKeys, recoveryFailures, err + } + + 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 + return recoveredKeys, recoveryFailures, err + } + + 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 + return recoveredKeys, recoveryFailures, err + } + + privateKey, err := s.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 + return recoveredKeys, recoveryFailures, err + } + + 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 + return recoveredKeys, recoveryFailures, err + } + + if err := s.manager.StoreValidatorKey(&privateKey, keystore.Path); err != nil { + recoveryFailures[keystore.Pubkey] = err + return recoveredKeys, recoveryFailures, err + } + recoveredKeys[reconstructedPublicKey] = true + + delete(publicKeys, keystore.Pubkey) + } + } + + return recoveredKeys, recoveryFailures, nil +} + +func (s *StrictRecoveryManager) 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 := s.manager.GetValidatorKeys(bucketStart, bucketEnd-bucketStart) + if err != nil { + return recoveredPublicKeys, map[beacon.ValidatorPubkey]error{beacon.ValidatorPubkey{}: fmt.Errorf("error getting node's validator keys")} + } + + for _, validatorKey := range keys { + if exists := publicKeys[validatorKey.PublicKey]; exists { + delete(publicKeys, validatorKey.PublicKey) + if err := s.manager.SaveValidatorKey(validatorKey); err != nil { + unrecoverablePublicKeys[validatorKey.PublicKey] = err + return recoveredPublicKeys, unrecoverablePublicKeys + } else { + recoveredPublicKeys = append(recoveredPublicKeys, validatorKey.PublicKey) + } + } else { + err := fmt.Errorf("keystore for pubkey %s not found in minipool keyset", validatorKey.PublicKey) + unrecoverablePublicKeys[validatorKey.PublicKey] = err + return recoveredPublicKeys, unrecoverablePublicKeys + } + } + + if len(publicKeys) == 0 { + // All keys have been recovered. + break + } + + bucketStart = bucketEnd + } + + return recoveredPublicKeys, unrecoverablePublicKeys +} diff --git a/rocketpool-daemon/common/validator/recover-keys.go b/rocketpool-daemon/common/validator/recover-keys.go index 37f32a1c5..e6b8f2e27 100644 --- a/rocketpool-daemon/common/validator/recover-keys.go +++ b/rocketpool-daemon/common/validator/recover-keys.go @@ -19,12 +19,6 @@ import ( "gopkg.in/yaml.v3" ) -const ( - bucketSize uint64 = 20 - bucketLimit uint64 = 2000 - pubkeyBatchSize int = 500 -) - func (m *ValidatorManager) RecoverMinipoolKeys(testOnly bool) ([]beacon.ValidatorPubkey, error) { status, err := m.wallet.GetStatus() if err != nil { diff --git a/rocketpool-daemon/common/validator/validator-manager.go b/rocketpool-daemon/common/validator/validator-manager.go index e891b09d9..329a200c9 100644 --- a/rocketpool-daemon/common/validator/validator-manager.go +++ b/rocketpool-daemon/common/validator/validator-manager.go @@ -3,21 +3,33 @@ package validator import ( "bytes" "fmt" - + "github.com/goccy/go-json" + batch "github.com/rocket-pool/batch-query" "github.com/rocket-pool/node-manager-core/beacon" "github.com/rocket-pool/node-manager-core/eth" "github.com/rocket-pool/node-manager-core/node/validator" - "github.com/rocket-pool/node-manager-core/node/wallet" + walletnode "github.com/rocket-pool/node-manager-core/node/wallet" + walletcore "github.com/rocket-pool/node-manager-core/wallet" + "github.com/rocket-pool/rocketpool-go/v2/minipool" + "github.com/rocket-pool/rocketpool-go/v2/node" "github.com/rocket-pool/rocketpool-go/v2/rocketpool" "github.com/rocket-pool/smartnode/v2/rocketpool-daemon/common/utils" "github.com/rocket-pool/smartnode/v2/shared/config" + "github.com/rocket-pool/smartnode/v2/shared/types/api" eth2types "github.com/wealdtech/go-eth2-types/v2" types "github.com/wealdtech/go-eth2-types/v2" + eth2ks "github.com/wealdtech/go-eth2-wallet-encryptor-keystorev4" + "gopkg.in/yaml.v3" + "os" + "path/filepath" ) // Config const ( MaxValidatorKeyRecoverAttempts uint64 = 1000 + bucketSize uint64 = 20 + bucketLimit uint64 = 2000 + pubkeyBatchSize int = 500 ) // A validator private/public key pair @@ -32,13 +44,13 @@ type ValidatorKey struct { type ValidatorManager struct { cfg *config.SmartNodeConfig rp *rocketpool.RocketPool - wallet *wallet.Wallet + wallet *walletnode.Wallet queryMgr *eth.QueryManager keystoreManager *validator.ValidatorManager nextAccount uint64 } -func NewValidatorManager(cfg *config.SmartNodeConfig, rp *rocketpool.RocketPool, walletImpl *wallet.Wallet, queryMgr *eth.QueryManager) (*ValidatorManager, error) { +func NewValidatorManager(cfg *config.SmartNodeConfig, rp *rocketpool.RocketPool, walletImpl *walletnode.Wallet, queryMgr *eth.QueryManager) (*ValidatorManager, error) { // Make a validator manager validatorManager := validator.NewValidatorManager(cfg.GetValidatorsFolderPath()) @@ -272,3 +284,139 @@ func (m *ValidatorManager) checkIfReady() error { } return utils.CheckIfWalletReady(status) } + +func (m *ValidatorManager) GetWalletStatus() (walletcore.WalletStatus, error) { + status, err := m.wallet.GetStatus() + if err != nil { + return status, err + } + if !walletcore.IsWalletReady(status) { + return status, fmt.Errorf("wallet is not ready") + } + return status, nil +} + +func (m *ValidatorManager) InitializeBindings(status walletcore.WalletStatus) (*node.Node, *minipool.MinipoolManager, error) { + address := status.Wallet.WalletAddress + rpNode, err := node.NewNode(m.rp, address) + + if err != nil { + return nil, nil, err + } + + mpMgr, err := minipool.NewMinipoolManager(m.rp) + if err != nil { + return nil, nil, err + } + return rpNode, mpMgr, nil +} + +func (m *ValidatorManager) GetMinipools(node *node.Node, mpMgr *minipool.MinipoolManager) (map[beacon.ValidatorPubkey]bool, error) { + err := m.queryMgr.Query(nil, nil, node.ValidatingMinipoolCount) + if err != nil { + return nil, fmt.Errorf("error getting node's validating minipool count: %w", err) + } + + addresses, err := node.GetValidatingMinipoolAddresses(node.ValidatingMinipoolCount.Formatted(), nil) + if err != nil { + return nil, fmt.Errorf("error getting node's validating minipool addresses: %w", err) + } + + mps, err := mpMgr.CreateMinipoolsFromAddresses(addresses, false, nil) + if err != nil { + return nil, fmt.Errorf("error creating bindings for node's validating minipools: %w", err) + } + + publicKeySet := map[beacon.ValidatorPubkey]bool{} + zeroPublicKey := beacon.ValidatorPubkey{} + + err = m.queryMgr.BatchQuery(len(mps), pubkeyBatchSize, func(mc *batch.MultiCaller, i int) error { + mps[i].Common().Pubkey.AddToQuery(mc) + return nil + }, nil) + if err != nil { + return nil, fmt.Errorf("error getting node's validating minipool pubkeys: %w", err) + } + + for _, mp := range mps { + publicKey := mp.Common().Pubkey.Get() + if publicKey != zeroPublicKey { + publicKeySet[publicKey] = true + } + } + + return publicKeySet, nil +} + +func (m *ValidatorManager) LoadCustomKeyPasswords() (map[string]string, error) { + passwordFile := m.cfg.GetCustomKeyPasswordFilePath() + fileBytes, err := os.ReadFile(passwordFile) + if err != nil { + return nil, fmt.Errorf("password file could not be loaded: %w", err) + } + passwords := map[string]string{} + err = yaml.Unmarshal(fileBytes, &passwords) + if err != nil { + return nil, fmt.Errorf("error unmarshalling custom keystore password file: %w", err) + } + return passwords, nil +} + +func (m *ValidatorManager) ReadCustomKeystore(file os.DirEntry) (api.ValidatorKeystore, error) { + bytes, err := os.ReadFile(filepath.Join(m.cfg.GetCustomKeyPath(), file.Name())) + if err != nil { + return api.ValidatorKeystore{}, fmt.Errorf("error reading custom keystore %s: %w", file.Name(), err) + } + keystore := api.ValidatorKeystore{} + err = json.Unmarshal(bytes, &keystore) + if err != nil { + return api.ValidatorKeystore{}, fmt.Errorf("error deserializing custom keystore %s: %w", file.Name(), err) + } + return keystore, nil +} + +func (m *ValidatorManager) DecryptCustomKeystore(keystore api.ValidatorKeystore, password string) (eth2types.BLSPrivateKey, error) { + kdf, exists := keystore.Crypto["kdf"] + if !exists { + return eth2types.BLSPrivateKey{}, fmt.Errorf("error processing custom keystore: \"crypto\" didn't contain a subkey named \"kdf\"") + } + kdfMap := kdf.(map[string]interface{}) + function, exists := kdfMap["function"] + if !exists { + return eth2types.BLSPrivateKey{}, fmt.Errorf("error processing custom keystore: \"crypto.kdf\" didn't contain a subkey named \"function\"") + } + functionString := function.(string) + + encryptor := eth2ks.New(eth2ks.WithCipher(functionString)) + decryptedKey, err := encryptor.Decrypt(keystore.Crypto, password) + if err != nil { + return eth2types.BLSPrivateKey{}, fmt.Errorf("error decrypting keystore: %w", err) + } + privateKey, err := eth2types.BLSPrivateKeyFromBytes(decryptedKey) + if err != nil { + return eth2types.BLSPrivateKey{}, fmt.Errorf("error recreating private key: %w", err) + } + return *privateKey, nil +} + +func (m *ValidatorManager) LoadFiles() ([]os.DirEntry, error) { + customKeyDir := m.cfg.GetCustomKeyPath() + info, err := os.Stat(customKeyDir) + + if os.IsNotExist(err) || !info.IsDir() { + err := fmt.Errorf("error loading custom keystore location: %w", err) + return nil, err + } + + keyFiles, err := os.ReadDir(customKeyDir) + if err != nil { + err := fmt.Errorf("error enumerating custom keystores: %w", err) + return nil, err + } + + if err := eth2types.InitBLS(); err != nil { + err := fmt.Errorf("error initializing BLS: %w", err) + return nil, err + } + return keyFiles, nil +} diff --git a/shared/types/api/wallet.go b/shared/types/api/wallet.go index f9c44eec1..fbc4675f4 100644 --- a/shared/types/api/wallet.go +++ b/shared/types/api/wallet.go @@ -30,6 +30,7 @@ type WalletInitializeData struct { type WalletRecoverData struct { AccountAddress common.Address `json:"accountAddress"` ValidatorKeys []beacon.ValidatorPubkey `json:"validatorKeys"` + FailureReasons map[beacon.ValidatorPubkey]error } type WalletSearchAndRecoverData struct { @@ -38,10 +39,12 @@ type WalletSearchAndRecoverData struct { DerivationPath string `json:"derivationPath"` Index uint `json:"index"` ValidatorKeys []beacon.ValidatorPubkey `json:"validatorKeys"` + FailureReasons map[beacon.ValidatorPubkey]error } type WalletRebuildData struct { - ValidatorKeys []beacon.ValidatorPubkey `json:"validatorKeys"` + RebuiltValidatorKeys []beacon.ValidatorPubkey `json:"validatorKeys"` + FailureReasons map[beacon.ValidatorPubkey]error } type WalletExportData struct { From 3aa32abcc5ef1976f2ca3fc725207920fa918d66 Mon Sep 17 00:00:00 2001 From: SolezOfScience Date: Fri, 30 Aug 2024 12:29:49 -0700 Subject: [PATCH 2/2] Refactoring ValidatorManager and parameterizing KeyRecoveryManager to support test and partial recovery modes. This change addresses all comments from the first draft of the PR to introduce partial recovery support. This removes multiple implementations of a KeyRecoveryManager interface and instead uses class fields to determine how to appropriately enact the recover code flow. --- rocketpool-cli/commands/wallet/rebuild.go | 12 +- rocketpool-cli/commands/wallet/utils.go | 2 +- rocketpool-daemon/api/wallet/rebuild.go | 11 +- .../common/validator/key-recovery-manager.go | 188 ++++++++++++++++++ .../dry-run-key-recovery-manager.go | 149 -------------- .../key-recovery-manager.go | 14 -- .../partial-key-recovery-manager.go | 160 --------------- .../strict-key-recovery-manager.go | 166 ---------------- .../common/validator/validator-manager.go | 128 +++++++----- 9 files changed, 273 insertions(+), 557 deletions(-) create mode 100644 rocketpool-daemon/common/validator/key-recovery-manager.go delete mode 100644 rocketpool-daemon/common/validator/key-recovery-manager/dry-run-key-recovery-manager.go delete mode 100644 rocketpool-daemon/common/validator/key-recovery-manager/key-recovery-manager.go delete mode 100644 rocketpool-daemon/common/validator/key-recovery-manager/partial-key-recovery-manager.go delete mode 100644 rocketpool-daemon/common/validator/key-recovery-manager/strict-key-recovery-manager.go diff --git a/rocketpool-cli/commands/wallet/rebuild.go b/rocketpool-cli/commands/wallet/rebuild.go index 14627f689..678eb0820 100644 --- a/rocketpool-cli/commands/wallet/rebuild.go +++ b/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" ) @@ -61,13 +61,18 @@ func rebuildWallet(c *cli.Context) error { // Log fmt.Println("Rebuilding node validator keystores...") - fmt.Printf("Partial rebuild enabled: %s.\n", enablePartialRebuild.Name) + fmt.Printf("Partial rebuild enabled: %s.\n", enablePartialRebuild.Value) // Rebuild wallet - response, _ := rp.Api.Wallet.Rebuild(enablePartialRebuildValue) + response, err := rp.Api.Wallet.Rebuild(enablePartialRebuildValue) + if err != nil { + return err + } // Handle and print failure reasons with associated public keys if len(response.Data.FailureReasons) > 0 { + fmt.Println("Some keys could not be recovered. You may need to import them manually, as they are not " + + "associated with your node wallet mnemonic. See the documentation for more details.") fmt.Println("Failure reasons:") for pubkey, reason := range response.Data.FailureReasons { fmt.Printf("Public Key: %s - Failure Reason: %s\n", pubkey.Hex(), reason) @@ -76,7 +81,6 @@ func rebuildWallet(c *cli.Context) error { fmt.Println("No failures reported.") } - 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.RebuiltValidatorKeys { diff --git a/rocketpool-cli/commands/wallet/utils.go b/rocketpool-cli/commands/wallet/utils.go index de0025999..e21b733a3 100644 --- a/rocketpool-cli/commands/wallet/utils.go +++ b/rocketpool-cli/commands/wallet/utils.go @@ -59,7 +59,7 @@ var ( } enablePartialRebuild = &cli.StringSliceFlag{ Name: "enable-partial-rebuild", - Aliases: []string{"epr"}, + Aliases: []string{"p"}, Usage: "Allows the wallet rebuild process to partially succeed, responding with public keys for successfully rebuilt targets and errors for rebuild failures", } ) diff --git a/rocketpool-daemon/api/wallet/rebuild.go b/rocketpool-daemon/api/wallet/rebuild.go index 8eca93ede..cb99d3f27 100644 --- a/rocketpool-daemon/api/wallet/rebuild.go +++ b/rocketpool-daemon/api/wallet/rebuild.go @@ -5,7 +5,7 @@ import ( "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" + "github.com/rocket-pool/smartnode/v2/rocketpool-daemon/common/validator" "net/url" "github.com/rocket-pool/node-manager-core/api/server" @@ -48,8 +48,7 @@ type walletRebuildContext struct { 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) + keyRecoveryManager := validator.NewKeyRecoveryManager(vMgr, c.enablePartialRebuild, false) // Requirements err := sp.RequireWalletReady() @@ -62,11 +61,7 @@ func (c *walletRebuildContext) PrepareData(data *api.WalletRebuildData, opts *bi } // Recover validator keys - if c.enablePartialRebuild { - data.RebuiltValidatorKeys, data.FailureReasons, err = partialKeyRecoveryManager.RecoverMinipoolKeys() - } else { - data.RebuiltValidatorKeys, data.FailureReasons, err = strictKeyRecoveryManager.RecoverMinipoolKeys() - } + data.RebuiltValidatorKeys, data.FailureReasons, err = keyRecoveryManager.RecoverMinipoolKeys() if err != nil { return types.ResponseStatus_Error, fmt.Errorf("error recovering minipool keys: %w", err) } diff --git a/rocketpool-daemon/common/validator/key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager.go new file mode 100644 index 000000000..c01ddf60d --- /dev/null +++ b/rocketpool-daemon/common/validator/key-recovery-manager.go @@ -0,0 +1,188 @@ +package validator + +import ( + "fmt" + "github.com/rocket-pool/node-manager-core/beacon" + "github.com/rocket-pool/node-manager-core/utils" + "golang.org/x/exp/maps" + "strings" +) + +type KeyRecoveryManager struct { + manager *ValidatorManager + partialEnabled bool + testOnly bool +} + +func NewKeyRecoveryManager(m *ValidatorManager, partialEnabled bool, testOnly bool) *KeyRecoveryManager { + return &KeyRecoveryManager{ + manager: m, + partialEnabled: partialEnabled, + testOnly: testOnly, + } +} + +func (s *KeyRecoveryManager) RecoverMinipoolKeys() ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error, error) { + publicKeys, err := s.manager.GetMinipools() + if err != nil { + return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err + } + + recoveredCustomPublicKeys, unrecoverableCustomPublicKeys, err := s.checkForAndRecoverCustomKeys(publicKeys) + if err != nil && !s.partialEnabled { + return maps.Keys(recoveredCustomPublicKeys), unrecoverableCustomPublicKeys, err + } + + recoveredConventionalPublicKeys, unrecoveredPublicKeys := s.recoverConventionalKeys(publicKeys) + + var allRecoveredPublicKeys []beacon.ValidatorPubkey + allRecoveredPublicKeys = append(allRecoveredPublicKeys, maps.Keys(recoveredCustomPublicKeys)...) + allRecoveredPublicKeys = append(allRecoveredPublicKeys, recoveredConventionalPublicKeys...) + + for publicKey, err := range unrecoveredPublicKeys { + unrecoverableCustomPublicKeys[publicKey] = err + } + + return allRecoveredPublicKeys, unrecoveredPublicKeys, nil +} + +func (s *KeyRecoveryManager) 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 := s.manager.LoadFiles() + if err != nil { + return recoveredKeys, recoveryFailures, err + } + + if len(keyFiles) == 0 { + return recoveredKeys, recoveryFailures, nil + } + + passwords, err = s.manager.LoadCustomKeyPasswords() + if err != nil { + return recoveredKeys, recoveryFailures, err + } + + for _, file := range keyFiles { + keystore, err := s.manager.ReadCustomKeystore(file) + if err != nil { + if s.partialEnabled { + continue + } + return recoveredKeys, recoveryFailures, err + } + + 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 + if s.partialEnabled { + continue + } + return recoveredKeys, recoveryFailures, err + } + + 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 + if s.partialEnabled { + continue + } + return recoveredKeys, recoveryFailures, err + } + + privateKey, err := s.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 + if s.partialEnabled { + continue + } + return recoveredKeys, recoveryFailures, err + } + + 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 + if s.partialEnabled { + continue + } + return recoveredKeys, recoveryFailures, err + } + + if !s.testOnly { + if err := s.manager.StoreValidatorKey(&privateKey, keystore.Path); err != nil { + recoveryFailures[keystore.Pubkey] = err + if s.partialEnabled { + continue + } + return recoveredKeys, recoveryFailures, err + } + } + recoveredKeys[reconstructedPublicKey] = true + + delete(publicKeys, keystore.Pubkey) + } + + return recoveredKeys, recoveryFailures, nil +} + +func (s *KeyRecoveryManager) recoverConventionalKeys(publicKeys map[beacon.ValidatorPubkey]bool) ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error) { + var 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 := s.manager.GetValidatorKeys(bucketStart, bucketEnd-bucketStart) + if err != nil { + return recoveredPublicKeys, map[beacon.ValidatorPubkey]error{beacon.ValidatorPubkey{}: fmt.Errorf("error getting node's validator keys")} + } + + for _, validatorKey := range keys { + if exists := publicKeys[validatorKey.PublicKey]; exists { + delete(publicKeys, validatorKey.PublicKey) + if !s.testOnly { + if err := s.manager.SaveValidatorKey(validatorKey); err != nil { + unrecoverablePublicKeys[validatorKey.PublicKey] = err + if s.partialEnabled { + continue + } + return recoveredPublicKeys, unrecoverablePublicKeys + } + } + recoveredPublicKeys = append(recoveredPublicKeys, validatorKey.PublicKey) + + } else { + err := fmt.Errorf("keystore for pubkey %s not found in minipool keyset", validatorKey.PublicKey) + unrecoverablePublicKeys[validatorKey.PublicKey] = err + if !s.partialEnabled { + return recoveredPublicKeys, unrecoverablePublicKeys + } + } + } + + if len(publicKeys) == 0 { + // All keys have been recovered. + break + } + + bucketStart = bucketEnd + } + + return recoveredPublicKeys, unrecoverablePublicKeys +} diff --git a/rocketpool-daemon/common/validator/key-recovery-manager/dry-run-key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager/dry-run-key-recovery-manager.go deleted file mode 100644 index fd854232e..000000000 --- a/rocketpool-daemon/common/validator/key-recovery-manager/dry-run-key-recovery-manager.go +++ /dev/null @@ -1,149 +0,0 @@ -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 -} diff --git a/rocketpool-daemon/common/validator/key-recovery-manager/key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager/key-recovery-manager.go deleted file mode 100644 index dab48bdd5..000000000 --- a/rocketpool-daemon/common/validator/key-recovery-manager/key-recovery-manager.go +++ /dev/null @@ -1,14 +0,0 @@ -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 -) diff --git a/rocketpool-daemon/common/validator/key-recovery-manager/partial-key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager/partial-key-recovery-manager.go deleted file mode 100644 index c96a93c9a..000000000 --- a/rocketpool-daemon/common/validator/key-recovery-manager/partial-key-recovery-manager.go +++ /dev/null @@ -1,160 +0,0 @@ -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 PartialRecoveryManager struct { - manager *validator.ValidatorManager -} - -func NewPartialRecoveryManager(m *validator.ValidatorManager) *PartialRecoveryManager { - return &PartialRecoveryManager{ - manager: m, - } -} - -func (p PartialRecoveryManager) RecoverMinipoolKeys() ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error, error) { - status, err := p.manager.GetWalletStatus() - if err != nil { - return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err - } - - rpNode, mpMgr, err := p.manager.InitializeBindings(status) - if err != nil { - return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err - } - - publicKeys, err := p.manager.GetMinipools(rpNode, mpMgr) - if err != nil { - return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err - } - - recoveredCustomPublicKeys, unrecoverableCustomPublicKeys, _ := p.checkForAndRecoverCustomKeys(publicKeys) - recoveredPublicKeys, unrecoverablePublicKeys := p.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 (p PartialRecoveryManager) 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 := p.manager.LoadFiles() - if err != nil { - return recoveredKeys, recoveryFailures, err - } - - if len(keyFiles) > 0 { - passwords, err = p.manager.LoadCustomKeyPasswords() - if err != nil { - return recoveredKeys, recoveryFailures, err - } - - for _, file := range keyFiles { - keystore, err := p.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 := p.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 - } - - if err := p.manager.StoreValidatorKey(&privateKey, keystore.Path); err != nil { - recoveryFailures[reconstructedPublicKey] = fmt.Errorf("error storing private keystore for %s: %w", reconstructedPublicKey.Hex(), err) - } else { - recoveredKeys[reconstructedPublicKey] = true - } - - delete(publicKeys, keystore.Pubkey) - } - } - - return recoveredKeys, recoveryFailures, nil -} - -func (p *PartialRecoveryManager) 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 := p.manager.GetValidatorKeys(bucketStart, bucketEnd-bucketStart) - if err != nil { - continue - } - - for _, validatorKey := range keys { - delete(publicKeys, validatorKey.PublicKey) - if exists := publicKeys[validatorKey.PublicKey]; exists { - if err := p.manager.SaveValidatorKey(validatorKey); err != nil { - unrecoverablePublicKeys[validatorKey.PublicKey] = err - } else { - 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 -} diff --git a/rocketpool-daemon/common/validator/key-recovery-manager/strict-key-recovery-manager.go b/rocketpool-daemon/common/validator/key-recovery-manager/strict-key-recovery-manager.go deleted file mode 100644 index ff5079c70..000000000 --- a/rocketpool-daemon/common/validator/key-recovery-manager/strict-key-recovery-manager.go +++ /dev/null @@ -1,166 +0,0 @@ -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 StrictRecoveryManager struct { - manager *validator.ValidatorManager -} - -func NewStrictRecoveryManager(m *validator.ValidatorManager) *StrictRecoveryManager { - return &StrictRecoveryManager{ - manager: m, - } -} - -func (s *StrictRecoveryManager) RecoverMinipoolKeys() ([]beacon.ValidatorPubkey, map[beacon.ValidatorPubkey]error, error) { - status, err := s.manager.GetWalletStatus() - if err != nil { - return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err - } - - rpNode, mpMgr, err := s.manager.InitializeBindings(status) - if err != nil { - return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err - } - - publicKeys, err := s.manager.GetMinipools(rpNode, mpMgr) - if err != nil { - return []beacon.ValidatorPubkey{}, map[beacon.ValidatorPubkey]error{}, err - } - - recoveredCustomPublicKeys, unrecoverableCustomPublicKeys, err := s.checkForAndRecoverCustomKeys(publicKeys) - if err != nil { - return maps.Keys(recoveredCustomPublicKeys), unrecoverableCustomPublicKeys, err - } - - recoveredPublicKeys, unrecoverablePublicKeys := s.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 (s *StrictRecoveryManager) 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 := s.manager.LoadFiles() - if err != nil { - return recoveredKeys, recoveryFailures, err - } - - if len(keyFiles) > 0 { - passwords, err = s.manager.LoadCustomKeyPasswords() - if err != nil { - return recoveredKeys, recoveryFailures, err - } - - for _, file := range keyFiles { - keystore, err := s.manager.ReadCustomKeystore(file) - if err != nil { - return recoveredKeys, recoveryFailures, err - } - - 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 - return recoveredKeys, recoveryFailures, err - } - - 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 - return recoveredKeys, recoveryFailures, err - } - - privateKey, err := s.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 - return recoveredKeys, recoveryFailures, err - } - - 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 - return recoveredKeys, recoveryFailures, err - } - - if err := s.manager.StoreValidatorKey(&privateKey, keystore.Path); err != nil { - recoveryFailures[keystore.Pubkey] = err - return recoveredKeys, recoveryFailures, err - } - recoveredKeys[reconstructedPublicKey] = true - - delete(publicKeys, keystore.Pubkey) - } - } - - return recoveredKeys, recoveryFailures, nil -} - -func (s *StrictRecoveryManager) 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 := s.manager.GetValidatorKeys(bucketStart, bucketEnd-bucketStart) - if err != nil { - return recoveredPublicKeys, map[beacon.ValidatorPubkey]error{beacon.ValidatorPubkey{}: fmt.Errorf("error getting node's validator keys")} - } - - for _, validatorKey := range keys { - if exists := publicKeys[validatorKey.PublicKey]; exists { - delete(publicKeys, validatorKey.PublicKey) - if err := s.manager.SaveValidatorKey(validatorKey); err != nil { - unrecoverablePublicKeys[validatorKey.PublicKey] = err - return recoveredPublicKeys, unrecoverablePublicKeys - } else { - recoveredPublicKeys = append(recoveredPublicKeys, validatorKey.PublicKey) - } - } else { - err := fmt.Errorf("keystore for pubkey %s not found in minipool keyset", validatorKey.PublicKey) - unrecoverablePublicKeys[validatorKey.PublicKey] = err - return recoveredPublicKeys, unrecoverablePublicKeys - } - } - - if len(publicKeys) == 0 { - // All keys have been recovered. - break - } - - bucketStart = bucketEnd - } - - return recoveredPublicKeys, unrecoverablePublicKeys -} diff --git a/rocketpool-daemon/common/validator/validator-manager.go b/rocketpool-daemon/common/validator/validator-manager.go index 329a200c9..8d10bb394 100644 --- a/rocketpool-daemon/common/validator/validator-manager.go +++ b/rocketpool-daemon/common/validator/validator-manager.go @@ -3,6 +3,8 @@ package validator import ( "bytes" "fmt" + "os" + "github.com/goccy/go-json" batch "github.com/rocket-pool/batch-query" "github.com/rocket-pool/node-manager-core/beacon" @@ -20,7 +22,6 @@ import ( types "github.com/wealdtech/go-eth2-types/v2" eth2ks "github.com/wealdtech/go-eth2-wallet-encryptor-keystorev4" "gopkg.in/yaml.v3" - "os" "path/filepath" ) @@ -48,6 +49,8 @@ type ValidatorManager struct { queryMgr *eth.QueryManager keystoreManager *validator.ValidatorManager nextAccount uint64 + node *node.Node + minipoolManager *minipool.MinipoolManager } func NewValidatorManager(cfg *config.SmartNodeConfig, rp *rocketpool.RocketPool, walletImpl *walletnode.Wallet, queryMgr *eth.QueryManager) (*ValidatorManager, error) { @@ -62,9 +65,12 @@ func NewValidatorManager(cfg *config.SmartNodeConfig, rp *rocketpool.RocketPool, queryMgr: queryMgr, keystoreManager: validatorManager, } + err := mgr.initializeBindings() + if err != nil { + return nil, err + } // Load the next account - var err error mgr.nextAccount, err = loadNextAccount(cfg.GetNextAccountFilePath()) if err != nil { return nil, err @@ -259,70 +265,23 @@ func (m *ValidatorManager) TestRecoverValidatorKey(pubkey beacon.ValidatorPubkey return index + startIndex, nil } -// Get a validator private key by index -func (m *ValidatorManager) getValidatorPrivateKey(index uint64) (*eth2types.BLSPrivateKey, string, error) { - // Get derivation path - derivationPath := fmt.Sprintf(ValidatorKeyPath, index) - - // Get private key - privateKeyBytes, err := m.wallet.GenerateValidatorKey(derivationPath) - if err != nil { - return nil, "", fmt.Errorf("error getting validator %d private key: %w", index, err) - } - privateKey, err := types.BLSPrivateKeyFromBytes(privateKeyBytes) - if err != nil { - return nil, "", fmt.Errorf("error converting validator %d private key: %w", index, err) - } - return privateKey, derivationPath, nil -} - -// Checks if the wallet is ready for validator key processing -func (m *ValidatorManager) checkIfReady() error { - status, err := m.wallet.GetStatus() - if err != nil { - return err - } - return utils.CheckIfWalletReady(status) -} - -func (m *ValidatorManager) GetWalletStatus() (walletcore.WalletStatus, error) { - status, err := m.wallet.GetStatus() - if err != nil { - return status, err - } - if !walletcore.IsWalletReady(status) { - return status, fmt.Errorf("wallet is not ready") - } - return status, nil -} - -func (m *ValidatorManager) InitializeBindings(status walletcore.WalletStatus) (*node.Node, *minipool.MinipoolManager, error) { - address := status.Wallet.WalletAddress - rpNode, err := node.NewNode(m.rp, address) - - if err != nil { - return nil, nil, err - } - - mpMgr, err := minipool.NewMinipoolManager(m.rp) +func (m *ValidatorManager) GetMinipools() (map[beacon.ValidatorPubkey]bool, error) { + err := m.initializeBindings() if err != nil { - return nil, nil, err + return nil, err } - return rpNode, mpMgr, nil -} -func (m *ValidatorManager) GetMinipools(node *node.Node, mpMgr *minipool.MinipoolManager) (map[beacon.ValidatorPubkey]bool, error) { - err := m.queryMgr.Query(nil, nil, node.ValidatingMinipoolCount) + err = m.queryMgr.Query(nil, nil, m.node.ValidatingMinipoolCount) if err != nil { return nil, fmt.Errorf("error getting node's validating minipool count: %w", err) } - addresses, err := node.GetValidatingMinipoolAddresses(node.ValidatingMinipoolCount.Formatted(), nil) + addresses, err := m.node.GetValidatingMinipoolAddresses(m.node.ValidatingMinipoolCount.Formatted(), nil) if err != nil { return nil, fmt.Errorf("error getting node's validating minipool addresses: %w", err) } - mps, err := mpMgr.CreateMinipoolsFromAddresses(addresses, false, nil) + mps, err := m.minipoolManager.CreateMinipoolsFromAddresses(addresses, false, nil) if err != nil { return nil, fmt.Errorf("error creating bindings for node's validating minipools: %w", err) } @@ -420,3 +379,62 @@ func (m *ValidatorManager) LoadFiles() ([]os.DirEntry, error) { } return keyFiles, nil } + +// Get a validator private key by index +func (m *ValidatorManager) getValidatorPrivateKey(index uint64) (*eth2types.BLSPrivateKey, string, error) { + // Get derivation path + derivationPath := fmt.Sprintf(ValidatorKeyPath, index) + + // Get private key + privateKeyBytes, err := m.wallet.GenerateValidatorKey(derivationPath) + if err != nil { + return nil, "", fmt.Errorf("error getting validator %d private key: %w", index, err) + } + privateKey, err := types.BLSPrivateKeyFromBytes(privateKeyBytes) + if err != nil { + return nil, "", fmt.Errorf("error converting validator %d private key: %w", index, err) + } + return privateKey, derivationPath, nil +} + +// Checks if the wallet is ready for validator key processing +func (m *ValidatorManager) checkIfReady() error { + status, err := m.wallet.GetStatus() + if err != nil { + return err + } + return utils.CheckIfWalletReady(status) +} + +func (m *ValidatorManager) initializeBindings() error { + status, err := m.getWalletStatus() + if err != nil { + return err + } + + rpNode, err := node.NewNode(m.rp, status.Wallet.WalletAddress) + if err != nil { + return err + } + + mpMgr, err := minipool.NewMinipoolManager(m.rp) + if err != nil { + return err + } + + m.node = rpNode + m.minipoolManager = mpMgr + + return nil +} + +func (m *ValidatorManager) getWalletStatus() (*walletcore.WalletStatus, error) { + status, err := m.wallet.GetStatus() + if err != nil { + return &status, err + } + if !walletcore.IsWalletReady(status) { + return &status, fmt.Errorf("wallet is not ready") + } + return &status, nil +}