multi: decode zero-length onion message payloads

Since the onion message payload can be zero-length, we need to decode it
correctly. This commit adds a boolean flag to the HopPayload Decode that
tells whether the payload is an onion message payload or not. If it is,
the payload is decoded as a tlv payload also if the first byte is 0x00.

sphinx_test: Add zero-length payload om test

Author: gijswijs

payload.go

@@ -44,6 +44,10 @@ type HopPayload struct {
 	// HMAC is an HMAC computed over the entire per-hop payload that also
 	// includes the higher-level (optional) associated data bytes.
 	HMAC [HMACSize]byte
+
+	// TLVPayloadGuaranteed is set to true if the payload is guaranteed to
+	// be a TLVPayload. E.g. in the case of an onion message.
+	TLVPayloadGuaranteed bool
 }
 
 // NewTLVHopPayload creates a new TLV encoded HopPayload. The payload will be
@@ -99,36 +103,13 @@ func (hp *HopPayload) Decode(r io.Reader) error {
 		return err
 	}
 
-	var (
-		legacyPayload = isLegacyPayloadByte(peekByte[0])
-		payloadSize   uint16
-	)
-
-	if legacyPayload {
-		payloadSize = legacyPayloadSize()
-		hp.Type = PayloadLegacy
-	} else {
-		payloadSize, err = tlvPayloadSize(bufReader)
-		if err != nil {
-			return err
-		}
-
-		hp.Type = PayloadTLV
+	// If the HopPayload is guaranteed to be a TLV payload, we can skip the
+	// check for the legacy payload byte.
+	if !hp.TLVPayloadGuaranteed && isLegacyPayloadByte(peekByte[0]) {
+		return decodeLegacyHopPayload(hp, bufReader)
 	}
 
-	// Now that we know the payload size, we'll create a  new buffer to
-	// read it out in full.
-	//
-	// TODO(roasbeef): can avoid all these copies
-	hp.Payload = make([]byte, payloadSize)
-	if _, err := io.ReadFull(bufReader, hp.Payload[:]); err != nil {
-		return err
-	}
-	if _, err := io.ReadFull(bufReader, hp.HMAC[:]); err != nil {
-		return err
-	}
-
-	return nil
+	return decodeTLVHopPayload(hp, bufReader)
 }
 
 // HopData attempts to extract a set of forwarding instructions from the target
@@ -146,6 +127,42 @@ func (hp *HopPayload) HopData() (*HopData, error) {
 	return nil, nil
 }
 
