lnwallet+peer: move internalKeyForAddr to lnwallet package

This way we can re-use it. We also make it slightly more generalized.

lnwallet/interface.go

@@ -11,6 +11,7 @@ import (
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/btcutil/hdkeychain"
 	"github.com/btcsuite/btcd/btcutil/psbt"
+	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
@@ -599,6 +600,68 @@ type MessageSigner interface {
 		doubleHash bool) (*ecdsa.Signature, error)
 }
 
+// AddrWithKey wraps a normal addr, but also includes the internal key for the
+// delivery addr if known.
+type AddrWithKey struct {
+	lnwire.DeliveryAddress
+
+	InternalKey fn.Option[btcec.PublicKey]
+
+	// TODO(roasbeef): consolidate w/ instance in chan closer
+}
+
+// InternalKeyForAddr returns the internal key associated with a taproot
+// address.
+func InternalKeyForAddr(wallet WalletController, netParams *chaincfg.Params,
+	deliveryScript []byte) (fn.Option[keychain.KeyDescriptor], error) {
+
+	none := fn.None[keychain.KeyDescriptor]()
+
+	pkScript, err := txscript.ParsePkScript(deliveryScript)
+	if err != nil {
+		return none, err
+	}
+	addr, err := pkScript.Address(netParams)
+	if err != nil {
+		return none, err
+	}
+
+	// If it's not a taproot address, we don't require to know the internal
+	// key in the first place. So we don't return an error here, but also no
+	// internal key.
+	_, isTaproot := addr.(*btcutil.AddressTaproot)
+	if !isTaproot {
+		return none, nil
+	}
+
+	walletAddr, err := wallet.AddressInfo(addr)
+	if err != nil {
+		return none, err
+	}
+
+	// No wallet addr. No error, but we'll return an nil error value here,
+	// as callers can use the .Option() method to get an option value.
+	if walletAddr == nil {
+		return none, nil
+	}
+
+	pubKeyAddr, ok := walletAddr.(waddrmgr.ManagedPubKeyAddress)
+	if !ok {
+		return none, fmt.Errorf("expected pubkey addr, got %T",
+			pubKeyAddr)
+	}
+
+	_, derivationPath, _ := pubKeyAddr.DerivationInfo()
+
+	return fn.Some[keychain.KeyDescriptor](keychain.KeyDescriptor{
+		KeyLocator: keychain.KeyLocator{
+			Family: keychain.KeyFamily(derivationPath.Account),
+			Index:  derivationPath.Index,
+		},
+		PubKey: pubKeyAddr.PubKey(),
+	}), nil
+}
+
 // WalletDriver represents a "driver" for a particular concrete
 // WalletController implementation. A driver is identified by a globally unique
 // string identifier along with a 'New()' method which is responsible for

peer/brontide.go

