Make in-memory Merkle tree immutable

integration/log.go

@@ -526,7 +526,7 @@ func buildMemoryMerkleTree(leafMap map[int64]*trillian.LogLeaf, params TestParam
 		if err := cr.Append(hasher.HashLeaf(leafMap[l].LeafValue), nil); err != nil {
 			return nil, err
 		}
-		merkleTree.AppendData(leafMap[l].LeafValue)
+		merkleTree = merkleTree.AppendData(leafMap[l].LeafValue)
 	}
 
 	// If the two reference results disagree there's no point in continuing the
@@ -543,7 +543,7 @@ func buildMemoryMerkleTree(leafMap map[int64]*trillian.LogLeaf, params TestParam
 		return nil, fmt.Errorf("different root hash results from merkle tree building: %v and %v", got, want)
 	}
 
-	return merkleTree, nil
+	return &merkleTree, nil
 }
 
 func getLatestSignedLogRoot(client trillian.TrillianLogClient, params TestParameters) (*trillian.GetLatestSignedLogRootResponse, error) {

internal/merkle/inmemory/tree.go

@@ -22,59 +22,87 @@ import (
 )
 
 // Tree implements an append-only Merkle tree. For testing.
+//
+// This type is immutable, but, if Append* methods are not used carefully, the
+// hashes can be corrupted. The semantics of memory reuse and reallocation is
+// similar to that of the append built-in for Go slices: overlapping Trees can
+// share memory.
+//
+// It is recommended to reuse old versions of Tree only for read operations, or
+// making sure that the newer Trees are no longer used before writing to an
+// older Tree. One scenario when rolling back to an older Tree can be useful is
+// implementing transaction semantics, rollbacks, and snapshots.
 type Tree struct {
 	hasher merkle.LogHasher
 	size   uint64
 	hashes [][][]byte // Node hashes, indexed by node (level, index).
 }
 
 // New returns a new empty Merkle tree.
-func New(hasher merkle.LogHasher) *Tree {
-	return &Tree{hasher: hasher}
+func New(hasher merkle.LogHasher) Tree {
+	return Tree{hasher: hasher}
+}
+
+// AppendData returns a new Tree which is the result of appending the hashes of
+// the given entries and the dependent Merkle tree nodes to the current tree.
+//
+// See Append method comment for details on safety.
+func (t Tree) AppendData(entries ...[]byte) Tree {
+	for _, entry := range entries {
+		t.appendImpl(t.hasher.HashLeaf(entry))
+	}
+	return t
 }
 
-// AppendData adds the leaf hash of the given entry to the end of the tree.
-func (t *Tree) AppendData(data []byte) {
-	t.Append(t.hasher.HashLeaf(data))
+// Append returns a new Tree which is the result of appending the given leaf
+// hashes and the dependent Merkle tree nodes to the current tree.
+//
+// The new Tree likely shares data with the old Tree, but in such a way that
+// both objects are valid. It is safe to reuse / roll back to the older Tree
+// objects, but Append should be called on them with caution because it may
+// corrupt hashes in the newer Tree objects.
+func (t Tree) Append(hashes ...[]byte) Tree {
+	for _, hash := range hashes {
+		t.appendImpl(hash)
+	}
+	return t
 }
 
-// Append adds the given leaf hash to the end of the tree.
-func (t *Tree) Append(hash []byte) {
-	level := 0
-	for ; (t.size>>level)&1 == 1; level++ {
-		row := append(t.hashes[level], hash)
-		hash = t.hasher.HashChildren(row[len(row)-2], hash)
-		t.hashes[level] = row
+func (t *Tree) appendImpl(hash []byte) {
+	level, width := 0, t.size
+	for ; width&1 == 1; width, level = width/2, level+1 {
+		t.hashes[level] = append(t.hashes[level][:width], hash)
+		hash = t.hasher.HashChildren(t.hashes[level][width-1], hash)
 	}
 	if level > len(t.hashes) {
 		panic("gap in tree appends")
 	} else if level == len(t.hashes) {
 		t.hashes = append(t.hashes, nil)
 	}
 
-	t.hashes[level] = append(t.hashes[level], hash)
+	t.hashes[level] = append(t.hashes[level][:width], hash)
 	t.size++
 }
 
 // Size returns the current number of leaves in the tree.
-func (t *Tree) Size() uint64 {
+func (t Tree) Size() uint64 {
 	return t.size
 }
 
 // LeafHash returns the leaf hash at the given index.
 // Requires 0 <= index < Size(), otherwise panics.
-func (t *Tree) LeafHash(index uint64) []byte {
+func (t Tree) LeafHash(index uint64) []byte {
 	return t.hashes[0][index]
 }
 
 // Hash returns the current root hash of the tree.
-func (t *Tree) Hash() []byte {
+func (t Tree) Hash() []byte {
 	return t.HashAt(t.size)
 }
 
 // HashAt returns the root hash at the given size.
 // Requires 0 <= size <= Size(), otherwise panics.
-func (t *Tree) HashAt(size uint64) []byte {
+func (t Tree) HashAt(size uint64) []byte {
 	if size == 0 {
 		return t.hasher.EmptyRoot()
 	}
@@ -90,7 +118,7 @@ func (t *Tree) HashAt(size uint64) []byte {
 // InclusionProof returns the inclusion proof for the given leaf index in the
 // tree of the given size. Requires 0 <= index < size <= Size(), otherwise may
 // panic.
-func (t *Tree) InclusionProof(index, size uint64) ([][]byte, error) {
+func (t Tree) InclusionProof(index, size uint64) ([][]byte, error) {
 	nodes, err := proof.Inclusion(index, size)
 	if err != nil {
 		return nil, err
@@ -100,15 +128,15 @@ func (t *Tree) InclusionProof(index, size uint64) ([][]byte, error) {
 
 // ConsistencyProof returns the consistency proof between the two given tree
 // sizes. Requires 0 <= size1 <= size2 <= Size(), otherwise may panic.
-func (t *Tree) ConsistencyProof(size1, size2 uint64) ([][]byte, error) {
+func (t Tree) ConsistencyProof(size1, size2 uint64) ([][]byte, error) {
 	nodes, err := proof.Consistency(size1, size2)
 	if err != nil {
 		return nil, err
 	}
 	return nodes.Rehash(t.getNodes(nodes.IDs), t.hasher.HashChildren)
 }
 
-func (t *Tree) getNodes(ids []compact.NodeID) [][]byte {
+func (t Tree) getNodes(ids []compact.NodeID) [][]byte {
 	hashes := make([][]byte, len(ids))
 	for i, id := range ids {
 		hashes[i] = t.hashes[id.Level][id.Index]

internal/merkle/inmemory/tree_test.go

@@ -35,9 +35,9 @@ var fuzzTestSize = int64(256)
 var emptyTreeHashValue = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
 
 // Inputs to the reference tree, which has eight leaves.
-var leafInputs = []string{
-	"", "00", "10", "2021", "3031", "40414243",
-	"5051525354555657", "606162636465666768696a6b6c6d6e6f",
+var leafInputs = [][]byte{
+	hx(""), hx("00"), hx("10"), hx("2021"), hx("3031"), hx("40414243"),
+	hx("5051525354555657"), hx("606162636465666768696a6b6c6d6e6f"),
 }
 
 // Incremental roots from building the reference tree from inputs leaf-by-leaf.
@@ -124,16 +124,16 @@ var testProofs = []proofTestVector{
 	}},
 }
 
-func decodeHexStringOrPanic(hs string) []byte {
+// hx decodes a hex string or panics.
+func hx(hs string) []byte {
 	data, err := hex.DecodeString(hs)
 	if err != nil {
 		panic(fmt.Errorf("failed to decode test data: %s", hs))
 	}
-
 	return data
 }
 
-func makeEmptyTree() *Tree {
+func makeEmptyTree() Tree {
 	return New(rfc6962.DefaultHasher)
 }
 
@@ -148,7 +148,7 @@ func makeFuzzTestData() [][]byte {
 	return data
 }
 
-func getRootAsString(mt *Tree, size uint64) string {
+func getRootAsString(mt Tree, size uint64) string {
 	if size > mt.Size() {
 		// Doesn't matter what this is as long as it could never be a valid
 		// hex encoding of a hash
@@ -328,7 +328,7 @@ func TestEmptyTreeHash(t *testing.T) {
 	}
 }
 
-func validateTree(mt *Tree, l uint64, t *testing.T) {
+func validateTree(mt Tree, l uint64, t *testing.T) {
 	if got, want := mt.Size(), l+1; got != want {
 		t.Errorf("Incorrect leaf count %d, expecting %d", got, want)
 	}
@@ -359,24 +359,19 @@ func TestBuildTreeBuildOneAtATime(t *testing.T) {
 
 	// Add to the tree, checking after each leaf
 	for l := uint64(0); l < 8; l++ {
-		mt.AppendData(decodeHexStringOrPanic(leafInputs[l]))
+		mt = mt.AppendData(leafInputs[l])
 		validateTree(mt, l, t)
 	}
 }
 
 func TestBuildTreeBuildAllAtOnce(t *testing.T) {
 	mt := makeEmptyTree()
-
-	for l := 0; l < 3; l++ {
-		mt.AppendData(decodeHexStringOrPanic(leafInputs[l]))
-	}
+	mt = mt.AppendData(leafInputs[:3]...)
 
 	// Check the intermediate state
 	validateTree(mt, 2, t)
 
-	for l := 3; l < 8; l++ {
-		mt.AppendData(decodeHexStringOrPanic(leafInputs[l]))
-	}
+	mt = mt.AppendData(leafInputs[3:8]...)
 
 	// Check the final state
 	validateTree(mt, 7, t)
@@ -386,9 +381,7 @@ func TestBuildTreeBuildTwoChunks(t *testing.T) {
 	mt := makeEmptyTree()
 
 	// Add to the tree, checking after each leaf
-	for l := 0; l < 8; l++ {
-		mt.AppendData(decodeHexStringOrPanic(leafInputs[l]))
-	}
+	mt = mt.AppendData(leafInputs[:8]...)
 
 	validateTree(mt, 7, t)
 }
@@ -409,13 +402,8 @@ func TestDownToPowerOfTwoSanity(t *testing.T) {
 }
 
 func TestReferenceMerklePathSanity(t *testing.T) {
-	var data [][]byte
-
 	mt := makeEmptyTree()
-
-	for s := 0; s < 8; s++ {
-		data = append(data, decodeHexStringOrPanic(leafInputs[s]))
-	}
+	data := append([][]byte{}, leafInputs[:8]...)
 
 	for _, path := range testPaths {
 		referencePath, err := referenceMerklePath(data[:path.snapshot], path.leaf, mt.hasher)
@@ -429,7 +417,7 @@ func TestReferenceMerklePathSanity(t *testing.T) {
 		}
 
 		for i := 0; i < len(path.testVector); i++ {
-			if !bytes.Equal(referencePath[i], decodeHexStringOrPanic(path.testVector[i])) {
+			if !bytes.Equal(referencePath[i], hx(path.testVector[i])) {
 				t.Errorf("Path mismatch: %s, %s", hex.EncodeToString(referencePath[i]),
 					path.testVector[i])
 			}
@@ -444,7 +432,7 @@ func TestMerkleTreeRootFuzz(t *testing.T) {
 		mt := makeEmptyTree()
 
 		for l := int64(0); l < treeSize; l++ {
-			mt.AppendData(data[l])
+			mt = mt.AppendData(data[l])
 		}
 
 		// Since the tree is evaluated lazily, the order of queries is significant.
@@ -475,7 +463,7 @@ func TestMerkleTreePathFuzz(t *testing.T) {
 		mt := makeEmptyTree()
 
 		for l := int64(0); l < treeSize; l++ {
-			mt.AppendData(data[l])
+			mt = mt.AppendData(data[l])
 		}
 
 		// Since the tree is evaluated lazily, the order of queries is significant.
@@ -518,7 +506,7 @@ func TestMerkleTreeConsistencyFuzz(t *testing.T) {
 		mt := makeEmptyTree()
 
 		for l := int64(0); l < treeSize; l++ {
-			mt.AppendData(data[l])
+			mt = mt.AppendData(data[l])
 		}
 
 		// Since the tree is evaluated lazily, the order of queries is significant.
@@ -557,17 +545,14 @@ func TestMerkleTreeConsistencyFuzz(t *testing.T) {
 func TestMerkleTreePathBuildOnce(t *testing.T) {
 	// First tree: build in one go.
 	mt := makeEmptyTree()
-
-	for i := 0; i < 8; i++ {
-		mt.AppendData(decodeHexStringOrPanic(leafInputs[i]))
-	}
+	mt = mt.AppendData(leafInputs[:8]...)
 
 	if size := mt.Size(); size != 8 {
 		t.Fatalf("8 leaves added but tree size is %d", size)
 	}
 
 	hash := mt.Hash()
-	if got, want := hash, decodeHexStringOrPanic(rootsAtSize[7]); !bytes.Equal(got, want) {
+	if got, want := hash, hx(rootsAtSize[7]); !bytes.Equal(got, want) {
 		t.Fatalf("Got unexpected root hash: %x %x", got, want)
 	}
 
@@ -589,7 +574,7 @@ func TestMerkleTreePathBuildOnce(t *testing.T) {
 		}
 
 		for j := range p2 {
-			if got, want := p1[j], decodeHexStringOrPanic(testPaths[i].testVector[j]); !bytes.Equal(got, want) {
+			if got, want := p1[j], hx(testPaths[i].testVector[j]); !bytes.Equal(got, want) {
 				t.Errorf("Path mismatch: got: %v want: %v", got, want)
 			}
 		}
@@ -600,15 +585,12 @@ func TestMerkleTreePathBuildIncrementally(t *testing.T) {
 	// Second tree: build incrementally.
 	// First tree: build in one go.
 	mt := makeEmptyTree()
-
-	for i := 0; i < 8; i++ {
-		mt.AppendData(decodeHexStringOrPanic(leafInputs[i]))
-	}
+	mt = mt.AppendData(leafInputs[:8]...)
 
 	mt2 := makeEmptyTree()
 
 	for i := uint64(0); i < 8; i++ {
-		mt2.AppendData(decodeHexStringOrPanic(leafInputs[i]))
+		mt2 = mt2.AppendData(leafInputs[i])
 
 		for j := uint64(0); j < i+1; j++ {
 			p1, err := mt.InclusionProof(j, i+1)
@@ -643,17 +625,14 @@ func TestMerkleTreePathBuildIncrementally(t *testing.T) {
 
 func TestProofConsistencyTestVectors(t *testing.T) {
 	mt := makeEmptyTree()
-
-	for i := 0; i < 8; i++ {
-		mt.AppendData(decodeHexStringOrPanic(leafInputs[i]))
-	}
+	mt = mt.AppendData(leafInputs[:8]...)
 
 	if size := mt.Size(); size != 8 {
 		t.Fatalf("8 leaves added but tree size is %d", size)
 	}
 
 	hash := mt.Hash()
-	if got, want := hash, decodeHexStringOrPanic(rootsAtSize[7]); !bytes.Equal(got, want) {
+	if got, want := hash, hx(rootsAtSize[7]); !bytes.Equal(got, want) {
 		t.Fatalf("Got unexpected root hash: %x %x", got, want)
 	}
 
@@ -671,7 +650,7 @@ func TestProofConsistencyTestVectors(t *testing.T) {
 		}
 
 		for j := 0; j < len(p2); j++ {
-			if got, want := p1[j], decodeHexStringOrPanic(testProofs[i].proof[j]); !bytes.Equal(got, want) {
+			if got, want := p1[j], hx(testProofs[i].proof[j]); !bytes.Equal(got, want) {
 				t.Errorf("Path mismatch: got: %v want: %v", got, want)
 			}
 		}