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 +}