From 36c713bbe3c9dbadd11e3f2bf1935c0b09febac3 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Wed, 28 Apr 2021 10:20:13 +0200
Subject: [PATCH] amp: add counter for child index

---
 amp/shard_tracker.go | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/amp/shard_tracker.go b/amp/shard_tracker.go
index 473447e4e41..cad50cd47e6 100644
--- a/amp/shard_tracker.go
+++ b/amp/shard_tracker.go
@@ -1,8 +1,6 @@
 package amp
 
 import (
-	"crypto/rand"
-	"encoding/binary"
 	"fmt"
 	"sync"
 
@@ -52,7 +50,8 @@ type ShardTracker struct {
 
 	sharer Sharer
 
-	shards map[uint64]*Child
+	shards     map[uint64]*Child
+	childIndex uint32
 	sync.Mutex
 }
 
@@ -90,19 +89,17 @@ func (s *ShardTracker) NewShard(pid uint64, last bool) (shards.PaymentShard,
 	s.Lock()
 	defer s.Unlock()
 
-	// Use a random child index.
-	var childIndex [4]byte
-	if _, err := rand.Read(childIndex[:]); err != nil {
-		return nil, err
-	}
-	idx := binary.BigEndian.Uint32(childIndex[:])
-
 	// Depending on whether we are requesting the last shard or not, either
 	// split the current share into two, or get a Child directly from the
 	// current sharer.
 	var child *Child
 	if last {
-		child = s.sharer.Child(idx)
+		// Use the current child index counter when generating the
+		// chilld, then increment it. This ensures we are creating
+		// unique hashes even when reusing the same share for the last
+		// shard.
+		child = s.sharer.Child(s.childIndex)
+		s.childIndex++
 
 		// If this was the last shard, set the current share to the
 		// zero share to indicate we cannot split it further.
@@ -114,7 +111,11 @@ func (s *ShardTracker) NewShard(pid uint64, last bool) (shards.PaymentShard,
 		}
 
 		s.sharer = sharer
-		child = left.Child(idx)
+
+		// Since splitting the current share added new 32-byte
+		// randomness to the shares, we reset the childIndex.
+		child = left.Child(0)
+		s.childIndex = 0
 	}
 
 	// Track the new child and return the shard.