lnwire: refactor custom vs official tlv parsing into new func

We'll use this to update CommitSig in the next commit.

Author: Roasbeef

lnwire/extra_bytes.go

@@ -1,9 +1,15 @@
 package lnwire
 
+// For some reason golangci-lint has a false positive on the sort order of the
+// imports for the new "maps" package... We need the nolint directive here to
+// ignore that.
+//
+//nolint:gci
 import (
 	"bytes"
 	"fmt"
 	"io"
+	"maps"
 
 	"github.com/lightningnetwork/lnd/tlv"
 )
@@ -194,3 +200,57 @@ func EncodeMessageExtraData(extraData *ExtraOpaqueData,
 	// are all properly sorted.
 	return extraData.PackRecords(recordProducers...)
 }
+
+// wireTlvMap is a struct that holds the official records and custom records in
+// a TLV type map. This is useful for ensuring that the set of custom TLV
+// records are handled properly and don't overlap with the official records.
+type wireTlvMap struct {
+	// officialTypes is the set of official records that are defined in the
+	// spec.
+	officialTypes tlv.TypeMap
+
+	// customTypes is the set of custom records that are not defined in
+	// spec, and are used by higher level applications.
+	customTypes tlv.TypeMap
+}
+
+// newWireTlvMap creates a new tlv.TypeMap from the given set of parsed TLV
+// records. A struct with two maps are returned:
+//
+//  1. officialTypes: the set of official records that are defined in the
+//     spec.
+//
+//  2. customTypes: the set of custom records that are not defined in
+//     the spec.
+func newWireTlvMap(typeMap tlv.TypeMap) wireTlvMap {
+	officialRecords := maps.Clone(typeMap)
+
+	// Any records from the extra data TLV map which are in the custom
+	// records TLV type range will be included in the custom records field
+	// and removed from the extra data field.
+	customRecordsTlvMap := make(tlv.TypeMap, len(typeMap))
+	for k, v := range typeMap {
+		// Skip records that are not in the custom records TLV type
+		// range.
+		if k < MinCustomRecordsTlvType {
+			continue
+		}
+
+		// Include the record in the custom records map.
+		customRecordsTlvMap[k] = v
+
+		// Now that the record is included in the custom records map,
+		// we can remove it from the extra data TLV map.
+		delete(officialRecords, k)
+	}
+
+	return wireTlvMap{
+		officialTypes: officialRecords,
+		customTypes:   customRecordsTlvMap,
+	}
+}
+
+// Len returns the total number of records in the wireTlvMap.
+func (w *wireTlvMap) Len() int {
+	return len(w.officialTypes) + len(w.customTypes)
+}

lnwire/extra_bytes_test.go

@@ -7,8 +7,11 @@ import (
 	"testing"
 	"testing/quick"
 
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/tlv"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/maps"
+	"pgregory.net/rapid"
 )
 
 // TestExtraOpaqueDataEncodeDecode tests that we're able to encode/decode
@@ -206,3 +209,47 @@ func TestPackRecords(t *testing.T) {
 	require.Equal(t, recordBytes2, extractedRecords[tlvType2.TypeVal()])
 	require.Equal(t, recordBytes3, extractedRecords[tlvType3.TypeVal()])
 }
