sphinx_test: onion message packet creation

TestTLVPayloadMessagePacket creates a onion message with payload and the
blinded route from the test vector. It then checks if the onion packet
we create is equal to the one provided in the test vector.

Author: gijswijs

sphinx_test.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/davecgh/go-spew/spew"
+	"github.com/lightningnetwork/lnd/tlv"
 	"github.com/stretchr/testify/require"
 )
 
@@ -162,6 +163,97 @@ func TestBolt4Packet(t *testing.T) {
 	}
 }
 
+// TestTLVPayloadMessagePacket tests the creation and encoding of an onion
+// message packet that uses a TLV payload for each hop in the route. This test
+// uses the test vectors defined in the BOLT 4 specification. The test reads a
+// JSON file containing a predefined route, session key, and the expected final
+// onion packet. It then constructs the route hop-by-hop, manually creating the
+// TLV payload for each, before creating a new onion packet with NewOnionPacket.
+// The test concludes by asserting that the newly encoded packet is identical to
+// the one specified in the test vector.
+func TestTLVPayloadMessagePacket(t *testing.T) {
+	t.Parallel()
+
+	// First, we'll read out the raw JSON file at the target location.
+	jsonBytes, err := os.ReadFile(testOnionMessageFileName)
+	require.NoError(t, err)
+
+	// Once we have the raw file, we'll unpack it into our jsonTestCase
+	// struct defined above.
+	testCase := &onionMessageJsonTestCase{}
+	require.NoError(t, json.Unmarshal(jsonBytes, testCase))
+
+	// Next, we'll populate a new OnionHop using the information included
+	// in this test case.
+	var route PaymentPath
+	for i, hop := range testCase.Route.Hops {
+		pubKeyBytes, err := hex.DecodeString(hop.BlindedNodeID)
+		require.NoError(t, err)
+
+		pubKey, err := btcec.ParsePubKey(pubKeyBytes)
+		require.NoError(t, err)
+
+		encryptedRecipientData, err := hex.DecodeString(
+			hop.EncryptedRecipientData,
+		)
+		require.NoError(t, err)
+
+		// Manually encode our onion payload
+		records := []tlv.Record{}
+
+		if i == len(testCase.Route.Hops)-1 {
+			helloBytes := []byte("hello")
+			records = append(records, tlv.MakePrimitiveRecord(
+				1, &helloBytes,
+			))
+		}
+
+		records = append(records, tlv.MakePrimitiveRecord(
+			4, &encryptedRecipientData,
+		))
+
+		stream, err := tlv.NewStream(records...)
+		require.NoError(t, err, "new stream")
+
+		b := new(bytes.Buffer)
+		require.NoError(t, stream.Encode(b), "encode")
+
+		route[i] = OnionHop{
+			NodePub: *pubKey,
+			HopPayload: HopPayload{
+				Type:    PayloadTLV,
+				Payload: b.Bytes(),
+			},
+		}
+	}
+
+	finalPacket, err := hex.DecodeString(
+		testCase.OnionMessage.OnionMessagePacket,
+	)
+	require.NoError(t, err)
+
+	sessionKeyBytes, err := hex.DecodeString(testCase.Generate.SessionKey)
+
+	require.NoError(t, err)
+
+	// With all the required data assembled, we'll craft a new packet.
+	sessionKey, _ := btcec.PrivKeyFromBytes(sessionKeyBytes)
+
+	pkt, err := NewOnionPacket(
+		&route, sessionKey, nil, DeterministicPacketFiller,
+	)
+	require.NoError(t, err)
+
+	var b bytes.Buffer
+	require.NoError(t, pkt.Encode(&b))
+
+	// Finally, we expect that our packet matches the packet included in
+	// the spec's test vectors.
+	require.Equalf(t, finalPacket, b.Bytes(), "final packet does not "+
+		"match expected BOLT 4 packet, want: %s, got %s",
+		hex.EncodeToString(finalPacket), hex.EncodeToString(b.Bytes()))
+}
+
 func TestSphinxCorrectness(t *testing.T) {
 	nodes, _, hopDatas, fwdMsg, err := newTestRoute(testLegacyRouteNumHops)
 	if err != nil {
@@ -755,6 +847,9 @@ const (
 
 	// testTLVFileName is the name of the tlv-payload-only onion test file.
 	testTLVFileName = "testdata/onion-test.json"
+
+	// testOnionMessageFileName is the name of the onion message test file.
+	testOnionMessageFileName = "testdata/blinded-onion-message-onion-test.json"
 )
 
 type jsonHop struct {