From 2029a06918f138064dedeff1019b9180e757a912 Mon Sep 17 00:00:00 2001
From: Carla Kirk-Cohen <kirkcohenc@gmail.com>
Date: Tue, 2 Apr 2024 10:59:57 -0400
Subject: [PATCH] multi: return parsed types from payload

To separate blinded route parsing from payload parsing, we need to
return the parsed types map so that we can properly validate blinded
data payloads against what we saw in the onion.
---
 htlcswitch/hop/fuzz_test.go    |  4 ++--
 htlcswitch/hop/iterator.go     |  4 +++-
 htlcswitch/hop/payload.go      | 17 ++++++++++-------
 htlcswitch/hop/payload_test.go |  2 +-
 4 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/htlcswitch/hop/fuzz_test.go b/htlcswitch/hop/fuzz_test.go
index 7d528a1c4..cbbe88260 100644
--- a/htlcswitch/hop/fuzz_test.go
+++ b/htlcswitch/hop/fuzz_test.go
@@ -117,7 +117,7 @@ func fuzzPayload(f *testing.F, finalPayload bool) {
 
 		r := bytes.NewReader(data)
 
-		payload1, err := NewPayloadFromReader(r, finalPayload)
+		payload1, _, err := NewPayloadFromReader(r, finalPayload)
 		if err != nil {
 			return
 		}
@@ -146,7 +146,7 @@ func fuzzPayload(f *testing.F, finalPayload bool) {
 			require.NoError(t, err)
 		}
 
-		payload2, err := NewPayloadFromReader(&b, finalPayload)
+		payload2, _, err := NewPayloadFromReader(&b, finalPayload)
 		require.NoError(t, err)
 
 		require.Equal(t, payload1, payload2)
diff --git a/htlcswitch/hop/iterator.go b/htlcswitch/hop/iterator.go
index 35fe5e44e..e14954723 100644
--- a/htlcswitch/hop/iterator.go
+++ b/htlcswitch/hop/iterator.go
@@ -106,11 +106,13 @@ func (r *sphinxHopIterator) HopPayload() (*Payload, error) {
 	// Otherwise, if this is the TLV payload, then we'll make a new stream
 	// to decode only what we need to make routing decisions.
 	case sphinx.PayloadTLV:
-		return NewPayloadFromReader(
+		payload, _, err := NewPayloadFromReader(
 			bytes.NewReader(r.processedPacket.Payload.Payload),
 			r.processedPacket.Action == sphinx.ExitNode,
 		)
 
+		return payload, err
+
 	default:
 		return nil, fmt.Errorf("unknown sphinx payload type: %v",
 			r.processedPacket.Payload.Type)
diff --git a/htlcswitch/hop/payload.go b/htlcswitch/hop/payload.go
index 99c36db6d..70fdb1403 100644
--- a/htlcswitch/hop/payload.go
+++ b/htlcswitch/hop/payload.go
@@ -133,11 +133,14 @@ func NewLegacyPayload(f *sphinx.HopData) *Payload {
 	}
 }
 
-// NewPayloadFromReader builds a new Hop from the passed io.Reader. The reader
+// NewPayloadFromReader builds a new Hop from the passed io.Reader and returns
+// a map of all the types that were found in the payload. The reader
 // should correspond to the bytes encapsulated in a TLV onion payload. The
 // final hop bool signals that this payload was the final packet parsed by
 // sphinx.
-func NewPayloadFromReader(r io.Reader, finalHop bool) (*Payload, error) {
+func NewPayloadFromReader(r io.Reader, finalHop bool) (*Payload,
+	map[tlv.Type][]byte, error) {
+
 	var (
 		cid           uint64
 		amt           uint64
@@ -162,27 +165,27 @@ func NewPayloadFromReader(r io.Reader, finalHop bool) (*Payload, error) {
 		record.NewTotalAmtMsatBlinded(&totalAmtMsat),
 	)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	// Since this data is provided by a potentially malicious peer, pass it
 	// into the P2P decoding variant.
 	parsedTypes, err := tlvStream.DecodeWithParsedTypesP2P(r)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	// Validate whether the sender properly included or omitted tlv records
 	// in accordance with BOLT 04.
 	err = ValidateParsedPayloadTypes(parsedTypes, finalHop)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	// Check for violation of the rules for mandatory fields.
 	violatingType := getMinRequiredViolation(parsedTypes)
 	if violatingType != nil {
-		return nil, ErrInvalidPayload{
+		return nil, nil, ErrInvalidPayload{
 			Type:      *violatingType,
 			Violation: RequiredViolation,
 			FinalHop:  finalHop,
@@ -229,7 +232,7 @@ func NewPayloadFromReader(r io.Reader, finalHop bool) (*Payload, error) {
 		blindingPoint: blindingPoint,
 		customRecords: customRecords,
 		totalAmtMsat:  lnwire.MilliSatoshi(totalAmtMsat),
-	}, nil
+	}, nil, nil
 }
 
 // ForwardingInfo returns the basic parameters required for HTLC forwarding,
diff --git a/htlcswitch/hop/payload_test.go b/htlcswitch/hop/payload_test.go
index 84e1e7d73..301e57716 100644
--- a/htlcswitch/hop/payload_test.go
+++ b/htlcswitch/hop/payload_test.go
@@ -479,7 +479,7 @@ func testDecodeHopPayloadValidation(t *testing.T, test decodePayloadTest) {
 		testChildIndex = uint32(9)
 	)
 
-	p, err := hop.NewPayloadFromReader(
+	p, _, err := hop.NewPayloadFromReader(
 		bytes.NewReader(test.payload), test.isFinalHop,
 	)
 	if !reflect.DeepEqual(test.expErr, err) {