+
+// TestNewWireTlvMap tests the newWireTlvMap function using property-based
+// testing.
+func TestNewWireTlvMap(t *testing.T) {
+	rapid.Check(t, func(t *rapid.T) {
+		// Make a random type map, using the generic Make which'll
+		// figure out what type to generate.
+		tlvTypeMap := rapid.Make[tlv.TypeMap]().Draw(t, "typeMap")
+
+		// Create a wireTlvMap from the generated type map, this'll
+		// operate on our random input.
+		result := newWireTlvMap(tlvTypeMap)
+
+		// Property 1: The sum of lengths of officialTypes and
+		// customTypes should equal the length of the input typeMap.
+		require.Equal(t, len(tlvTypeMap), result.Len())
+
+		// Property 2: All types in customTypes should be >=
+		// MinCustomRecordsTlvType.
+		require.True(t, fn.All(func(k tlv.Type) bool {
+			return uint64(k) >= uint64(MinCustomRecordsTlvType)
+		}, maps.Keys(result.customTypes)))
+
+		// Property 3: All types in officialTypes should be <
+		// MinCustomRecordsTlvType.
+		require.True(t, fn.All(func(k tlv.Type) bool {
+			return uint64(k) < uint64(MinCustomRecordsTlvType)
+		}, maps.Keys(result.officialTypes)))
+
+		// Property 4: The union of officialTypes and customTypes
+		// should equal the input typeMap.
+		unionMap := make(tlv.TypeMap)
+		maps.Copy(unionMap, result.officialTypes)
+		maps.Copy(unionMap, result.customTypes)
+		require.Equal(t, tlvTypeMap, unionMap)
+
+		// Property 5: No type should appear in both officialTypes and
+		// customTypes.
+		require.True(t, fn.All(func(k tlv.Type) bool {
+			_, exists := result.officialTypes[k]
+			return !exists
+		}, maps.Keys(result.customTypes)))
+	})
+}

lnwire/update_add_htlc.go

@@ -131,29 +131,14 @@ func (c *UpdateAddHTLC) Decode(r io.Reader, pver uint32) error {
 		delete(extraDataTlvMap, c.BlindingPoint.TlvType())
 	}
 
-	// Any records from the extra data TLV map which are in the custom
-	// records TLV type range will be included in the custom records field
-	// and removed from the extra data field.
-	customRecordsTlvMap := make(tlv.TypeMap, len(extraDataTlvMap))
-	for k, v := range extraDataTlvMap {
-		// Skip records that are not in the custom records TLV type
-		// range.
-		if k < MinCustomRecordsTlvType {
-			continue
-		}
-
-		// Include the record in the custom records map.
-		customRecordsTlvMap[k] = v
-
-		// Now that the record is included in the custom records map,
-		// we can remove it from the extra data TLV map.
-		delete(extraDataTlvMap, k)
-	}
+	// Parse through the remaining extra data map to separate the custom
+	// records, from the set of official records.
+	tlvTypes := newWireTlvMap(extraDataTlvMap)
 
 	// Set the custom records field to the custom records specific TLV
 	// record map.
 	customRecords, err := NewCustomRecordsFromTlvTypeMap(
-		customRecordsTlvMap,
+		tlvTypes.customTypes,
 	)
 	if err != nil {
 		return err
@@ -162,21 +147,23 @@ func (c *UpdateAddHTLC) Decode(r io.Reader, pver uint32) error {
 
 	// Set custom records to nil if we didn't parse anything out of it so
 	// that we can use assert.Equal in tests.
-	if len(customRecordsTlvMap) == 0 {
+	if len(customRecords) == 0 {
 		c.CustomRecords = nil
 	}
 
 	// Set extra data to nil if we didn't parse anything out of it so that
 	// we can use assert.Equal in tests.
-	if len(extraDataTlvMap) == 0 {
+	if len(tlvTypes.officialTypes) == 0 {
 		c.ExtraData = nil
 		return nil
 	}
 
 	// Encode the remaining records back into the extra data field. These
-	// records are not in the custom records TLV type range and do not
-	// have associated fields in the UpdateAddHTLC struct.
-	c.ExtraData, err = NewExtraOpaqueDataFromTlvTypeMap(extraDataTlvMap)
+	// records are not in the custom records TLV type range and do not have
+	// associated fields in the UpdateAddHTLC struct.
+	c.ExtraData, err = NewExtraOpaqueDataFromTlvTypeMap(
+		tlvTypes.officialTypes,
+	)
 	if err != nil {
 		return err
 	}