From 4e36e2fd2b24d8e81e7c3503315d9e7b2743ae35 Mon Sep 17 00:00:00 2001
From: Carla Kirk-Cohen <kirkcohenc@gmail.com>
Date: Mon, 3 Apr 2023 11:53:02 +0200
Subject: [PATCH] multi: add custom feature bits to config options

---
 feature/manager.go                 | 21 +++++++++++++++++++
 lncfg/protocol_experimental_off.go | 10 +++++++++
 lncfg/protocol_experimental_on.go  | 33 ++++++++++++++++++++++++++++++
 server.go                          |  2 ++
 4 files changed, 66 insertions(+)

diff --git a/feature/manager.go b/feature/manager.go
index ecc7470da..c1f036fd9 100644
--- a/feature/manager.go
+++ b/feature/manager.go
@@ -49,6 +49,10 @@ type Config struct {
 	// NoAnySegwit unsets any bits that signal support for using other
 	// segwit witness versions for co-op closes.
 	NoAnySegwit bool
+
+	// CustomFeatures is a set of custom features to advertise in each
+	// set.
+	CustomFeatures map[Set][]lnwire.FeatureBit
 }
 
 // Manager is responsible for generating feature vectors for different requested
@@ -156,6 +160,23 @@ func newManager(cfg Config, desc setDesc) (*Manager, error) {
 			raw.Unset(lnwire.ShutdownAnySegwitRequired)
 		}
 
+		for _, custom := range cfg.CustomFeatures[set] {
+			if custom > set.Maximum() {
+				return nil, fmt.Errorf("feature bit: %v "+
+					"exceeds set: %v maximum: %v", custom,
+					set, set.Maximum())
+			}
+
+			if raw.IsSet(custom) {
+				return nil, fmt.Errorf("feature bit: %v "+
+					"already set", custom)
+			}
+
+			if err := raw.SafeSet(custom); err != nil {
+				return nil, fmt.Errorf("%w: could not set "+
+					"feature: %d", err, custom)
+			}
+		}
 		// Ensure that all of our feature sets properly set any
 		// dependent features.
 		fv := lnwire.NewFeatureVector(raw, lnwire.Features)
diff --git a/lncfg/protocol_experimental_off.go b/lncfg/protocol_experimental_off.go
index e88b4b52b..51c90ec38 100644
--- a/lncfg/protocol_experimental_off.go
+++ b/lncfg/protocol_experimental_off.go
@@ -3,6 +3,11 @@
 
 package lncfg
 
+import (
+	"github.com/lightningnetwork/lnd/feature"
+	"github.com/lightningnetwork/lnd/lnwire"
+)
+
 // ExperimentalProtocol is a sub-config that houses any experimental protocol
 // features that also require a build-tag to activate.
 type ExperimentalProtocol struct {
@@ -13,3 +18,8 @@ type ExperimentalProtocol struct {
 func (p ExperimentalProtocol) CustomMessageOverrides() []uint16 {
 	return nil
 }
+
+// CustomFeatures returns a custom set of feature bits to advertise.
+func (p ExperimentalProtocol) CustomFeatures() map[feature.Set][]lnwire.FeatureBit {
+	return map[feature.Set][]lnwire.FeatureBit{}
+}
diff --git a/lncfg/protocol_experimental_on.go b/lncfg/protocol_experimental_on.go
index cf49890ae..bb23e83c2 100644
--- a/lncfg/protocol_experimental_on.go
+++ b/lncfg/protocol_experimental_on.go
@@ -3,12 +3,23 @@
 
 package lncfg
 
+import (
+	"github.com/lightningnetwork/lnd/feature"
+	"github.com/lightningnetwork/lnd/lnwire"
+)
+
 // 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."`
+
+	CustomInit []uint16 `long:"custom-init" description:"custom feature bits to advertise in the node's init message"`
+
+	CustomNodeAnn []uint16 `long:"custom-nodeann" description:"custom feature bits to advertise in the node's announcement message"`
+
+	CustomInvoice []uint16 `long:"custom-invoice" description:"custom feature bits to advertise in the node's invoices"`
 }
 
 // CustomMessageOverrides returns the set of protocol messages that we override
@@ -16,3 +27,25 @@ type ExperimentalProtocol struct {
 func (p ExperimentalProtocol) CustomMessageOverrides() []uint16 {
 	return p.CustomMessage
 }
+
+// CustomFeatures returns a custom set of feature bits to advertise.
+//
+//nolint:lll
+func (p ExperimentalProtocol) CustomFeatures() map[feature.Set][]lnwire.FeatureBit {
+	customFeatures := make(map[feature.Set][]lnwire.FeatureBit)
+
+	setFeatures := func(set feature.Set, bits []uint16) {
+		for _, customFeature := range bits {
+			customFeatures[set] = append(
+				customFeatures[set],
+				lnwire.FeatureBit(customFeature),
+			)
+		}
+	}
+
+	setFeatures(feature.SetInit, p.CustomInit)
+	setFeatures(feature.SetNodeAnn, p.CustomNodeAnn)
+	setFeatures(feature.SetInvoice, p.CustomInvoice)
+
+	return customFeatures
+}
diff --git a/server.go b/server.go
index d70324d81..d698535e9 100644
--- a/server.go
+++ b/server.go
@@ -530,6 +530,7 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 		readBufferPool, cfg.Workers.Read, pool.DefaultWorkerTimeout,
 	)
 
+	//nolint:lll
 	featureMgr, err := feature.NewManager(feature.Config{
 		NoTLVOnion:               cfg.ProtocolOptions.LegacyOnion(),
 		NoStaticRemoteKey:        cfg.ProtocolOptions.NoStaticRemoteKey(),
@@ -540,6 +541,7 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 		NoOptionScidAlias:        !cfg.ProtocolOptions.ScidAlias(),
 		NoZeroConf:               !cfg.ProtocolOptions.ZeroConf(),
 		NoAnySegwit:              cfg.ProtocolOptions.NoAnySegwit(),
+		CustomFeatures:           cfg.ProtocolOptions.ExperimentalProtocol.CustomFeatures(),
 	})
 	if err != nil {
 		return nil, err