From c1e82e534d25cc09aba349d063f8703186d72d8e Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Thu, 8 Apr 2021 12:08:06 +0200
Subject: [PATCH] amp: add Merge and Zero to Sharer interface

---
 amp/derivation_test.go | 28 ++++++++++++++++++++++++++++
 amp/sharer.go          | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/amp/derivation_test.go b/amp/derivation_test.go
index 4ed533be..af8162d1 100644
--- a/amp/derivation_test.go
+++ b/amp/derivation_test.go
@@ -10,6 +10,7 @@ import (
 type sharerTest struct {
 	name      string
 	numShares int
+	merge     bool
 }
 
 var sharerTests = []sharerTest{
@@ -25,6 +26,16 @@ var sharerTests = []sharerTest{
 		name:      "many shares",
 		numShares: 10,
 	},
+	{
+		name:      "merge 4 shares",
+		numShares: 4,
+		merge:     true,
+	},
+	{
+		name:      "merge many shares",
+		numShares: 20,
+		merge:     true,
+	},
 }
 
 // TestSharer executes the end-to-end derivation between sender and receiver,
@@ -71,10 +82,27 @@ func testSharer(t *testing.T, test sharerTest) {
 
 	// Compute the final share and finalize the sharing.
 	child := sharer.Child(0)
+	sharer = sharer.Zero()
 
 	assertChildShare(t, child, 0)
 	children = append(children, child)
 
+	// If we are testing merging, merge half of the created children back
+	// into the sharer.
+	if test.merge {
+		for i := len(children) / 2; i < len(children); i++ {
+			sharer = sharer.Merge(children[i])
+		}
+		children = children[:len(children)/2]
+
+		// We must create a new last child from what we just merged
+		// back.
+		child := sharer.Child(0)
+
+		assertChildShare(t, child, 0)
+		children = append(children, child)
+	}
+
 	assertReconstruction(t, children...)
 }
 
diff --git a/amp/sharer.go b/amp/sharer.go
index de27550c..c8a8d059 100644
--- a/amp/sharer.go
+++ b/amp/sharer.go
@@ -2,8 +2,12 @@ package amp
 
 import (
 	"crypto/rand"
+	"fmt"
 )
 
+// zeroShare is the all-zero 32-byte share.
+var zeroShare = Share{}
+
 // Sharer facilitates dynamic splitting of a root share value and derivation of
 // child preimage and hashes for individual HTLCs in an AMP payment. A sharer
 // represents a specific node in an abstract binary tree that can generate up to
@@ -34,6 +38,16 @@ type Sharer interface {
 	// that the shares of all nodes descending from the parent will XOR to
 	// the parent's share.
 	Split() (Sharer, Sharer, error)
+
+	// Merge takes the given Child and "merges" it into the Sharer by
+	// XOR-ing its share with the Sharer's current share.
+	Merge(*Child) Sharer
+
+	// Zero returns a a new "zero Sharer" that has its current share set to
+	// zero, while keeping the root share. Merging a Child from the
+	// original Sharer into this zero-Sharer gives back the original
+	// Sharer.
+	Zero() Sharer
 }
 
 // SeedSharer orchestrates the sharing of the root AMP seed along multiple
@@ -81,6 +95,11 @@ func (s *SeedSharer) Root() Share {
 // parent share should no longer be used, and the caller should use the Child
 // method on each to derive preimage/hash pairs for the HTLCs.
 func (s *SeedSharer) Split() (Sharer, Sharer, error) {
+	// We cannot split the zero-Sharer.
+	if s.curr == zeroShare {
+		return nil, nil, fmt.Errorf("cannot split zero-Sharer")
+	}
+
 	shareLeft, shareRight, err := split(&s.curr)
 	if err != nil {
 		return nil, nil, err
@@ -92,6 +111,23 @@ func (s *SeedSharer) Split() (Sharer, Sharer, error) {
 	return left, right, nil
 }
 
+// Merge takes the given Child and "merges" it into the Sharer by XOR-ing its
+// share with the Sharer's current share.
+func (s *SeedSharer) Merge(child *Child) Sharer {
+	var curr Share
+	curr.Xor(&s.curr, &child.Share)
+
+	sharer := initSeedSharer(&s.root, &curr)
+	return sharer
+}
+
+// Zero returns a a new "zero Sharer" that has its current share set to zero,
+// while keeping the root share. Merging a Child from the original Sharer into
+// this zero-Sharer gives back the original Sharer.
+func (s *SeedSharer) Zero() Sharer {
+	return initSeedSharer(&s.root, &zeroShare)
+}
+
 // Child derives a preimage/hash pair to be used for an AMP HTLC.
 // All children of s will use the same underlying share, but have unique
 // preimage and hash. This can be used to rerandomize the preimage/hash pair for