From ca6d414308c13100a2d30beeac1c636399cfecf3 Mon Sep 17 00:00:00 2001
From: Carla Kirk-Cohen <kirkcohenc@gmail.com>
Date: Tue, 2 Apr 2024 15:20:25 -0400
Subject: [PATCH] multi: validate contents in blinded data against payload

---
 htlcswitch/hop/iterator.go      | 11 +++++-
 htlcswitch/hop/iterator_test.go |  1 +
 htlcswitch/hop/payload.go       | 34 +++++++++++++++++
 htlcswitch/hop/payload_test.go  | 65 +++++++++++++++++++++++++++++++++
 4 files changed, 110 insertions(+), 1 deletion(-)

diff --git a/htlcswitch/hop/iterator.go b/htlcswitch/hop/iterator.go
index 7ddb9af77..87691d192 100644
--- a/htlcswitch/hop/iterator.go
+++ b/htlcswitch/hop/iterator.go
@@ -200,7 +200,8 @@ func (b *BlindingKit) validateBlindingPoint(payloadBlinding *btcec.PublicKey,
 // DecryptAndValidateFwdInfo performs all operations required to decrypt and
 // validate a blinded route.
 func (b *BlindingKit) DecryptAndValidateFwdInfo(payload *Payload,
-	isFinalHop bool) (*ForwardingInfo, error) {
+	isFinalHop bool, payloadParsed map[tlv.Type][]byte) (
+	*ForwardingInfo, error) {
 
 	// We expect this function to be called when we have encrypted data
 	// present, and a blinding key is set either in the payload or the
@@ -227,6 +228,14 @@ func (b *BlindingKit) DecryptAndValidateFwdInfo(payload *Payload,
 			ErrDecodeFailed, err)
 	}
 
+	// Validate the contents of the payload against the values we've
+	// just pulled out of the encrypted data blob.
+	err = ValidatePayloadWithBlinded(isFinalHop, payloadParsed)
+	if err != nil {
+		return nil, err
+	}
+	// Validate the data in the blinded route against our incoming htlc's
+	// information.
 	if err := ValidateBlindedRouteData(
 		routeData, b.IncomingAmount, b.IncomingCltv,
 	); err != nil {
diff --git a/htlcswitch/hop/iterator_test.go b/htlcswitch/hop/iterator_test.go
index b8b8b4f72..e39dbb151 100644
--- a/htlcswitch/hop/iterator_test.go
+++ b/htlcswitch/hop/iterator_test.go
@@ -287,6 +287,7 @@ func TestDecryptAndValidateFwdInfo(t *testing.T) {
 					encryptedData: testCase.data,
 					blindingPoint: testCase.payloadBlinding,
 				}, false,
+				make(map[tlv.Type][]byte),
 			)
 			require.ErrorIs(t, err, testCase.expectedErr)
 		})
diff --git a/htlcswitch/hop/payload.go b/htlcswitch/hop/payload.go
index cbd8d08a5..99c36db6d 100644
--- a/htlcswitch/hop/payload.go
+++ b/htlcswitch/hop/payload.go
@@ -484,3 +484,37 @@ func ValidateBlindedRouteData(blindedData *record.BlindedRouteData,
 
 	return nil
 }
+
+// ValidatePayloadWithBlinded validates a payload against the contents of
+// its encrypted data blob.
+func ValidatePayloadWithBlinded(isFinalHop bool,
+	payloadParsed map[tlv.Type][]byte) error {
+
+	// Blinded routes restrict the presence of TLVs more strictly than
+	// regular routes, check that intermediate and final hops only have
+	// the TLVs the spec allows them to have.
+	allowedTLVs := map[tlv.Type]bool{
+		record.EncryptedDataOnionType: true,
+		record.BlindingPointOnionType: true,
+	}
+
+	if isFinalHop {
+		allowedTLVs[record.AmtOnionType] = true
+		allowedTLVs[record.LockTimeOnionType] = true
+		allowedTLVs[record.TotalAmtMsatBlindedType] = true
+	}
+
+	for tlvType := range payloadParsed {
+		if _, ok := allowedTLVs[tlvType]; ok {
+			continue
+		}
+
+		return ErrInvalidPayload{
+			Type:      tlvType,
+			Violation: IncludedViolation,
+			FinalHop:  isFinalHop,
+		}
+	}
+
+	return nil
+}
diff --git a/htlcswitch/hop/payload_test.go b/htlcswitch/hop/payload_test.go
index 148b806f9..84e1e7d73 100644
--- a/htlcswitch/hop/payload_test.go
+++ b/htlcswitch/hop/payload_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/lightningnetwork/lnd/htlcswitch/hop"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/record"
+	"github.com/lightningnetwork/lnd/tlv"
 	"github.com/stretchr/testify/require"
 )
 
@@ -695,3 +696,67 @@ func TestValidateBlindedRouteData(t *testing.T) {
 		})
 	}
 }
+
+// TestValidatePayloadWithBlinded tests validation of the contents of a
+// payload when it's for a blinded payment.
+func TestValidatePayloadWithBlinded(t *testing.T) {
+	t.Parallel()
+
+	finalHopMap := map[tlv.Type][]byte{
+		record.AmtOnionType:            nil,
+		record.LockTimeOnionType:       nil,
+		record.TotalAmtMsatBlindedType: nil,
+	}
+
+	tests := []struct {
+		name    string
+		isFinal bool
+		parsed  map[tlv.Type][]byte
+		err     bool
+	}{
+		{
+			name:    "final hop, valid",
+			isFinal: true,
+			parsed:  finalHopMap,
+		},
+		{
+			name:    "intermediate hop, invalid",
+			isFinal: false,
+			parsed:  finalHopMap,
+			err:     true,
+		},
+		{
+			name:    "intermediate hop, invalid",
+			isFinal: false,
+			parsed: map[tlv.Type][]byte{
+				record.EncryptedDataOnionType: nil,
+				record.BlindingPointOnionType: nil,
+			},
+		},
+		{
+			name:    "unknown record, invalid",
+			isFinal: false,
+			parsed: map[tlv.Type][]byte{
+				tlv.Type(99): nil,
+			},
+			err: true,
+		},
+	}
+
+	for _, testCase := range tests {
+		t.Run(testCase.name, func(t *testing.T) {
+			err := hop.ValidatePayloadWithBlinded(
+				testCase.isFinal, testCase.parsed,
+			)
+
+			// We can't determine our exact error because we
+			// iterate through a map (non-deterministic) in the
+			// function.
+			if testCase.err {
+				require.NotNil(t, err)
+			} else {
+				require.Nil(t, err)
+			}
+		})
+	}
+}