staticaddr: add restore and reconciliation hooks

Author: hieblmi

staticaddr/address/manager.go

@@ -3,7 +3,9 @@ package address
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
+	"strings"
 	"sync"
 	"sync/atomic"
 
@@ -21,6 +23,12 @@ import (
 	"github.com/lightningnetwork/lnd/lnwallet"
 )
 
+var (
+	// ErrNoStaticAddress is returned when no static address parameters are
+	// present in the store.
+	ErrNoStaticAddress = errors.New("no static address parameters found")
+)
+
 // ManagerConfig holds the configuration for the address manager.
 type ManagerConfig struct {
 	// AddressClient is the client that communicates with the loop server
@@ -199,27 +207,143 @@ func (m *Manager) NewAddress(ctx context.Context) (*btcutil.AddressTaproot,
 		return nil, 0, err
 	}
 
+	_, err = m.importAddressTapscript(ctx, staticAddress)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	address, err := m.GetTaprootAddress(
+		clientPubKey.PubKey, serverPubKey, int64(serverParams.Expiry),
+	)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	return address, int64(serverParams.Expiry), nil
+}
+
+// RestoreAddress recreates a static address record locally and makes sure the
+// corresponding tapscript is imported into lnd. If the same address already
+// exists locally, the call is idempotent.
+func (m *Manager) RestoreAddress(ctx context.Context,
+	addrParams *Parameters) (*btcutil.AddressTaproot, bool, error) {
+
+	if addrParams == nil {
+		return nil, false, fmt.Errorf("missing static address parameters")
+	}
+
+	staticAddress, err := script.NewStaticAddress(
+		input.MuSig2Version100RC2, int64(addrParams.Expiry),
+		addrParams.ClientPubkey, addrParams.ServerPubkey,
+	)
+	if err != nil {
+		return nil, false, err
+	}
+
+	pkScript, err := staticAddress.StaticAddressScript()
+	if err != nil {
+		return nil, false, err
+	}
+
+	if len(addrParams.PkScript) != 0 &&
+		!bytes.Equal(addrParams.PkScript, pkScript) {
+
+		return nil, false, fmt.Errorf("static address pk script mismatch")
+	}
+
+	addrParams.PkScript = pkScript
+	if addrParams.InitiationHeight <= 0 {
+		addrParams.InitiationHeight = m.currentHeight.Load()
+	}
+
+	m.Lock()
+	existing, err := m.cfg.Store.GetAllStaticAddresses(ctx)
+	if err != nil {
+		m.Unlock()
+
+		return nil, false, err
+	}
+
+	changed := false
+	switch {
+	case len(existing) == 0:
+		err = m.cfg.Store.CreateStaticAddress(ctx, addrParams)
+		if err != nil {
+			m.Unlock()
+
+			return nil, false, err
+		}
+		changed = true
+
+	case len(existing) > 1:
+		m.Unlock()
+
+		return nil, false, fmt.Errorf("more than one static address found")
+
+	case !sameAddressParameters(existing[0], addrParams):
+		m.Unlock()
+
+		return nil, false, fmt.Errorf("existing static address differs from " +
+			"backup")
+	}
+	m.Unlock()
+
+	imported, err := m.importAddressTapscript(ctx, staticAddress)
+	if err != nil {
+		return nil, false, err
+	}
+
+	changed = changed || imported
+
+	addr, err := m.GetTaprootAddress(
+		addrParams.ClientPubkey, addrParams.ServerPubkey,
+		int64(addrParams.Expiry),
+	)
+	if err != nil {
+		return nil, false, err
+	}
+
+	return addr, changed, nil
+}
+
+func (m *Manager) importAddressTapscript(ctx context.Context,
+	staticAddress *script.StaticAddress) (bool, error) {
+
 	// Import the static address tapscript into our lnd wallet, so we can
 	// track unspent outputs of it.
 	tapScript := input.TapscriptFullTree(
 		staticAddress.InternalPubKey, *staticAddress.TimeoutLeaf,
 	)
 	addr, err := m.cfg.WalletKit.ImportTaprootScript(ctx, tapScript)
 	if err != nil {
-		return nil, 0, err
+		// Restoring into an lnd instance that already imported the script is
+		// expected. Treat the duplicate import as success.
+		if strings.Contains(err.Error(), "already exists") {
+			log.Infof("Static address tapscript already imported")
+			return false, nil
+		}
+
+		return false, err
 	}
 
 	log.Infof("Imported static address taproot script to lnd wallet: %v",
 		addr)
 
-	address, err := m.GetTaprootAddress(
-		clientPubKey.PubKey, serverPubKey, int64(serverParams.Expiry),
-	)
-	if err != nil {
-		return nil, 0, err
+	return true, nil
+}
+
+func sameAddressParameters(a, b *Parameters) bool {
+	if a == nil || b == nil {
+		return false
 	}
 
-	return address, int64(serverParams.Expiry), nil
+	return a.ClientPubkey.IsEqual(b.ClientPubkey) &&
+		a.ServerPubkey.IsEqual(b.ServerPubkey) &&
+		a.Expiry == b.Expiry &&
+		bytes.Equal(a.PkScript, b.PkScript) &&
+		a.KeyLocator == b.KeyLocator &&
+		a.ProtocolVersion == b.ProtocolVersion &&
+		a.InitiationHeight == b.InitiationHeight
 }
 
 // GetTaprootAddress returns a taproot address for the given client and server
@@ -297,7 +421,7 @@ func (m *Manager) GetStaticAddressParameters(ctx context.Context) (*Parameters,
 	}
 
 	if len(params) == 0 {
-		return nil, fmt.Errorf("no static address parameters found")
+		return nil, ErrNoStaticAddress
 	}
 
 	return params[0], nil

staticaddr/address/manager_test.go

@@ -128,6 +128,110 @@ func TestManager(t *testing.T) {
 	require.EqualValues(t, defaultExpiry, expiry)
 }
 
+// TestRestoreAddress verifies that restoring an address recreates the same
+// static address locally without requiring a server call.
+func TestRestoreAddress(t *testing.T) {
+	ctxb := t.Context()
+
+	testContext := NewAddressManagerTestContext(t)
+
+	keyDesc, err := testContext.mockLnd.WalletKit.DeriveKey(
+		ctxb, &keychain.KeyLocator{
+			Family: keychain.KeyFamily(swap.StaticAddressKeyFamily),
+			Index:  7,
+		},
+	)
+	require.NoError(t, err)
+
+	staticAddress, err := script.NewStaticAddress(
+		input.MuSig2Version100RC2, int64(defaultExpiry),
+		keyDesc.PubKey, defaultServerPubkey,
+	)
+	require.NoError(t, err)
+
+	pkScript, err := staticAddress.StaticAddressScript()
+	require.NoError(t, err)
+
+	addressParams := &Parameters{
+		ClientPubkey:     keyDesc.PubKey,
+		ServerPubkey:     defaultServerPubkey,
+		Expiry:           defaultExpiry,
+		PkScript:         pkScript,
+		KeyLocator:       keyDesc.KeyLocator,
+		ProtocolVersion:  0,
+		InitiationHeight: 123,
+	}
+
+	taprootAddress, restored, err := testContext.manager.RestoreAddress(
+		ctxb, addressParams,
+	)
+	require.NoError(t, err)
+	require.True(t, restored)
+
+	expectedAddress, err := btcutil.NewAddressTaproot(
+		schnorr.SerializePubKey(staticAddress.TaprootKey),
+		testContext.manager.cfg.ChainParams,
+	)
+	require.NoError(t, err)
+	require.Equal(t, expectedAddress.String(), taprootAddress.String())
+
+	storedParams, err := testContext.manager.GetStaticAddressParameters(ctxb)
+	require.NoError(t, err)
+	require.True(t, sameAddressParameters(storedParams, addressParams))
+
+	taprootAddress, restored, err = testContext.manager.RestoreAddress(
+		ctxb, addressParams,
+	)
+	require.NoError(t, err)
+	require.False(t, restored)
+	require.Equal(t, expectedAddress.String(), taprootAddress.String())
+}
+
+// TestRestoreAddressRejectsDifferentInitiationHeight verifies that a restore
+// request with the same address material but a different initiation height is
+// rejected instead of being treated as idempotent.
+func TestRestoreAddressRejectsDifferentInitiationHeight(t *testing.T) {
+	ctxb := t.Context()
+
+	testContext := NewAddressManagerTestContext(t)
+
+	keyDesc, err := testContext.mockLnd.WalletKit.DeriveKey(
+		ctxb, &keychain.KeyLocator{
+			Family: keychain.KeyFamily(swap.StaticAddressKeyFamily),
+			Index:  7,
+		},
+	)
+	require.NoError(t, err)
+
+	staticAddress, err := script.NewStaticAddress(
+		input.MuSig2Version100RC2, int64(defaultExpiry),
+		keyDesc.PubKey, defaultServerPubkey,
+	)
+	require.NoError(t, err)
+
+	pkScript, err := staticAddress.StaticAddressScript()
+	require.NoError(t, err)
+
+	addressParams := &Parameters{
+		ClientPubkey:     keyDesc.PubKey,
+		ServerPubkey:     defaultServerPubkey,
+		Expiry:           defaultExpiry,
+		PkScript:         pkScript,
+		KeyLocator:       keyDesc.KeyLocator,
+		ProtocolVersion:  0,
+		InitiationHeight: 123,
+	}
+
+	_, _, err = testContext.manager.RestoreAddress(ctxb, addressParams)
+	require.NoError(t, err)
+
+	differentHeight := *addressParams
+	differentHeight.InitiationHeight = 456
+
+	_, _, err = testContext.manager.RestoreAddress(ctxb, &differentHeight)
+	require.ErrorContains(t, err, "existing static address differs from backup")
+}
+
 // GenerateExpectedTaprootAddress generates the expected taproot address that
 // the predefined parameters are supposed to generate.
 func GenerateExpectedTaprootAddress(t *ManagerTestContext) (

staticaddr/deposit/manager.go

@@ -64,6 +64,10 @@ type Manager struct {
 	// mu guards access to the activeDeposits map.
 	mu sync.Mutex
 
+	// reconcileMu serializes deposit recovery and reconciliation so restore
+	// requests can't race the background polling loop.
+	reconcileMu sync.Mutex
+
 	// activeDeposits contains all the active static address outputs.
 	activeDeposits map[wire.OutPoint]*FSM
 
@@ -108,7 +112,7 @@ func (m *Manager) Run(ctx context.Context, initChan chan struct{}) error {
 
 	// Reconcile immediately on startup so deposits are available
 	// before the first ticker fires.
-	err = m.reconcileDeposits(ctx)
+	_, err = m.ReconcileDeposits(ctx)
 	if err != nil {
 		log.Errorf("unable to reconcile deposits: %v", err)
 	}
@@ -162,6 +166,9 @@ func (m *Manager) Run(ctx context.Context, initChan chan struct{}) error {
 // recoverDeposits recovers static address parameters, previous deposits and
 // state machines from the database and starts the deposit notifier.
 func (m *Manager) recoverDeposits(ctx context.Context) error {
+	m.reconcileMu.Lock()
+	defer m.reconcileMu.Unlock()
+
 	log.Infof("Recovering static address parameters and deposits...")
 
 	// Recover deposits.
@@ -218,7 +225,7 @@ func (m *Manager) pollDeposits(ctx context.Context) {
 		for {
 			select {
 			case <-ticker.C:
-				err := m.reconcileDeposits(ctx)
+				_, err := m.ReconcileDeposits(ctx)
 				if err != nil {
 					log.Errorf("unable to reconcile "+
 						"deposits: %v", err)
@@ -235,38 +242,47 @@ func (m *Manager) pollDeposits(ctx context.Context) {
 // wallet and matches it against the deposits in our memory that we've seen so
 // far. It picks the newly identified deposits and starts a state machine per
 // deposit to track its progress.
-func (m *Manager) reconcileDeposits(ctx context.Context) error {
+func (m *Manager) reconcileDeposits(ctx context.Context) (int, error) {
 	log.Tracef("Reconciling new deposits...")
 
 	utxos, err := m.cfg.AddressManager.ListUnspent(
 		ctx, MinConfs, MaxConfs,
 	)
 	if err != nil {
-		return fmt.Errorf("unable to list new deposits: %w", err)
+		return 0, fmt.Errorf("unable to list new deposits: %w", err)
 	}
 
 	newDeposits := m.filterNewDeposits(utxos)
 	if len(newDeposits) == 0 {
 		log.Tracef("No new deposits...")
-		return nil
+		return 0, nil
 	}
 
 	for _, utxo := range newDeposits {
 		deposit, err := m.createNewDeposit(ctx, utxo)
 		if err != nil {
-			return fmt.Errorf("unable to retain new deposit: %w",
+			return 0, fmt.Errorf("unable to retain new deposit: %w",
 				err)
 		}
 
 		log.Debugf("Received deposit: %v", deposit)
 		err = m.startDepositFsm(ctx, deposit)
 		if err != nil {
-			return fmt.Errorf("unable to start new deposit FSM: %w",
+			return 0, fmt.Errorf("unable to start new deposit FSM: %w",
 				err)
 		}
 	}
 
-	return nil
+	return len(newDeposits), nil
+}
+
+// ReconcileDeposits triggers a best-effort reconciliation pass and returns the
+// number of newly discovered deposits.
+func (m *Manager) ReconcileDeposits(ctx context.Context) (int, error) {
+	m.reconcileMu.Lock()
+	defer m.reconcileMu.Unlock()
+
+	return m.reconcileDeposits(ctx)
 }
 
 // createNewDeposit transforms the wallet utxo into a deposit struct and stores