From b9786e1f207eea394869b7e833940dac98fb01a0 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 |  5 +++++
 peer/brontide.go  | 15 ++++++++++++++-
 server.go         |  1 +
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/config_builder.go b/config_builder.go
index 7af327382..e5839208d 100644
--- a/config_builder.go
+++ b/config_builder.go
@@ -44,6 +44,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnwallet/btcwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/rpcwallet"
 	"github.com/lightningnetwork/lnd/macaroons"
+	"github.com/lightningnetwork/lnd/peer"
 	"github.com/lightningnetwork/lnd/rpcperms"
 	"github.com/lightningnetwork/lnd/signal"
 	"github.com/lightningnetwork/lnd/sqldb"
@@ -157,6 +158,10 @@ type AuxComponents struct {
 	// AuxLeafStore is an optional data source that can be used by custom
 	// channels to fetch+store various data.
 	AuxLeafStore fn.Option[lnwallet.AuxLeafStore]
+
+	// MsgRouter is an optional message router that if set will be used in
+	// place of a new balnk default message router.
+	MsgRouter fn.Option[peer.MsgRouter]
 }
 
 // DefaultWalletImpl is the default implementation of our normal, btcwallet
diff --git a/peer/brontide.go b/peer/brontide.go
index 230f3cb41..b86a2db62 100644
--- a/peer/brontide.go
+++ b/peer/brontide.go
@@ -374,6 +374,11 @@ type Config struct {
 	// invalid.
 	DisallowRouteBlinding bool
 
+	// 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[MsgRouter]
+
 	// Quit is the server's quit channel. If this is closed, we halt operation.
 	Quit chan struct{}
 }
@@ -512,6 +517,14 @@ 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.
+	//
+	// TODO(roasbeef): extend w/ source peer info?
+	msgRouter := cfg.MsgRouter.Alt(
+		fn.Some[MsgRouter](NewMultiMsgRouter()),
+	)
+
 	p := &Brontide{
 		cfg:           cfg,
 		activeSignal:  make(chan struct{}),
@@ -534,7 +547,7 @@ func NewBrontide(cfg Config) *Brontide {
 		startReady:         make(chan struct{}),
 		quit:               make(chan struct{}),
 		log:                build.NewPrefixLog(logPrefix, peerLog),
-		msgRouter:          fn.Some[MsgRouter](NewMultiMsgRouter()),
+		msgRouter:          msgRouter,
 	}
 
 	var (
diff --git a/server.go b/server.go
index 650cdbced..e53a7bb07 100644
--- a/server.go
+++ b/server.go
@@ -3912,6 +3912,7 @@ func (s *server) peerConnected(conn net.Conn, connReq *connmgr.ConnReq,
 		DisallowRouteBlinding:  s.cfg.ProtocolOptions.NoRouteBlinding(),
 		Quit:                   s.quit,
 		AuxLeafStore:           s.implCfg.AuxLeafStore,
+		MsgRouter:              s.implCfg.MsgRouter,
 	}
 
 	copy(pCfg.PubKeyBytes[:], peerAddr.IdentityKey.SerializeCompressed())