From b028af1836b22a2550cb4ad47b1e1d52f500d5d0 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Thu, 4 Apr 2024 16:13:30 -0700
Subject: [PATCH] multi: make MsgRouter available in the ImplementationCfg

With this commit, we allow the `MsgRouter` to be available in the
`ImplementationCfg`. With this, programs outside of lnd itself are able
to now hook into the message processing flow to direct handle custom
messages, and even normal wire messages.
---
 config_builder.go | 14 ++++++++++++++
 lnd.go            |  1 +
 peer/brontide.go  | 15 ++++++++++++---
 server.go         |  7 ++++++-
 4 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/config_builder.go b/config_builder.go
index bf6274cdf..bef59b9a0 100644
--- a/config_builder.go
+++ b/config_builder.go
@@ -33,6 +33,7 @@ import (
 	"github.com/lightningnetwork/lnd/chainreg"
 	"github.com/lightningnetwork/lnd/channeldb"
 	"github.com/lightningnetwork/lnd/clock"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/invoices"
 	"github.com/lightningnetwork/lnd/keychain"
 	"github.com/lightningnetwork/lnd/kvdb"
@@ -42,6 +43,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnwallet/btcwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/rpcwallet"
 	"github.com/lightningnetwork/lnd/macaroons"
+	"github.com/lightningnetwork/lnd/msgmux"
 	"github.com/lightningnetwork/lnd/rpcperms"
 	"github.com/lightningnetwork/lnd/signal"
 	"github.com/lightningnetwork/lnd/sqldb"
@@ -118,6 +120,14 @@ type ChainControlBuilder interface {
 		*btcwallet.Config) (*chainreg.ChainControl, func(), error)
 }
 
+// AuxComponents is a set of auxiliary components that can be used by lnd for
+// certain custom channel types.
+type AuxComponents struct {
+	// MsgRouter is an optional message router that if set will be used in
+	// place of a new blank default message router.
+	MsgRouter fn.Option[msgmux.Router]
+}
+
 // ImplementationCfg is a struct that holds all configuration items for
 // components that can be implemented outside lnd itself.
 type ImplementationCfg struct {
@@ -144,6 +154,10 @@ type ImplementationCfg struct {
 	// ChainControlBuilder is a type that can provide a custom wallet
 	// implementation.
 	ChainControlBuilder
+
+	// AuxComponents is a set of auxiliary components that can be used by
+	// lnd for certain custom channel types.
+	AuxComponents
 }
 
 // DefaultWalletImpl is the default implementation of our normal, btcwallet
diff --git a/lnd.go b/lnd.go
index 38f5c0d75..e483d5512 100644
--- a/lnd.go
+++ b/lnd.go
@@ -600,6 +600,7 @@ func Main(cfg *Config, lisCfg ListenerCfg, implCfg *ImplementationCfg,
 		cfg, cfg.Listeners, dbs, activeChainControl, &idKeyDesc,
 		activeChainControl.Cfg.WalletUnlockParams.ChansToRestore,
 		multiAcceptor, torController, tlsManager, leaderElector,
+		implCfg,
 	)
 	if err != nil {
 		return mkErr("unable to create server: %v", err)
diff --git a/peer/brontide.go b/peer/brontide.go
index d56194850..a6cb4da88 100644
--- a/peer/brontide.go
+++ b/peer/brontide.go
@@ -387,6 +387,11 @@ type Config struct {
 	// This value will be passed to created links.
 	MaxFeeExposure lnwire.MilliSatoshi
 
+	// MsgRouter is an optional instance of the main message router that
+	// the peer will use. If None, then a new default version will be used
+	// in place.
+	MsgRouter fn.Option[msgmux.Router]
+
 	// Quit is the server's quit channel. If this is closed, we halt operation.
 	Quit chan struct{}
 }
@@ -542,6 +547,12 @@ var _ lnpeer.Peer = (*Brontide)(nil)
 func NewBrontide(cfg Config) *Brontide {
 	logPrefix := fmt.Sprintf("Peer(%x):", cfg.PubKeyBytes)
 
+	// We'll either use the msg router instance passed in, or create a new
+	// blank instance.
+	msgRouter := cfg.MsgRouter.Alt(fn.Some[msgmux.Router](
+		msgmux.NewMultiMsgRouter(),
+	))
+
 	p := &Brontide{
 		cfg:           cfg,
 		activeSignal:  make(chan struct{}),
@@ -564,9 +575,7 @@ func NewBrontide(cfg Config) *Brontide {
 		startReady:         make(chan struct{}),
 		quit:               make(chan struct{}),
 		log:                build.NewPrefixLog(logPrefix, peerLog),
-		msgRouter: fn.Some[msgmux.Router](
-			msgmux.NewMultiMsgRouter(),
-		),
+		msgRouter:          msgRouter,
 	}
 
 	if cfg.Conn != nil && cfg.Conn.RemoteAddr() != nil {
diff --git a/server.go b/server.go
index 555ee2627..0cedb2f48 100644
--- a/server.go
+++ b/server.go
@@ -160,6 +160,8 @@ type server struct {
 
 	cfg *Config
 
+	implCfg *ImplementationCfg
+
 	// identityECDH is an ECDH capable wrapper for the private key used
 	// to authenticate any incoming connections.
 	identityECDH keychain.SingleKeyECDH
@@ -486,7 +488,8 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 	chansToRestore walletunlocker.ChannelsToRecover,
 	chanPredicate chanacceptor.ChannelAcceptor,
 	torController *tor.Controller, tlsManager *TLSManager,
-	leaderElector cluster.LeaderElector) (*server, error) {
+	leaderElector cluster.LeaderElector,
+	implCfg *ImplementationCfg) (*server, error) {
 
 	var (
 		err         error
@@ -571,6 +574,7 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 
 	s := &server{
 		cfg:            cfg,
+		implCfg:        implCfg,
 		graphDB:        dbs.GraphDB.ChannelGraph(),
 		chanStateDB:    dbs.ChanStateDB.ChannelStateDB(),
 		addrSource:     dbs.ChanStateDB,
@@ -3986,6 +3990,7 @@ func (s *server) peerConnected(conn net.Conn, connReq *connmgr.ConnReq,
 		DisallowRouteBlinding:  s.cfg.ProtocolOptions.NoRouteBlinding(),
 		MaxFeeExposure:         thresholdMSats,
 		Quit:                   s.quit,
+		MsgRouter:              s.implCfg.MsgRouter,
 	}
 
 	copy(pCfg.PubKeyBytes[:], peerAddr.IdentityKey.SerializeCompressed())