From 86b7d92520d049c681902bc1c7b81581f79ec37e Mon Sep 17 00:00:00 2001
From: Paul Hadfield <hadfieldp@google.com>
Date: Mon, 3 Oct 2016 13:55:28 +0100
Subject: [PATCH] Add CompactMerkleTree tests checking the root hash at varying
 tree sizes.

Test data generated with the code to write nil in 'nodes []trillian.Hash'
disabled, so that it is known to be good with the existing CompactMerkleTree.

With the new 'write nil' code enabled this acts as a regression check.
---
 merkle/compact_merkle_tree.go      |  2 ++
 merkle/compact_merkle_tree_test.go | 33 ++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/merkle/compact_merkle_tree.go b/merkle/compact_merkle_tree.go
index e38d054c73..d9af06c1a9 100644
--- a/merkle/compact_merkle_tree.go
+++ b/merkle/compact_merkle_tree.go
@@ -199,6 +199,8 @@ func (c *CompactMerkleTree) AddLeafHash(leafHash trillian.Hash, f setNodeFunc) (
 		hash = c.hasher.HashChildren(c.nodes[bit], hash)
 		// Store the resulting parent hash.
 		f(bit+1, index, hash)
+		// Now, clear this position in the nodes list as the hash it formerly contained will be propogated upwards.
+		c.nodes[bit] = nil
 		// Figure out if we're done:
 		if bit+1 >= len(c.nodes) {
 			// If we're extending the node list then add a new entry with our
diff --git a/merkle/compact_merkle_tree_test.go b/merkle/compact_merkle_tree_test.go
index 4ebd2cc9a1..39e38d3ad9 100644
--- a/merkle/compact_merkle_tree_test.go
+++ b/merkle/compact_merkle_tree_test.go
@@ -2,6 +2,7 @@ package merkle
 
 import (
 	"bytes"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"reflect"
@@ -187,3 +188,35 @@ func TestCompactVsFullTree(t *testing.T) {
 
 	}
 }
+
+func TestRootHashForVariousTreeSizes(t *testing.T) {
+	tests := []struct{
+		size int64
+		wantRoot trillian.Hash
+	}{
+		{10, testonly.MustDecodeBase64("VjWMPSYNtCuCNlF/RLnQy6HcwSk6CIipfxm+hettA+4=")},
+		{15, testonly.MustDecodeBase64("j4SulYmocFuxdeyp12xXCIgK6PekBcxzAIj4zbQzNEI=")},
+		{16, testonly.MustDecodeBase64("c+4Uc6BCMOZf/v3NZK1kqTUJe+bBoFtOhP+P3SayKRE=")},
+		{100, testonly.MustDecodeBase64("dUh9hYH88p0CMoHkdr1wC2szbhcLAXOejWpINIooKUY=")},
+		{255, testonly.MustDecodeBase64("SmdsuKUqiod3RX2jyF2M6JnbdE4QuTwwipfAowI4/i0=")},
+		{256, testonly.MustDecodeBase64("qFI0t/tZ1MdOYgyPpPzHFiZVw86koScXy9q3FU5casA=")},
+		{1000, testonly.MustDecodeBase64("RXrgb8xHd55Y48FbfotJwCbV82Kx22LZfEbmBGAvwlQ=")},
+		{4095, testonly.MustDecodeBase64("cWRFdQhPcjn9WyBXE/r1f04ejxIm5lvg40DEpRBVS0w=")},
+		{4096, testonly.MustDecodeBase64("6uU/phfHg1n/GksYT6TO9aN8EauMCCJRl3dIK0HDs2M=")},
+		{10000, testonly.MustDecodeBase64("VZcav65F9haHVRk3wre2axFoBXRNeUh/1d9d5FQfxIg=")},
+		{65535, testonly.MustDecodeBase64("iPuVYJhP6SEE4gUFp8qbafd2rYv9YTCDYqAxCj8HdLM=")},
+	}
+
+	b64e := func(b []byte) string { return base64.StdEncoding.EncodeToString(b) }
+
+	for _, test := range tests {
+		tree := NewCompactMerkleTree(NewRFC6962TreeHasher(trillian.NewSHA256()))
+		for i := int64(0); i < test.size; i++ {
+			l := []byte{ byte(i & 0xff), byte((i >> 8) & 0xff) }
+			tree.AddLeaf(l, func(int, int64, trillian.Hash) {})
+		}
+		if got, want := tree.CurrentRoot(), test.wantRoot; !bytes.Equal(got, want) {
+			t.Errorf("Test (treesize=%v) got root %v, want %v", test.size, b64e(got), b64e(want))
+		}
+	}
+}
\ No newline at end of file