From de88a4b17427eb489c57d316d3073e6cbc11597f Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Tue, 28 Jan 2020 06:43:07 -0800
Subject: [PATCH] record: add AMP record and encode/decode methods

---
 record/amp.go         | 107 ++++++++++++++++++++++++++++++++++++++++++
 record/record_test.go |  30 +++++++++++-
 2 files changed, 135 insertions(+), 2 deletions(-)
 create mode 100644 record/amp.go

diff --git a/record/amp.go b/record/amp.go
new file mode 100644
index 00000000..72b4cbf4
--- /dev/null
+++ b/record/amp.go
@@ -0,0 +1,107 @@
+package record
+
+import (
+	"fmt"
+	"io"
+
+	"github.com/lightningnetwork/lnd/tlv"
+)
+
+// AMPOnionType is the type used in the onion to reference the AMP fields:
+// root_share, set_id, and child_index.
+const AMPOnionType tlv.Type = 10
+
+// AMP is a record that encodes the fields necessary for atomic multi-path
+// payments.
+type AMP struct {
+	rootShare  [32]byte
+	setID      [32]byte
+	childIndex uint16
+}
+
+// NewAMP generate a new AMP record with the given root_share, set_id, and
+// child_index.
+func NewAMP(rootShare, setID [32]byte, childIndex uint16) *AMP {
+	return &AMP{
+		rootShare:  rootShare,
+		setID:      setID,
+		childIndex: childIndex,
+	}
+}
+
+// RootShare returns the root share contained in the AMP record.
+func (a *AMP) RootShare() [32]byte {
+	return a.rootShare
+}
+
+// SetID returns the set id contained in the AMP record.
+func (a *AMP) SetID() [32]byte {
+	return a.setID
+}
+
+// ChildIndex returns the child index contained in the AMP record.
+func (a *AMP) ChildIndex() uint16 {
+	return a.childIndex
+}
+
+// AMPEncoder writes the AMP record to the provided io.Writer.
+func AMPEncoder(w io.Writer, val interface{}, buf *[8]byte) error {
+	if v, ok := val.(*AMP); ok {
+		if err := tlv.EBytes32(w, &v.rootShare, buf); err != nil {
+			return err
+		}
+
+		if err := tlv.EBytes32(w, &v.setID, buf); err != nil {
+			return err
+		}
+
+		return tlv.ETUint16T(w, v.childIndex, buf)
+	}
+	return tlv.NewTypeForEncodingErr(val, "AMP")
+}
+
+const (
+	// minAMPLength is the minimum length of a serialized AMP TLV record,
+	// which occurs when the truncated encoding of child_index takes 0
+	// bytes, leaving only the root_share and set_id.
+	minAMPLength = 64
+
+	// maxAMPLength is the maximum legnth of a serialized AMP TLV record,
+	// which occurs when the truncated endoing of a child_index takes 2
+	// bytes.
+	maxAMPLength = 66
+)
+
+// AMPDecoder reads the AMP record from the provided io.Reader.
+func AMPDecoder(r io.Reader, val interface{}, buf *[8]byte, l uint64) error {
+	if v, ok := val.(*AMP); ok && minAMPLength <= l && l <= maxAMPLength {
+		if err := tlv.DBytes32(r, &v.rootShare, buf, 32); err != nil {
+			return err
+		}
+
+		if err := tlv.DBytes32(r, &v.setID, buf, 32); err != nil {
+			return err
+		}
+
+		return tlv.DTUint16(r, &v.childIndex, buf, l-64)
+	}
+	return tlv.NewTypeForDecodingErr(val, "AMP", l, maxAMPLength)
+}
+
+// Record returns a tlv.Record that can be used to encode or decode this record.
+func (a *AMP) Record() tlv.Record {
+	return tlv.MakeDynamicRecord(
+		AMPOnionType, a, a.PayloadSize, AMPEncoder, AMPDecoder,
+	)
+}
+
+// PayloadSize returns the size this record takes up in encoded form.
+func (a *AMP) PayloadSize() uint64 {
+	return 32 + 32 + tlv.SizeTUint16(a.childIndex)
+}
+
+// String returns a human-readble description of the amp payload fields.
+func (a *AMP) String() string {
+	return fmt.Sprintf("root_share=%x set_id=%x child_index=%d",
+		a.rootShare, a.setID, a.childIndex)
+}
diff --git a/record/record_test.go b/record/record_test.go
index 052e2f1f..8c39790e 100644
--- a/record/record_test.go
+++ b/record/record_test.go
@@ -17,8 +17,11 @@ type recordEncDecTest struct {
 }
 
 var (
-	testTotal = lnwire.MilliSatoshi(45)
-	testAddr  = [32]byte{0x01, 0x02}
+	testTotal      = lnwire.MilliSatoshi(45)
+	testAddr       = [32]byte{0x01, 0x02}
+	testShare      = [32]byte{0x03, 0x04}
+	testSetID      = [32]byte{0x05, 0x06}
+	testChildIndex = uint16(17)
 )
 
 var recordEncDecTests = []recordEncDecTest{
@@ -40,6 +43,29 @@ var recordEncDecTests = []recordEncDecTest{
 			}
 		},
 	},
+	{
+		name: "amp",
+		encRecord: func() tlv.RecordProducer {
+			return record.NewAMP(
+				testShare, testSetID, testChildIndex,
+			)
+		},
+		decRecord: func() tlv.RecordProducer {
+			return new(record.AMP)
+		},
+		assert: func(t *testing.T, r interface{}) {
+			amp := r.(*record.AMP)
+			if amp.RootShare() != testShare {
+				t.Fatal("incorrect root share")
+			}
+			if amp.SetID() != testSetID {
+				t.Fatal("incorrect set id")
+			}
+			if amp.ChildIndex() != testChildIndex {
+				t.Fatal("incorrect child index")
+			}
+		},
+	},
 }
 
 // TestRecordEncodeDecode is a generic test framework for custom TLV records. It