From c2ce8635e739bd64b4991bf6c43e4807644dbde3 Mon Sep 17 00:00:00 2001
From: Carla Kirk-Cohen <kirkcohenc@gmail.com>
Date: Wed, 6 Jul 2022 09:26:47 +0200
Subject: [PATCH] multi: add custom message overrides to protocol config
 options

---
 config.go                          | 12 ++++++++++++
 lncfg/protocol_experimental_off.go |  6 ++++++
 lncfg/protocol_experimental_on.go  |  9 +++++++++
 3 files changed, 27 insertions(+)

diff --git a/config.go b/config.go
index 4a7123bf2..39818397c 100644
--- a/config.go
+++ b/config.go
@@ -39,6 +39,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnrpc/routerrpc"
 	"github.com/lightningnetwork/lnd/lnrpc/signrpc"
 	"github.com/lightningnetwork/lnd/lnwallet"
+	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/routing"
 	"github.com/lightningnetwork/lnd/signal"
 	"github.com/lightningnetwork/lnd/sweep"
@@ -1656,6 +1657,17 @@ func ValidateConfig(cfg Config, interceptor signal.Interceptor, fileParser,
 			lncfg.DefaultIncomingBroadcastDelta)
 	}
 
+	// If the experimental protocol options specify any protocol messages
+	// that we want to handle as custom messages, set them now.
+	//nolint:lll
+	customMsg := cfg.ProtocolOptions.ExperimentalProtocol.CustomMessageOverrides()
+
+	// We can safely set our custom override values during startup because
+	// startup is blocked on config parsing.
+	if err := lnwire.SetCustomOverrides(customMsg); err != nil {
+		return nil, mkErr("custom-message: %v", err)
+	}
+
 	// Validate the subconfigs for workers, caches, and the tower client.
 	err = lncfg.Validate(
 		cfg.Workers,
diff --git a/lncfg/protocol_experimental_off.go b/lncfg/protocol_experimental_off.go
index 34143a224..e88b4b52b 100644
--- a/lncfg/protocol_experimental_off.go
+++ b/lncfg/protocol_experimental_off.go
@@ -7,3 +7,9 @@ package lncfg
 // features that also require a build-tag to activate.
 type ExperimentalProtocol struct {
 }
+
+// CustomMessageOverrides returns the set of protocol messages that we override
+// to allow custom handling.
+func (p ExperimentalProtocol) CustomMessageOverrides() []uint16 {
+	return nil
+}
diff --git a/lncfg/protocol_experimental_on.go b/lncfg/protocol_experimental_on.go
index b7d74acfe..cf49890ae 100644
--- a/lncfg/protocol_experimental_on.go
+++ b/lncfg/protocol_experimental_on.go
@@ -5,5 +5,14 @@ package lncfg
 
 // ExperimentalProtocol is a sub-config that houses any experimental protocol
 // features that also require a build-tag to activate.
+//
+//nolint:lll
 type ExperimentalProtocol struct {
+	CustomMessage []uint16 `long:"custom-message" description:"allows the custom message apis to send and report messages with the protocol number provided that fall outside of the custom message number range."`
+}
+
+// CustomMessageOverrides returns the set of protocol messages that we override
+// to allow custom handling.
+func (p ExperimentalProtocol) CustomMessageOverrides() []uint16 {
+	return p.CustomMessage
 }