From 138b0798f5b7ea4988dec3c34d18087c2fdddbf1 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Fri, 30 Aug 2019 14:11:57 -0700
Subject: [PATCH] htlcswitch/iterator: use hop.Payload constructors

---
 htlcswitch/hop/payload.go | 75 +++++++++++++++++++++++++++++++++++++++
 htlcswitch/iterator.go    | 43 +++-------------------
 2 files changed, 79 insertions(+), 39 deletions(-)
 create mode 100644 htlcswitch/hop/payload.go

diff --git a/htlcswitch/hop/payload.go b/htlcswitch/hop/payload.go
new file mode 100644
index 000000000..1c7e563be
--- /dev/null
+++ b/htlcswitch/hop/payload.go
@@ -0,0 +1,75 @@
+package hop
+
+import (
+	"encoding/binary"
+	"io"
+
+	"github.com/lightningnetwork/lightning-onion"
+	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/record"
+	"github.com/lightningnetwork/lnd/tlv"
+)
+
+// Payload encapsulates all information delivered to a hop in an onion payload.
+// A Hop can represent either a TLV or legacy payload. The primary forwarding
+// instruction can be accessed via ForwardingInfo, and additional records can be
+// accessed by other member functions.
+type Payload struct {
+	// FwdInfo holds the basic parameters required for HTLC forwarding, e.g.
+	// amount, cltv, and next hop.
+	FwdInfo ForwardingInfo
+}
+
+// NewLegacyPayload builds a Payload from the amount, cltv, and next hop
+// parameters provided by leegacy onion payloads.
+func NewLegacyPayload(f *sphinx.HopData) *Payload {
+	nextHop := binary.BigEndian.Uint64(f.NextAddress[:])
+
+	return &Payload{
+		FwdInfo: ForwardingInfo{
+			Network:         BitcoinNetwork,
+			NextHop:         lnwire.NewShortChanIDFromInt(nextHop),
+			AmountToForward: lnwire.MilliSatoshi(f.ForwardAmount),
+			OutgoingCTLV:    f.OutgoingCltv,
+		},
+	}
+}
+
+// NewPayloadFromReader builds a new Hop from the passed io.Reader. The reader
+// should correspond to the bytes encapsulated in a TLV onion payload.
+func NewPayloadFromReader(r io.Reader) (*Payload, error) {
+	var (
+		cid  uint64
+		amt  uint64
+		cltv uint32
+	)
+
+	tlvStream, err := tlv.NewStream(
+		record.NewAmtToFwdRecord(&amt),
+		record.NewLockTimeRecord(&cltv),
+		record.NewNextHopIDRecord(&cid),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	err = tlvStream.Decode(r)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Payload{
+		FwdInfo: ForwardingInfo{
+			Network:         BitcoinNetwork,
+			NextHop:         lnwire.NewShortChanIDFromInt(cid),
+			AmountToForward: lnwire.MilliSatoshi(amt),
+			OutgoingCTLV:    cltv,
+		},
+	}, nil
+}
+
+// ForwardingInfo returns the basic parameters required for HTLC forwarding,
+// e.g. amount, cltv, and next hop.
+func (h *Payload) ForwardingInfo() ForwardingInfo {
+	return h.FwdInfo
+}
diff --git a/htlcswitch/iterator.go b/htlcswitch/iterator.go
index 4e04508a5..9dfe3a1af 100644
--- a/htlcswitch/iterator.go
+++ b/htlcswitch/iterator.go
@@ -2,7 +2,6 @@ package htlcswitch
 
 import (
 	"bytes"
-	"encoding/binary"
 	"fmt"
 	"io"
 
@@ -10,8 +9,6 @@ import (
 	"github.com/lightningnetwork/lightning-onion"
 	"github.com/lightningnetwork/lnd/htlcswitch/hop"
 	"github.com/lightningnetwork/lnd/lnwire"
-	"github.com/lightningnetwork/lnd/record"
-	"github.com/lightningnetwork/lnd/tlv"
 )
 
 // HopIterator is an interface that abstracts away the routing information
@@ -85,64 +82,32 @@ func (r *sphinxHopIterator) EncodeNextHop(w io.Writer) error {
 func (r *sphinxHopIterator) ForwardingInstructions() (
 	hop.ForwardingInfo, error) {
 
-	var (
-		nextHop lnwire.ShortChannelID
-		amt     uint64
-		cltv    uint32
-	)
-
 	switch r.processedPacket.Payload.Type {
 	// If this is the legacy payload, then we'll extract the information
 	// directly from the pre-populated ForwardingInstructions field.
 	case sphinx.PayloadLegacy:
 		fwdInst := r.processedPacket.ForwardingInstructions
+		p := hop.NewLegacyPayload(fwdInst)
 
-		switch r.processedPacket.Action {
-		case sphinx.ExitNode:
-			nextHop = hop.Exit
-		case sphinx.MoreHops:
-			s := binary.BigEndian.Uint64(fwdInst.NextAddress[:])
-			nextHop = lnwire.NewShortChanIDFromInt(s)
-		}
-
-		amt = fwdInst.ForwardAmount
-		cltv = fwdInst.OutgoingCltv
+		return p.ForwardingInfo(), nil
 
 	// 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:
-		var cid uint64
-
-		tlvStream, err := tlv.NewStream(
-			record.NewAmtToFwdRecord(&amt),
-			record.NewLockTimeRecord(&cltv),
-			record.NewNextHopIDRecord(&cid),
-		)
-		if err != nil {
-			return hop.ForwardingInfo{}, err
-		}
-
-		err = tlvStream.Decode(bytes.NewReader(
+		p, err := hop.NewPayloadFromReader(bytes.NewReader(
 			r.processedPacket.Payload.Payload,
 		))
 		if err != nil {
 			return hop.ForwardingInfo{}, err
 		}
 
-		nextHop = lnwire.NewShortChanIDFromInt(cid)
+		return p.ForwardingInfo(), nil
 
 	default:
 		return hop.ForwardingInfo{}, fmt.Errorf("unknown "+
 			"sphinx payload type: %v",
 			r.processedPacket.Payload.Type)
 	}
-
-	return hop.ForwardingInfo{
-		Network:         hop.BitcoinNetwork,
-		NextHop:         nextHop,
-		AmountToForward: lnwire.MilliSatoshi(amt),
-		OutgoingCTLV:    cltv,
-	}, nil
 }
 
 // ExtraOnionBlob returns the additional EOB data (if available).