+// readPayloadAndHMAC reads the payload and HMAC from the reader into the
+// HopPayload.
+func readPayloadAndHMAC(hp *HopPayload, r io.Reader, payloadSize uint16) error {
+	// Now that we know the payload size, we'll create a new buffer to read
+	// it out in full.
+	hp.Payload = make([]byte, payloadSize)
+	if _, err := io.ReadFull(r, hp.Payload[:]); err != nil {
+		return err
+	}
+	if _, err := io.ReadFull(r, hp.HMAC[:]); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// decodeTLVHopPayload decodes a TLV hop payload from the passed reader.
+func decodeTLVHopPayload(hp *HopPayload, r io.Reader) error {
+	payloadSize, err := tlvPayloadSize(r)
+	if err != nil {
+		return err
+	}
+
+	hp.Type = PayloadTLV
+
+	return readPayloadAndHMAC(hp, r, payloadSize)
+}
+
+// decodeLegacyHopPayload decodes a legacy hop payload from the passed reader.
+func decodeLegacyHopPayload(hp *HopPayload, r io.Reader) error {
+	payloadSize := legacyPayloadSize()
+	hp.Type = PayloadLegacy
+
+	return readPayloadAndHMAC(hp, r, payloadSize)
+}
+
 // tlvPayloadSize uses the passed reader to extract the payload length encoded
 // as a var-int.
 func tlvPayloadSize(r io.Reader) (uint16, error) {
@@ -314,8 +331,12 @@ func legacyNumBytes() int {
 	return LegacyHopDataSize
 }
 
-// isLegacyPayload returns true if the given byte is equal to the 0x00 byte
-// which indicates that the payload should be decoded as a legacy payload.
+// isLegacyPayloadByte determines if the first byte of a hop payload indicates
+// that it is a legacy payload. The first byte of a legacy payload will always
+// be 0x00, as this is the realm. For TLV payloads, the first byte is a
+// var-int encoding the length of the payload. A TLV stream can be empty, in
+// which case its length is 0, which is also encoded as a 0x00 byte. This
+// creates an ambiguity between a legacy payload and an empty TLV payload.
 func isLegacyPayloadByte(b byte) bool {
 	return b == 0x00
 }

sphinx.go

@@ -510,7 +510,8 @@ func (r *Router) Stop() {
 // processOnionCfg is a set of config values that can be used to modify how an
 // onion is processed.
 type processOnionCfg struct {
-	blindingPoint *btcec.PublicKey
+	blindingPoint  *btcec.PublicKey
+	isOnionMessage bool
 }
 
 // ProcessOnionOpt defines the signature of a function option that can be used
@@ -525,6 +526,14 @@ func WithBlindingPoint(point *btcec.PublicKey) ProcessOnionOpt {
 	}
 }
 
+// WithIsOnionMessage is a functional option that signals that the onion packet
+// being processed is from onion message.
+func WithIsOnionMessage() ProcessOnionOpt {
+	return func(cfg *processOnionCfg) {
+		cfg.isOnionMessage = true
+	}
+}
+
 // ProcessOnionPacket processes an incoming onion packet which has been forward
 // to the target Sphinx router. If the encoded ephemeral key isn't on the
 // target Elliptic Curve, then the packet is rejected. Similarly, if the
@@ -560,7 +569,9 @@ func (r *Router) ProcessOnionPacket(onionPkt *OnionPacket, assocData []byte,
 	// Continue to optimistically process this packet, deferring replay
 	// protection until the end to reduce the penalty of multiple IO
 	// operations.
-	packet, err := processOnionPacket(onionPkt, &sharedSecret, assocData)
+	packet, err := processOnionPacket(
+		onionPkt, &sharedSecret, assocData, cfg.isOnionMessage,
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -594,7 +605,9 @@ func (r *Router) ReconstructOnionPacket(onionPkt *OnionPacket, assocData []byte,
 		return nil, err
 	}
 
-	return processOnionPacket(onionPkt, &sharedSecret, assocData)
+	return processOnionPacket(
+		onionPkt, &sharedSecret, assocData, cfg.isOnionMessage,
+	)
 }
 
 // DecryptBlindedHopData uses the router's private key to decrypt data encrypted
@@ -625,7 +638,8 @@ func (r *Router) OnionPublicKey() *btcec.PublicKey {
 // packet. This function returns the next inner onion packet layer, along with
 // the hop data extracted from the outer onion packet.
 func unwrapPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
-	assocData []byte) (*OnionPacket, *HopPayload, error) {
+	assocData []byte, isOnionMessage bool) (*OnionPacket, *HopPayload,
+	error) {
 
 	dhKey := onionPkt.EphemeralKey
 	routeInfo := onionPkt.RoutingInfo
@@ -660,8 +674,9 @@ func unwrapPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
 	// With the MAC checked, and the payload decrypted, we can now parse
 	// out the payload so we can derive the specified forwarding
 	// instructions.
-	var hopPayload HopPayload
-	if err := hopPayload.Decode(bytes.NewReader(hopInfo[:])); err != nil {
+	hopPayload := HopPayload{TLVPayloadGuaranteed: isOnionMessage}
+	err := hopPayload.Decode(bytes.NewReader(hopInfo[:]))
+	if err != nil {
 		return nil, nil, err
 	}
 
@@ -683,7 +698,7 @@ func unwrapPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
 // packets. The processed packets returned from this method should only be used
 // if the packet was not flagged as a replayed packet.
 func processOnionPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
-	assocData []byte) (*ProcessedPacket, error) {
+	assocData []byte, isOnionMessage bool) (*ProcessedPacket, error) {
 
 	// First, we'll unwrap an initial layer of the onion packet. Typically,
 	// we'll only have a single layer to unwrap, However, if the sender has
@@ -693,7 +708,7 @@ func processOnionPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
 	// they can properly check the HMAC and unwrap a layer for their
 	// handoff hop.
 	innerPkt, outerHopPayload, err := unwrapPacket(
-		onionPkt, sharedSecret, assocData,
+		onionPkt, sharedSecret, assocData, isOnionMessage,
 	)
 	if err != nil {
 		return nil, err
@@ -703,7 +718,7 @@ func processOnionPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
 	// However if the uncovered 'nextMac' is all zeroes, then this
 	// indicates that we're the final hop in the route.
 	var action ProcessCode = MoreHops
-	if bytes.Compare(zeroHMAC[:], outerHopPayload.HMAC[:]) == 0 {
+	if bytes.Equal(zeroHMAC[:], outerHopPayload.HMAC[:]) {
 		action = ExitNode
 	}
 
@@ -794,7 +809,9 @@ func (t *Tx) ProcessOnionPacket(seqNum uint16, onionPkt *OnionPacket,
 	// Continue to optimistically process this packet, deferring replay
 	// protection until the end to reduce the penalty of multiple IO
 	// operations.
-	packet, err := processOnionPacket(onionPkt, &sharedSecret, assocData)
+	packet, err := processOnionPacket(
+		onionPkt, &sharedSecret, assocData, cfg.isOnionMessage,
+	)
 	if err != nil {
 		return err
 	}

sphinx_test.go

@@ -288,6 +288,60 @@ func TestTLVPayloadMessagePacket(t *testing.T) {
 		hex.EncodeToString(finalPacket), hex.EncodeToString(b.Bytes()))
 }
 
+// TestProcessOnionMessageZeroLengthPayload tests that we can properly process an
+// onion message that has a zero-length payload.
+func TestProcessOnionMessageZeroLengthPayload(t *testing.T) {
+	t.Parallel()
+
+	// First, create a router that will be the destination of the onion
+	// message.
+	privKey, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+
+	router := NewRouter(&PrivKeyECDH{privKey}, NewMemoryReplayLog())
+	err = router.Start()
+	require.NoError(t, err)
+	defer router.Stop()
+
+	// Next, create a session key for the onion packet.
+	sessionKey, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+
+	// We'll create a simple one-hop path.
+	path := &PaymentPath{
+		{
+			NodePub: *privKey.PubKey(),
+		},
+	}
+
+	// The hop payload will be an empty TLV payload.
+	payload, err := NewTLVHopPayload(nil)
+	require.NoError(t, err)
+	path[0].HopPayload = payload
+
+	// Now, create the onion packet.
+	onionPacket, err := NewOnionPacket(
+		path, sessionKey, nil, DeterministicPacketFiller,
+	)
+	require.NoError(t, err)
+
+	// We'll now process the packet, making sure to indicate that this is
+	// an onion message.
+	processedPacket, err := router.ProcessOnionPacket(
+		onionPacket, nil, 0, WithIsOnionMessage(),
+	)
+	require.NoError(t, err)
+
+	// The packet should be decoded as an exit node.
+	require.EqualValues(t, ExitNode, processedPacket.Action)
+
+	// The payload should be of type TLV.
+	require.Equal(t, PayloadTLV, processedPacket.Payload.Type)
+
+	// And the payload should be empty.
+	require.Empty(t, processedPacket.Payload.Payload)
+}
+
 func TestSphinxCorrectness(t *testing.T) {
 	nodes, _, hopDatas, fwdMsg, err := newTestRoute(testLegacyRouteNumHops)
 	if err != nil {