From c1212d13692926dc509ba9e58e0de99b0bc50ed2 Mon Sep 17 00:00:00 2001
From: Tadge Dryja <tadge@lightning.network>
Date: Mon, 18 Jan 2016 21:26:28 -0800
Subject: [PATCH] off by 1 fix, a little more serdes check

---
 elkrem/elkrem.go      |  2 +-
 elkrem/serdes.go      | 12 ++++++++----
 elkrem/serdes_test.go | 33 ++++++++++++++++++++++++++++++++-
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/elkrem/elkrem.go b/elkrem/elkrem.go
index 83d4a58a..de54ec04 100644
--- a/elkrem/elkrem.go
+++ b/elkrem/elkrem.go
@@ -123,7 +123,7 @@ func (e *ElkremReceiver) AddNext(sha *wire.ShaHash) error {
 	var n ElkremNode
 	n.sha = sha
 	t := len(e.s) - 1 // top of stack
-	if t > 0 {        // if this is not the first hash
+	if t >= 0 {       // if this is not the first hash (>= because we -1'd)
 		n.i = e.s[t].i + 1 // incoming index is tip of stack index + 1
 	}
 	if t > 0 && e.s[t-1].h == e.s[t].h { // top 2 elements are equal height
diff --git a/elkrem/serdes.go b/elkrem/serdes.go
index 84054d6b..5212e2eb 100644
--- a/elkrem/serdes.go
+++ b/elkrem/serdes.go
@@ -45,7 +45,11 @@ func (e *ElkremSender) ToBytes() ([]byte, error) {
 // the index where it left off.
 func ElkremSenderFromBytes(b []byte) (ElkremSender, error) {
 	var e ElkremSender
+	e.root = new(wire.ShaHash)
 	buf := bytes.NewBuffer(b)
+	if buf.Len() != 41 {
+		return e, fmt.Errorf("Got %d bytes for sender, expect 41")
+	}
 	// read 1 byte height
 	err := binary.Read(buf, binary.BigEndian, &e.treeHeight)
 	if err != nil {
@@ -139,7 +143,6 @@ func ElkremReceiverFromBytes(b []byte) (ElkremReceiver, error) {
 	if e.treeHeight < 1 || e.treeHeight > 63 {
 		return e, fmt.Errorf("Read invalid receiver height: %d", e.treeHeight)
 	}
-
 	var max uint64 // maximum possible given height
 	for j := uint8(0); j <= e.treeHeight; j++ {
 		max = max<<1 | 1
@@ -154,9 +157,10 @@ func ElkremReceiverFromBytes(b []byte) (ElkremReceiver, error) {
 	if numOfNodes < 1 || numOfNodes > 64 {
 		return e, fmt.Errorf("Read invalid number of nodes: %d", numOfNodes)
 	}
-	if buf.Len() != (int(numOfNodes)*41)+2 {
-		return e, fmt.Errorf("Input buf wrong size, expect %d got %d",
-			(numOfNodes*41)+2, buf.Len())
+
+	if buf.Len() != (int(numOfNodes) * 41) {
+		return e, fmt.Errorf("Remaining buf wrong size, expect %d got %d",
+			(numOfNodes * 41), buf.Len())
 	}
 
 	for i := uint8(0); i < numOfNodes; i++ {
diff --git a/elkrem/serdes_test.go b/elkrem/serdes_test.go
index cb3e4173..8dab5e33 100644
--- a/elkrem/serdes_test.go
+++ b/elkrem/serdes_test.go
@@ -1,6 +1,9 @@
 package elkrem
 
-import "testing"
+import (
+	"bytes"
+	"testing"
+)
 
 func ReceiverSerdesTest(t *testing.T, er ElkremReceiver) {
 	b, err := er.ToBytes()
@@ -8,6 +11,20 @@ func ReceiverSerdesTest(t *testing.T, er ElkremReceiver) {
 		t.Fatal(err)
 	}
 	t.Logf("Serialized receiver; %d bytes, hex:\n%x\n", len(b), b)
+
+	rcv2, err := ElkremReceiverFromBytes(b)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b2, err := rcv2.ToBytes()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !bytes.Equal(b, b2) {
+		t.Fatalf("First and second serializations different")
+	}
 }
 
 func SenderSerdesTest(t *testing.T, es ElkremSender) {
@@ -16,4 +33,18 @@ func SenderSerdesTest(t *testing.T, es ElkremSender) {
 		t.Fatal(err)
 	}
 	t.Logf("Serialized sender; %d bytes, hex:\n%x\n", len(b), b)
+
+	sndr2, err := ElkremSenderFromBytes(b)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b2, err := sndr2.ToBytes()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !bytes.Equal(b, b2) {
+		t.Fatalf("First and second serializations different")
+	}
 }