From 019da315a8fe2cdb9d9873b9a00e476097f382b5 Mon Sep 17 00:00:00 2001
From: Carla Kirk-Cohen <kirkcohenc@gmail.com>
Date: Tue, 25 Apr 2023 09:38:41 -0400
Subject: [PATCH] itest: add coverage for custom features

---
 itest/list_on_test.go        |  4 ++
 itest/lnd_custom_features.go | 96 ++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)
 create mode 100644 itest/lnd_custom_features.go

diff --git a/itest/list_on_test.go b/itest/list_on_test.go
index ade6a80f6..cc735afad 100644
--- a/itest/list_on_test.go
+++ b/itest/list_on_test.go
@@ -531,4 +531,8 @@ var allTestCases = []*lntest.TestCase{
 		Name:     "htlc timeout resolver extract preimage local",
 		TestFunc: testHtlcTimeoutResolverExtractPreimageLocal,
 	},
+	{
+		Name:     "custom features",
+		TestFunc: testCustomFeatures,
+	},
 }
diff --git a/itest/lnd_custom_features.go b/itest/lnd_custom_features.go
new file mode 100644
index 000000000..4867ca2d1
--- /dev/null
+++ b/itest/lnd_custom_features.go
@@ -0,0 +1,96 @@
+package itest
+
+import (
+	"fmt"
+
+	"github.com/lightningnetwork/lnd/lnrpc"
+	"github.com/lightningnetwork/lnd/lnrpc/peersrpc"
+	"github.com/lightningnetwork/lnd/lntest"
+	"github.com/stretchr/testify/require"
+)
+
+// testCustomFeatures tests advertisement of custom features in various bolt 9
+// sets. For completeness, it also asserts that features aren't set in places
+// where they aren't intended to be.
+func testCustomFeatures(ht *lntest.HarnessTest) {
+	var (
+		// Odd custom features so that we don't need to worry about
+		// issues connecting to peers.
+		customInit    uint32 = 101
+		customNodeAnn uint32 = 103
+		customInvoice uint32 = 105
+	)
+
+	// Restart Alice with custom feature bits configured.
+	extraArgs := []string{
+		fmt.Sprintf("--protocol.custom-init=%v", customInit),
+		fmt.Sprintf("--protocol.custom-nodeann=%v", customNodeAnn),
+		fmt.Sprintf("--protocol.custom-invoice=%v", customInvoice),
+	}
+	ht.RestartNodeWithExtraArgs(ht.Alice, extraArgs)
+
+	// Connect nodes and open a channel so that Alice will be included
+	// in Bob's graph.
+	ht.ConnectNodes(ht.Alice, ht.Bob)
+	chanPoint := ht.OpenChannel(
+		ht.Alice, ht.Bob, lntest.OpenChannelParams{Amt: 1000000},
+	)
+
+	// Check that Alice's custom feature bit was sent to Bob in her init
+	// message.
+	peers := ht.Bob.RPC.ListPeers()
+	require.Len(ht, peers.Peers, 1)
+	require.Equal(ht, peers.Peers[0].PubKey, ht.Alice.PubKeyStr)
+
+	_, customInitSet := peers.Peers[0].Features[customInit]
+	require.True(ht, customInitSet)
+	assertFeatureNotInSet(
+		ht, []uint32{customNodeAnn, customInvoice},
+		peers.Peers[0].Features,
+	)
+
+	// Assert that Alice's custom feature bit is contained in the node
+	// announcement sent to Bob.
+	updates := ht.AssertNumNodeAnns(ht.Bob, ht.Alice.PubKeyStr, 1)
+	features := updates[len(updates)-1].Features
+	_, customFeature := features[customNodeAnn]
+	require.True(ht, customFeature)
+	assertFeatureNotInSet(
+		ht, []uint32{customInit, customInvoice}, features,
+	)
+
+	// Assert that Alice's custom feature bit is included in invoices.
+	invoice := ht.Alice.RPC.AddInvoice(&lnrpc.Invoice{})
+	payReq := ht.Alice.RPC.DecodePayReq(invoice.PaymentRequest)
+	_, customInvoiceSet := payReq.Features[customInvoice]
+	require.True(ht, customInvoiceSet)
+	assertFeatureNotInSet(
+		ht, []uint32{customInit, customNodeAnn}, payReq.Features,
+	)
+
+	// Check that Alice can't unset configured features via the node
+	// announcement update API. This is only checked for node announcement
+	// because it is the only set that can be updated via the API.
+	nodeAnnReq := &peersrpc.NodeAnnouncementUpdateRequest{
+		FeatureUpdates: []*peersrpc.UpdateFeatureAction{
+			{
+				Action:     peersrpc.UpdateAction_ADD,
+				FeatureBit: lnrpc.FeatureBit(customNodeAnn),
+			},
+		},
+	}
+	ht.Alice.RPC.UpdateNodeAnnouncementErr(nodeAnnReq)
+
+	ht.CloseChannel(ht.Alice, chanPoint)
+}
+
+// assertFeatureNotInSet checks that the features provided aren't contained in
+// a feature map.
+func assertFeatureNotInSet(ht *lntest.HarnessTest, features []uint32,
+	advertised map[uint32]*lnrpc.Feature) {
+
+	for _, feature := range features {
+		_, featureSet := advertised[feature]
+		require.False(ht, featureSet)
+	}
+}