@@ -13,13 +13,11 @@ import (
 	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2"
-	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/connmgr"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btclog"
-	"github.com/btcsuite/btcwallet/waddrmgr"
 	"github.com/davecgh/go-spew/spew"
 	"github.com/lightningnetwork/lnd/buffer"
 	"github.com/lightningnetwork/lnd/build"
@@ -37,6 +35,7 @@ import (
 	"github.com/lightningnetwork/lnd/htlcswitch/hop"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/invoices"
+	"github.com/lightningnetwork/lnd/keychain"
 	"github.com/lightningnetwork/lnd/lnpeer"
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnutils"
@@ -902,71 +901,31 @@ func (p *Brontide) QuitSignal() <-chan struct{} {
 	return p.quit
 }
 
-// internalKeyForAddr returns the internal key associated with a taproot
-// address.
-func internalKeyForAddr(wallet *lnwallet.LightningWallet,
-	deliveryScript []byte) (fn.Option[btcec.PublicKey], error) {
-
-	none := fn.None[btcec.PublicKey]()
-
-	pkScript, err := txscript.ParsePkScript(deliveryScript)
-	if err != nil {
-		return none, err
-	}
-	addr, err := pkScript.Address(&wallet.Cfg.NetParams)
-	if err != nil {
-		return none, err
-	}
-
-	// If it's not a taproot address, we don't require to know the internal
-	// key in the first place. So we don't return an error here, but also no
-	// internal key.
-	_, isTaproot := addr.(*btcutil.AddressTaproot)
-	if !isTaproot {
-		return none, nil
-	}
-
-	walletAddr, err := wallet.AddressInfo(addr)
-	if err != nil {
-		return none, err
-	}
-
-	// If the address isn't known to the wallet, we can't determine the
-	// internal key.
-	if walletAddr == nil {
-		return none, nil
-	}
-
-	pubKeyAddr, ok := walletAddr.(waddrmgr.ManagedPubKeyAddress)
-	if !ok {
-		return none, fmt.Errorf("expected pubkey addr, got %T",
-			pubKeyAddr)
-	}
-
-	return fn.Some(*pubKeyAddr.PubKey()), nil
-}
-
 // addrWithInternalKey takes a delivery script, then attempts to supplement it
 // with information related to the internal key for the addr, but only if it's
 // a taproot addr.
 func (p *Brontide) addrWithInternalKey(
-	deliveryScript []byte) fn.Result[chancloser.DeliveryAddrWithKey] {
+	deliveryScript []byte) (*chancloser.DeliveryAddrWithKey, error) {
 
-	// TODO(roasbeef): not compatible with external shutdown addr?
 	// Currently, custom channels cannot be created with external upfront
 	// shutdown addresses, so this shouldn't be an issue. We only require
 	// the internal key for taproot addresses to be able to provide a non
 	// inclusion proof of any scripts.
-
-	internalKey, err := internalKeyForAddr(p.cfg.Wallet, deliveryScript)
+	internalKeyDesc, err := lnwallet.InternalKeyForAddr(
+		p.cfg.Wallet, &p.cfg.Wallet.Cfg.NetParams, deliveryScript,
+	)
 	if err != nil {
-		return fn.Err[chancloser.DeliveryAddrWithKey](err)
+		return nil, fmt.Errorf("unable to fetch internal key: %w", err)
 	}
 
-	return fn.Ok(chancloser.DeliveryAddrWithKey{
+	return &chancloser.DeliveryAddrWithKey{
 		DeliveryAddress: deliveryScript,
-		InternalKey:     internalKey,
-	})
+		InternalKey: fn.MapOption(
+			func(desc keychain.KeyDescriptor) btcec.PublicKey {
+				return *desc.PubKey
+			},
+		)(internalKeyDesc),
+	}, nil
 }
 
 // loadActiveChannels creates indexes within the peer for tracking all active
@@ -1040,6 +999,8 @@ func (p *Brontide) loadActiveChannels(chans []*channeldb.OpenChannel) (
 			}
 		}
 
+		// TODO(roasbeef): also make aux resolver here
+
 		var chanOpts []lnwallet.ChannelOpt
 		p.cfg.AuxLeafStore.WhenSome(func(s lnwallet.AuxLeafStore) {
 			chanOpts = append(chanOpts, lnwallet.WithLeafStore(s))
@@ -1211,7 +1172,7 @@ func (p *Brontide) loadActiveChannels(chans []*channeldb.OpenChannel) (
 
 			addr, err := p.addrWithInternalKey(
 				info.DeliveryScript.Val,
-			).Unpack()
+			)
 			if err != nil {
 				shutdownInfoErr = fmt.Errorf("unable to make "+
 					"delivery addr: %w", err)
@@ -2977,7 +2938,7 @@ func (p *Brontide) fetchActiveChanCloser(chanID lnwire.ChannelID) (
 		return nil, fmt.Errorf("unable to estimate fee")
 	}
 
-	addr, err := p.addrWithInternalKey(deliveryScript).Unpack()
+	addr, err := p.addrWithInternalKey(deliveryScript)
 	if err != nil {
 		return nil, fmt.Errorf("unable to parse addr: %w", err)
 	}
@@ -3224,7 +3185,7 @@ func (p *Brontide) restartCoopClose(lnChan *lnwallet.LightningChannel) (
 		closingParty = lntypes.Local
 	}
 
-	addr, err := p.addrWithInternalKey(deliveryScript).Unpack()
+	addr, err := p.addrWithInternalKey(deliveryScript)
 	if err != nil {
 		return nil, fmt.Errorf("unable to parse addr: %w", err)
 	}
@@ -3256,7 +3217,7 @@ func (p *Brontide) restartCoopClose(lnChan *lnwallet.LightningChannel) (
 // createChanCloser constructs a ChanCloser from the passed parameters and is
 // used to de-duplicate code.
 func (p *Brontide) createChanCloser(channel *lnwallet.LightningChannel,
-	deliveryScript chancloser.DeliveryAddrWithKey,
+	deliveryScript *chancloser.DeliveryAddrWithKey,
 	fee chainfee.SatPerKWeight, req *htlcswitch.ChanClose,
 	closer lntypes.ChannelParty) (*chancloser.ChanCloser, error) {
 
@@ -3291,7 +3252,7 @@ func (p *Brontide) createChanCloser(channel *lnwallet.LightningChannel,
 			ChainParams: &p.cfg.Wallet.Cfg.NetParams,
 			Quit:        p.quit,
 		},
-		deliveryScript,
+		*deliveryScript,
 		fee,
 		uint32(startingHeight),
 		req,
@@ -3350,7 +3311,7 @@ func (p *Brontide) handleLocalCloseReq(req *htlcswitch.ChanClose) {
 				return
 			}
 		}
-		addr, err := p.addrWithInternalKey(deliveryScript).Unpack()
+		addr, err := p.addrWithInternalKey(deliveryScript)
 		if err != nil {
 			err = fmt.Errorf("unable to parse addr for channel "+
 				"%v: %w", req.ChanPoint, err)

server.go

@@ -18,6 +18,7 @@ import (
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcec/v2/ecdsa"
 	"github.com/btcsuite/btcd/btcutil"
+	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/connmgr"
 	"github.com/btcsuite/btcd/txscript"
@@ -520,6 +521,8 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 	var serializedPubKey [33]byte
 	copy(serializedPubKey[:], nodeKeyDesc.PubKey.SerializeCompressed())
 
+	netParams := cfg.ActiveNetParams.Params
+
 	// Initialize the sphinx router.
 	replayLog := htlcswitch.NewDecayedLog(
 		dbs.DecayedLogDB, cc.ChainNotifier,
@@ -1123,8 +1126,10 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 	})
 
 	s.sweeper = sweep.New(&sweep.UtxoSweeperConfig{
-		FeeEstimator:         cc.FeeEstimator,
-		GenSweepScript:       newSweepPkScriptGen(cc.Wallet),
+		FeeEstimator: cc.FeeEstimator,
+		GenSweepScript: newSweepPkScriptGen(
+			cc.Wallet, s.cfg.ActiveNetParams.Params,
+		),
 		Signer:               cc.Wallet.Cfg.Signer,
 		Wallet:               newSweeperWallet(cc.Wallet),
 		Mempool:              cc.MempoolNotifier,
@@ -1167,10 +1172,19 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 
 	s.breachArbitrator = contractcourt.NewBreachArbitrator(
 		&contractcourt.BreachConfig{
-			CloseLink:          closeLink,
-			DB:                 s.chanStateDB,
-			Estimator:          s.cc.FeeEstimator,
-			GenSweepScript:     newSweepPkScriptGen(cc.Wallet),
+			CloseLink: closeLink,
+			DB:        s.chanStateDB,
+			Estimator: s.cc.FeeEstimator,
+			GenSweepScript: func() ([]byte, error) {
+				addr, err := newSweepPkScriptGen(
+					cc.Wallet, netParams,
+				)().Unpack()
+				if err != nil {
+					return nil, err
+				}
+
+				return addr.DeliveryAddress, nil
+			},
 			Notifier:           cc.ChainNotifier,
 			PublishTransaction: cc.Wallet.PublishTransaction,
 			ContractBreaches:   contractBreaches,
@@ -1186,8 +1200,17 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 		ChainHash:              *s.cfg.ActiveNetParams.GenesisHash,
 		IncomingBroadcastDelta: lncfg.DefaultIncomingBroadcastDelta,
 		OutgoingBroadcastDelta: lncfg.DefaultOutgoingBroadcastDelta,
-		NewSweepAddr:           newSweepPkScriptGen(cc.Wallet),
-		PublishTx:              cc.Wallet.PublishTransaction,
+		NewSweepAddr: func() ([]byte, error) {
+			addr, err := newSweepPkScriptGen(
+				cc.Wallet, netParams,
+			)().Unpack()
+			if err != nil {
+				return nil, err
+			}
+
+			return addr.DeliveryAddress, nil
+		},
+		PublishTx: cc.Wallet.PublishTransaction,
 		DeliverResolutionMsg: func(msgs ...contractcourt.ResolutionMsg) error {
 			for _, msg := range msgs {
 				err := s.htlcSwitch.ProcessContractResolution(msg)
@@ -1665,8 +1688,17 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 				return s.channelNotifier.
 					SubscribeChannelEvents()
 			},
-			Signer:             cc.Wallet.Cfg.Signer,
-			NewAddress:         newSweepPkScriptGen(cc.Wallet),
+			Signer: cc.Wallet.Cfg.Signer,
+			NewAddress: func() ([]byte, error) {
+				addr, err := newSweepPkScriptGen(
+					cc.Wallet, netParams,
+				)().Unpack()
+				if err != nil {
+					return nil, err
+				}
+
+				return addr.DeliveryAddress, nil
+			},
 			SecretKeyRing:      s.cc.KeyRing,
 			Dial:               cfg.net.Dial,
 			AuthDial:           authDial,
@@ -4935,18 +4967,39 @@ func (s *server) SendCustomMessage(peerPub [33]byte, msgType lnwire.MessageType,
 // Specifically, the script generated is a version 0, pay-to-witness-pubkey-hash
 // (p2wkh) output.
 func newSweepPkScriptGen(
-	wallet lnwallet.WalletController) func() ([]byte, error) {
+	wallet lnwallet.WalletController,
+	netParams *chaincfg.Params) func() fn.Result[lnwallet.AddrWithKey] {
 
-	return func() ([]byte, error) {
+	return func() fn.Result[lnwallet.AddrWithKey] {
 		sweepAddr, err := wallet.NewAddress(
 			lnwallet.TaprootPubkey, false,
 			lnwallet.DefaultAccountName,
 		)
 		if err != nil {
-			return nil, err
+			return fn.Err[lnwallet.AddrWithKey](err)
 		}
 
-		return txscript.PayToAddrScript(sweepAddr)
+		addr, err := txscript.PayToAddrScript(sweepAddr)
+		if err != nil {
+			return fn.Err[lnwallet.AddrWithKey](err)
+		}
+
+		internalKeyDesc, err := lnwallet.InternalKeyForAddr(
+			wallet, netParams, addr,
+		)
+		if err != nil {
+			return fn.Err[lnwallet.AddrWithKey](err)
+		}
+
+		return fn.Ok(lnwallet.AddrWithKey{
+			DeliveryAddress: addr,
+			InternalKey: fn.MapOption(func(
+				desc keychain.KeyDescriptor) btcec.PublicKey {
+
+				return *desc.PubKey
+			},
+			)(internalKeyDesc),
+		})
 	}
 }
 

sweep/sweeper.go

@@ -318,7 +318,7 @@ type UtxoSweeper struct {
 type UtxoSweeperConfig struct {
 	// GenSweepScript generates a P2WKH script belonging to the wallet where
 	// funds can be swept.
-	GenSweepScript func() ([]byte, error)
+	GenSweepScript func() fn.Result[lnwallet.AddrWithKey]
 
 	// FeeEstimator is used when crafting sweep transactions to estimate
 	// the necessary fee relative to the expected size of the sweep
@@ -803,11 +803,11 @@ func (s *UtxoSweeper) signalResult(pi *SweeperInput, result Result) {
 func (s *UtxoSweeper) sweep(set InputSet) error {
 	// Generate an output script if there isn't an unused script available.
 	if s.currentOutputScript == nil {
-		pkScript, err := s.cfg.GenSweepScript()
+		pkScript, err := s.cfg.GenSweepScript().Unpack()
 		if err != nil {
 			return fmt.Errorf("gen sweep script: %w", err)
 		}
-		s.currentOutputScript = pkScript
+		s.currentOutputScript = pkScript.DeliveryAddress
 	}
 
 	// Create a fee bump request and ask the publisher to broadcast it. The

sweep/sweeper_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/lightningnetwork/lnd/chainntnfs"
 	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/input"
+	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -667,8 +668,11 @@ func TestSweepPendingInputs(t *testing.T) {
 		Wallet:     wallet,
 		Aggregator: aggregator,
 		Publisher:  publisher,
-		GenSweepScript: func() ([]byte, error) {
-			return testPubKey.SerializeCompressed(), nil
+		GenSweepScript: func() fn.Result[lnwallet.AddrWithKey] {
+			//nolint:lll
+			return fn.Ok(lnwallet.AddrWithKey{
+				DeliveryAddress: testPubKey.SerializeCompressed(),
+			})
 		},
 		NoDeadlineConfTarget: uint32(DefaultDeadlineDelta),
 	})