draft: proof of concept for plan-before-write validator-key recovery (#407, #572)

rocketpool/api/wallet/recover.go

@@ -276,8 +276,9 @@ func recoverNodeKeys(c *cli.Command, rp *rocketpool.RocketPool, bc beacon.Client
 		}
 	}
 
+	pubkeys = filteredPubkeys
 	pubkeyMap := map[types.ValidatorPubkey]bool{}
-	for _, pubkey := range pubkeys {
+	for _, pubkey := range filteredPubkeys {
 		pubkeyMap[pubkey] = true
 	}
 

shared/services/recovery/README.md

@@ -0,0 +1,55 @@
+# Smartnode Safe Recovery Planner
+
+This package implements the **two-phase (Plan-Before-Write) recovery architecture** and per-key diagnosis for Rocket Pool Smartnode validator keys, addressing issues **#407**, **#440**, and **#572**.
+
+---
+
+## The Problem
+
+Previously, Smartnode's key recovery routines (`wallet recover`, `wallet rebuild`, and `wallet test-recovery`) immediately persisted keys to disk as soon as each individual key was discovered. In mixed-key deployments (a combination of mnemonic-derived keys, imported custom EIP-2335 keystores, legacy minipools, and megapool validators), if recovery succeeded for several keys but subsequently failed on a later key (e.g., missing custom keystore, corrupt file, invalid password, or derivation limit exceeded), the node was left in an inconsistent, partially-recovered state.
+
+Additionally, inactive/exited validators on the Beacon chain were filtered into an active list, but the target discovery map was populated from the original unfiltered list, causing exited validators to trigger false recovery failures (Issue #407).
+
+---
+
+## Architectural Principles
+
+### 1. Two-Phase Execution (Plan-Before-Write)
+
+Recovery is strictly separated into two independent phases:
+
+1. **Plan Phase (`PlanValidatorRecovery`)**:
+   - Pure, read-only discovery and diagnosis.
+   - Evaluates all expected minipool and megapool public keys against Beacon chain validator states, installed client keystores, imported custom keys (`custom_keys/`), and mnemonic derivation limits.
+   - Assigns a deterministic `DiscoveryOutcome` to every single public key.
+   - **Zero file mutations, zero keystore modifications.**
+
+2. **Commit Phase (`CommitPlan`)**:
+   - Evaluates the safety invariants of the plan before touching disk.
+   - By default, **any** unresolved or invalid in-scope validator key causes the commit phase to abort immediately before any keystore is touched (`ErrUnresolvedKeysRefusingCommit`).
+   - Only when all in-scope keys are resolved (or when the operator explicitly passes `--allow-partial-recover`), the validated keys are committed to disk via `KeyWriter`.
+
+### 2. Deterministic Per-Key Outcomes
+
+Each public key is assigned a deterministic diagnosis:
+- **`mnemonic_valid`**: Successfully derived from the node mnemonic and verified against the public key.
+- **`custom_valid`**: Loaded from an imported EIP-2335 keystore and decrypted with the provided password.
+- **`installed_valid`**: Key material already exists and is valid in the node's validator client keystores.
+- **`unresolved`**: Not found in custom keys or within the derivation search window (`bucketLimit`).
+- **`invalid_material`**: Corrupt keystore, incorrect password, or pubkey/private key mismatch.
+- **`excluded_inactive`**: Exited or inactive on the Beacon chain; safely excluded from recovery.
+
+Commit status is tracked explicitly:
+- `written`: Successfully written to keystores.
+- `skipped_already_installed`: Key was already installed; avoided redundant overwrite.
+- `skipped_excluded`: Inactive validator; no keys written.
+- `skipped_partial`: Unresolved key omitted during an operator-approved partial recovery.
+- `not_attempted`: Write phase was blocked because plan was invalid or unapproved.
+
+### 3. Partial Recovery Behavior (`--allow-partial-recover`)
+
+When an operator cannot locate a lost custom key or wants to bring online only the validators that could be found, passing `--allow-partial-recover` allows Smartnode to commit all resolved keys while recording which keys were skipped (`skipped_partial`). This prevents operators from being stuck when one obsolete or migrated key cannot be derived.
+
+### 4. Exclusion of Cross-Client Rollback
+
+Smartnode interfaces with diverse external Validator Clients (Lighthouse, Nimbus, Prysm, Teku, Lodestar), each with its own local keystore formats and disk layouts. This architecture provides a **pre-mutation safety guarantee** (ensuring all keys are validated in memory before any writes begin). It does not attempt transactional cross-client atomic rollback in the event of hardware or OS-level disk failures mid-write. Operators can safely re-run recovery idempotently at any time.

shared/services/recovery/committer.go

@@ -0,0 +1,111 @@
+package recovery
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/rocket-pool/smartnode/shared/services/wallet"
+)
+
+// ErrUnresolvedKeysRefusingCommit is returned when default mode encounters unresolved or invalid keys.
+var ErrUnresolvedKeysRefusingCommit = errors.New("refusing to commit recovery: plan contains unresolved or invalid validator keys")
+
+// CommitPlan applies a pre-calculated RecoveryPlan using the provided KeyWriter.
+// By default, if any in-scope key is unresolved or invalid, zero writes occur and an error is returned.
+// When allowPartial is true, resolved keys are committed while unresolved ones are skipped and recorded.
+func CommitPlan(plan *RecoveryPlan, writer KeyWriter, allowPartial bool) error {
+	if err := validateCommitPreconditions(plan, writer, allowPartial); err != nil {
+		return err
+	}
+
+	for i := range plan.Entries {
+		if err := commitEntry(&plan.Entries[i], writer); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// validateCommitPreconditions verifies parameters and enforces the fail-fast safety invariant.
+func validateCommitPreconditions(plan *RecoveryPlan, writer KeyWriter, allowPartial bool) error {
+	if plan == nil {
+		return errors.New("nil recovery plan provided")
+	}
+	if writer == nil {
+		return errors.New("nil key writer provided")
+	}
+
+	if !allowPartial && plan.HasUnresolvedOrInvalid() {
+		return fmt.Errorf("%w: %d unresolved, %d invalid (run with --allow-partial-recover to write only resolved keys)",
+			ErrUnresolvedKeysRefusingCommit, plan.TotalUnresolved, plan.TotalInvalid)
+	}
+
+	return nil
+}
+
+// commitEntry routes key persistence or status updates for a single plan entry.
+func commitEntry(entry *ValidatorPlanEntry, writer KeyWriter) error {
+	switch entry.DiscoveryOutcome {
+	case OutcomeExcludedInactive:
+		entry.CommitOutcome = CommitSkippedExcluded
+		return nil
+
+	case OutcomeInstalledValid:
+		entry.CommitOutcome = CommitSkippedAlreadyInstalled
+		return nil
+
+	case OutcomeUnresolved, OutcomeInvalidMaterial:
+		entry.CommitOutcome = CommitSkippedPartial
+		return nil
+
+	case OutcomeMnemonicValid:
+		return commitMnemonicEntry(entry, writer)
+
+	case OutcomeCustomValid:
+		return commitCustomEntry(entry, writer)
+
+	default:
+		entry.CommitOutcome = CommitNotAttempted
+		return nil
+	}
+}
+
+// commitMnemonicEntry validates and persists a mnemonic-derived validator key.
+func commitMnemonicEntry(entry *ValidatorPlanEntry, writer KeyWriter) error {
+	if entry.WalletIndex == nil || entry.PrivateKey == nil {
+		entry.CommitOutcome = CommitFailed
+		return fmt.Errorf("incomplete key material for mnemonic-derived pubkey %s", entry.Pubkey.Hex())
+	}
+
+	vk := wallet.ValidatorKey{
+		PublicKey:      entry.Pubkey,
+		PrivateKey:     entry.PrivateKey,
+		DerivationPath: entry.DerivationPath,
+		WalletIndex:    *entry.WalletIndex,
+	}
+
+	if err := writer.SaveValidatorKey(vk); err != nil {
+		entry.CommitOutcome = CommitFailed
+		return fmt.Errorf("error writing mnemonic validator key %s: %w", entry.Pubkey.Hex(), err)
+	}
+
+	entry.CommitOutcome = CommitWritten
+	return nil
+}
+
+// commitCustomEntry validates and persists an imported custom validator keystore.
+func commitCustomEntry(entry *ValidatorPlanEntry, writer KeyWriter) error {
+	if entry.PrivateKey == nil {
+		entry.CommitOutcome = CommitFailed
+		return fmt.Errorf("incomplete key material for custom pubkey %s", entry.Pubkey.Hex())
+	}
+
+	if err := writer.StoreValidatorKey(entry.PrivateKey, entry.DerivationPath); err != nil {
+		entry.CommitOutcome = CommitFailed
+		return fmt.Errorf("error storing custom validator key %s: %w", entry.Pubkey.Hex(), err)
+	}
+
+	entry.CommitOutcome = CommitWritten
+	return nil
+}

shared/services/recovery/mock_writer_test.go

@@ -0,0 +1,157 @@
+package recovery
+
+import (
+	"sync"
+
+	eth2types "github.com/wealdtech/go-eth2-types/v2"
+
+	"github.com/rocket-pool/smartnode/bindings/types"
+	"github.com/rocket-pool/smartnode/shared/services/wallet"
+)
+
+// StoreCall records a call to StoreValidatorKey.
+type StoreCall struct {
+	Key  *eth2types.BLSPrivateKey
+	Path string
+}
+
+// MockKeyWriter is a thread-safe spy implementation of KeyWriter that records all write attempts.
+type MockKeyWriter struct {
+	mu                     sync.Mutex
+	saveValidatorKeyCalls  []wallet.ValidatorKey
+	storeValidatorKeyCalls []StoreCall
+
+	SaveError  error
+	StoreError error
+}
+
+func NewMockKeyWriter() *MockKeyWriter {
+	return &MockKeyWriter{
+		saveValidatorKeyCalls:  make([]wallet.ValidatorKey, 0),
+		storeValidatorKeyCalls: make([]StoreCall, 0),
+	}
+}
+
+func (m *MockKeyWriter) SaveValidatorKey(key wallet.ValidatorKey) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if m.SaveError != nil {
+		return m.SaveError
+	}
+	m.saveValidatorKeyCalls = append(m.saveValidatorKeyCalls, key)
+	return nil
+}
+
+func (m *MockKeyWriter) StoreValidatorKey(key *eth2types.BLSPrivateKey, path string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if m.StoreError != nil {
+		return m.StoreError
+	}
+	m.storeValidatorKeyCalls = append(m.storeValidatorKeyCalls, StoreCall{Key: key, Path: path})
+	return nil
+}
+
+func (m *MockKeyWriter) TotalWrites() int {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return len(m.saveValidatorKeyCalls) + len(m.storeValidatorKeyCalls)
+}
+
+func (m *MockKeyWriter) SaveCalls() []wallet.ValidatorKey {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	copied := make([]wallet.ValidatorKey, len(m.saveValidatorKeyCalls))
+	copy(copied, m.saveValidatorKeyCalls)
+	return copied
+}
+
+func (m *MockKeyWriter) StoreCalls() []StoreCall {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	copied := make([]StoreCall, len(m.storeValidatorKeyCalls))
+	copy(copied, m.storeValidatorKeyCalls)
+	return copied
+}
+
+func (m *MockKeyWriter) Reset() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.saveValidatorKeyCalls = m.saveValidatorKeyCalls[:0]
+	m.storeValidatorKeyCalls = m.storeValidatorKeyCalls[:0]
+	m.SaveError = nil
+	m.StoreError = nil
+}
+
+// MockKeyDeriver provides in-memory mock validator keys for derivation testing.
+type MockKeyDeriver struct {
+	mu          sync.RWMutex
+	keysByIndex map[uint]wallet.ValidatorKey
+}
+
+func NewMockKeyDeriver() *MockKeyDeriver {
+	return &MockKeyDeriver{
+		keysByIndex: make(map[uint]wallet.ValidatorKey),
+	}
+}
+
+func (m *MockKeyDeriver) AddKey(index uint, key wallet.ValidatorKey) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.keysByIndex[index] = key
+}
+
+func (m *MockKeyDeriver) GetValidatorKeys(startIndex uint, length uint) ([]wallet.ValidatorKey, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	res := make([]wallet.ValidatorKey, 0, length)
+	for i := startIndex; i < startIndex+length; i++ {
+		if k, ok := m.keysByIndex[i]; ok {
+			res = append(res, k)
+		}
+	}
+	return res, nil
+}
+
+// MockCustomKeyProvider provides custom keys for testing.
+type MockCustomKeyProvider struct {
+	mu   sync.RWMutex
+	keys []CustomKey
+	err  error
+}
+
+func NewMockCustomKeyProvider(keys []CustomKey, err error) *MockCustomKeyProvider {
+	return &MockCustomKeyProvider{keys: keys, err: err}
+}
+
+func (m *MockCustomKeyProvider) GetCustomKeys() ([]CustomKey, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	if m.err != nil {
+		return nil, m.err
+	}
+	return m.keys, nil
+}
+
+// MockInstalledKeyChecker checks installed keys for testing.
+type MockInstalledKeyChecker struct {
+	mu        sync.RWMutex
+	installed map[types.ValidatorPubkey]bool
+}
+
+func NewMockInstalledKeyChecker(installedPubkeys ...types.ValidatorPubkey) *MockInstalledKeyChecker {
+	m := &MockInstalledKeyChecker{installed: make(map[types.ValidatorPubkey]bool, len(installedPubkeys))}
+	for _, pk := range installedPubkeys {
+		m.installed[pk] = true
+	}
+	return m
+}
+
+func (m *MockInstalledKeyChecker) IsKeyInstalled(pubkey types.ValidatorPubkey) (bool, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.installed[pubkey], nil
+}