github.com/letsencrypt/trillian@v1.1.2-0.20180615153820-ae375a99d36a/merkle/compact_merkle_tree_test.go (about)

     1  // Copyright 2016 Google Inc. All Rights Reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package merkle
    16  
    17  import (
    18  	"bytes"
    19  	"encoding/base64"
    20  	"errors"
    21  	"fmt"
    22  	"math/bits"
    23  	"strings"
    24  	"testing"
    25  
    26  	"github.com/google/trillian/merkle/rfc6962"
    27  	"github.com/google/trillian/storage"
    28  	"github.com/google/trillian/testonly"
    29  	"github.com/kylelemons/godebug/pretty"
    30  )
    31  
    32  func checkUnusedNodesInvariant(c *CompactMerkleTree) error {
    33  	// The structure of this invariant check mirrors the structure in
    34  	// NewCompactMerkleTreeWithState in which only the nodes which
    35  	// should be present for a tree of given size are fetched from the
    36  	// backing store via GetNodeFunc.
    37  	size := c.size
    38  	sizeBits := bits.Len64(uint64(size))
    39  	if isPerfectTree(size) {
    40  		for i, n := range c.nodes {
    41  			expectNil := i != sizeBits-1
    42  			if expectNil && n != nil {
    43  				return fmt.Errorf("perfect Tree size %d has non-nil node at index %d, wanted nil", size, i)
    44  			}
    45  			if !expectNil && n == nil {
    46  				return fmt.Errorf("perfect Tree size %d has nil node at index %d, wanted non-nil", size, i)
    47  			}
    48  		}
    49  	} else {
    50  		for depth := 0; depth < sizeBits; depth++ {
    51  			if size&1 == 1 {
    52  				if c.nodes[depth] == nil {
    53  					return fmt.Errorf("imperfect Tree size %d has nil node at index %d, wanted non-nil", c.size, depth)
    54  				}
    55  			} else {
    56  				if c.nodes[depth] != nil {
    57  					return fmt.Errorf("imperfect Tree size %d has non-nil node at index %d, wanted nil", c.size, depth)
    58  				}
    59  			}
    60  			size >>= 1
    61  		}
    62  	}
    63  	return nil
    64  }
    65  
    66  func TestAddingLeaves(t *testing.T) {
    67  	inputs := testonly.MerkleTreeLeafTestInputs()
    68  	roots := testonly.MerkleTreeLeafTestRootHashes()
    69  	hashes := testonly.CompactMerkleTreeLeafTestNodeHashes()
    70  
    71  	// We test the "same" thing 3 different ways this is to ensure than any lazy
    72  	// update strategy being employed by the implementation doesn't affect the
    73  	// api-visible calculation of root & size.
    74  	{
    75  		// First tree, add nodes one-by-one
    76  		tree := NewCompactMerkleTree(rfc6962.DefaultHasher)
    77  		if got, want := tree.Size(), int64(0); got != want {
    78  			t.Errorf("Size()=%d, want %d", got, want)
    79  		}
    80  		if got, want := tree.CurrentRoot(), testonly.EmptyMerkleTreeRootHash(); !bytes.Equal(got, want) {
    81  			t.Errorf("CurrentRoot()=%x, want %x", got, want)
    82  		}
    83  
    84  		for i := 0; i < 8; i++ {
    85  			tree.AddLeaf(inputs[i], func(int, int64, []byte) error {
    86  				return nil
    87  			})
    88  			if err := checkUnusedNodesInvariant(tree); err != nil {
    89  				t.Fatalf("UnusedNodesInvariant check failed: %v", err)
    90  			}
    91  			if got, want := tree.Size(), int64(i+1); got != want {
    92  				t.Errorf("Size()=%d, want %d", got, want)
    93  			}
    94  			if got, want := tree.CurrentRoot(), roots[i]; !bytes.Equal(got, want) {
    95  				t.Errorf("CurrentRoot()=%v, want %v", got, want)
    96  			}
    97  			if diff := pretty.Compare(tree.Hashes(), hashes[i]); diff != "" {
    98  				t.Errorf("post-Hashes() diff:\n%v", diff)
    99  			}
   100  		}
   101  	}
   102  
   103  	{
   104  		// Second tree, add nodes all at once
   105  		tree := NewCompactMerkleTree(rfc6962.DefaultHasher)
   106  		for i := 0; i < 8; i++ {
   107  			tree.AddLeaf(inputs[i], func(int, int64, []byte) error {
   108  				return nil
   109  			})
   110  			if err := checkUnusedNodesInvariant(tree); err != nil {
   111  				t.Fatalf("UnusedNodesInvariant check failed: %v", err)
   112  			}
   113  		}
   114  		if got, want := tree.Size(), int64(8); got != want {
   115  			t.Errorf("Size()=%d, want %d", got, want)
   116  		}
   117  		if got, want := tree.CurrentRoot(), roots[7]; !bytes.Equal(got, want) {
   118  			t.Errorf("CurrentRoot()=%v, want %v", got, want)
   119  		}
   120  		if diff := pretty.Compare(tree.Hashes(), hashes[7]); diff != "" {
   121  			t.Errorf("post-Hashes() diff:\n%v", diff)
   122  		}
   123  	}
   124  
   125  	{
   126  		// Third tree, add nodes in two chunks
   127  		tree := NewCompactMerkleTree(rfc6962.DefaultHasher)
   128  		for i := 0; i < 3; i++ {
   129  			tree.AddLeaf(inputs[i], func(int, int64, []byte) error {
   130  				return nil
   131  			})
   132  			if err := checkUnusedNodesInvariant(tree); err != nil {
   133  				t.Fatalf("UnusedNodesInvariant check failed: %v", err)
   134  			}
   135  		}
   136  		if got, want := tree.Size(), int64(3); got != want {
   137  			t.Errorf("Size()=%d, want %d", got, want)
   138  		}
   139  		if got, want := tree.CurrentRoot(), roots[2]; !bytes.Equal(got, want) {
   140  			t.Errorf("CurrentRoot()=%v, want %v", got, want)
   141  		}
   142  		if diff := pretty.Compare(tree.Hashes(), hashes[2]); diff != "" {
   143  			t.Errorf("post-Hashes() diff:\n%v", diff)
   144  		}
   145  
   146  		for i := 3; i < 8; i++ {
   147  			tree.AddLeaf(inputs[i], func(int, int64, []byte) error {
   148  				return nil
   149  			})
   150  			if err := checkUnusedNodesInvariant(tree); err != nil {
   151  				t.Fatalf("UnusedNodesInvariant check failed: %v", err)
   152  			}
   153  		}
   154  		if got, want := tree.Size(), int64(8); got != want {
   155  			t.Errorf("Size()=%d, want %d", got, want)
   156  		}
   157  		if got, want := tree.CurrentRoot(), roots[7]; !bytes.Equal(got, want) {
   158  			t.Errorf("CurrentRoot()=%v, want %v", got, want)
   159  		}
   160  		if diff := pretty.Compare(tree.Hashes(), hashes[7]); diff != "" {
   161  			t.Errorf("post-Hashes() diff:\n%v", diff)
   162  		}
   163  	}
   164  }
   165  
   166  func failingGetNodeFunc(int, int64) ([]byte, error) {
   167  	return []byte{}, errors.New("bang")
   168  }
   169  
   170  // This returns something that won't result in a valid root hash match, doesn't really
   171  // matter what it is but it must be correct length for an SHA256 hash as if it was real
   172  func fixedHashGetNodeFunc(int, int64) ([]byte, error) {
   173  	return []byte("12345678901234567890123456789012"), nil
   174  }
   175  
   176  func TestLoadingTreeFailsNodeFetch(t *testing.T) {
   177  	_, err := NewCompactMerkleTreeWithState(rfc6962.DefaultHasher, 237, failingGetNodeFunc, []byte("notimportant"))
   178  
   179  	if err == nil || !strings.Contains(err.Error(), "bang") {
   180  		t.Errorf("Did not return correctly on failed node fetch: %v", err)
   181  	}
   182  }
   183  
   184  func TestLoadingTreeFailsBadRootHash(t *testing.T) {
   185  	// Supply a root hash that can't possibly match the result of the SHA 256 hashing on our dummy
   186  	// data
   187  	_, err := NewCompactMerkleTreeWithState(rfc6962.DefaultHasher, 237, fixedHashGetNodeFunc, []byte("nomatch!nomatch!nomatch!nomatch!"))
   188  	_, ok := err.(RootHashMismatchError)
   189  
   190  	if err == nil || !ok {
   191  		t.Errorf("Did not return correct error type on root mismatch: %v", err)
   192  	}
   193  }
   194  
   195  func nodeKey(d int, i int64) (string, error) {
   196  	n, err := storage.NewNodeIDForTreeCoords(int64(d), i, 64)
   197  	if err != nil {
   198  		return "", err
   199  	}
   200  	return n.String(), nil
   201  }
   202  
   203  func TestCompactVsFullTree(t *testing.T) {
   204  	imt := NewInMemoryMerkleTree(rfc6962.DefaultHasher)
   205  	nodes := make(map[string][]byte)
   206  
   207  	for i := int64(0); i < 1024; i++ {
   208  		cmt, err := NewCompactMerkleTreeWithState(
   209  			rfc6962.DefaultHasher,
   210  			imt.LeafCount(),
   211  			func(depth int, index int64) ([]byte, error) {
   212  				k, err := nodeKey(depth, index)
   213  				if err != nil {
   214  					t.Errorf("failed to create nodeID: %v", err)
   215  				}
   216  				h := nodes[k]
   217  				return h, nil
   218  			}, imt.CurrentRoot().Hash())
   219  
   220  		if err != nil {
   221  			t.Errorf("interation %d: failed to create CMT with state: %v", i, err)
   222  		}
   223  		if a, b := imt.CurrentRoot().Hash(), cmt.CurrentRoot(); !bytes.Equal(a, b) {
   224  			t.Errorf("iteration %d: Got in-memory root of %v, but compact tree has root %v", i, a, b)
   225  		}
   226  
   227  		newLeaf := []byte(fmt.Sprintf("Leaf %d", i))
   228  
   229  		iSeq, iHash, err := imt.AddLeaf(newLeaf)
   230  		if err != nil {
   231  			t.Errorf("AddLeaf(): %v", err)
   232  		}
   233  
   234  		cSeq, cHash, err := cmt.AddLeaf(newLeaf,
   235  			func(depth int, index int64, hash []byte) error {
   236  				k, err := nodeKey(depth, index)
   237  				if err != nil {
   238  					return fmt.Errorf("failed to create nodeID: %v", err)
   239  				}
   240  				nodes[k] = hash
   241  				return nil
   242  			})
   243  		if err != nil {
   244  			t.Fatalf("mt update failed: %v", err)
   245  		}
   246  
   247  		// In-Memory tree is 1-based for sequence numbers, since it's based on the original CT C++ impl.
   248  		if got, want := iSeq, i+1; got != want {
   249  			t.Errorf("iteration %d: Got in-memory sequence number of %d, expected %d", i, got, want)
   250  		}
   251  		if int64(iSeq) != cSeq+1 {
   252  			t.Errorf("iteration %d: Got in-memory sequence number of %d but %d (zero based) from compact tree", i, iSeq, cSeq)
   253  		}
   254  		if a, b := iHash.Hash(), cHash; !bytes.Equal(a, b) {
   255  			t.Errorf("iteration %d: Got leaf hash %v from in-memory tree, but %v from compact tree", i, a, b)
   256  		}
   257  		if a, b := imt.CurrentRoot().Hash(), cmt.CurrentRoot(); !bytes.Equal(a, b) {
   258  			t.Errorf("iteration %d: Got in-memory root of %v, but compact tree has root %v", i, a, b)
   259  		}
   260  
   261  	}
   262  }
   263  
   264  func TestRootHashForVariousTreeSizes(t *testing.T) {
   265  	tests := []struct {
   266  		size     int64
   267  		wantRoot []byte
   268  	}{
   269  		{10, testonly.MustDecodeBase64("VjWMPSYNtCuCNlF/RLnQy6HcwSk6CIipfxm+hettA+4=")},
   270  		{15, testonly.MustDecodeBase64("j4SulYmocFuxdeyp12xXCIgK6PekBcxzAIj4zbQzNEI=")},
   271  		{16, testonly.MustDecodeBase64("c+4Uc6BCMOZf/v3NZK1kqTUJe+bBoFtOhP+P3SayKRE=")},
   272  		{100, testonly.MustDecodeBase64("dUh9hYH88p0CMoHkdr1wC2szbhcLAXOejWpINIooKUY=")},
   273  		{255, testonly.MustDecodeBase64("SmdsuKUqiod3RX2jyF2M6JnbdE4QuTwwipfAowI4/i0=")},
   274  		{256, testonly.MustDecodeBase64("qFI0t/tZ1MdOYgyPpPzHFiZVw86koScXy9q3FU5casA=")},
   275  		{1000, testonly.MustDecodeBase64("RXrgb8xHd55Y48FbfotJwCbV82Kx22LZfEbmBGAvwlQ=")},
   276  		{4095, testonly.MustDecodeBase64("cWRFdQhPcjn9WyBXE/r1f04ejxIm5lvg40DEpRBVS0w=")},
   277  		{4096, testonly.MustDecodeBase64("6uU/phfHg1n/GksYT6TO9aN8EauMCCJRl3dIK0HDs2M=")},
   278  		{10000, testonly.MustDecodeBase64("VZcav65F9haHVRk3wre2axFoBXRNeUh/1d9d5FQfxIg=")},
   279  		{65535, testonly.MustDecodeBase64("iPuVYJhP6SEE4gUFp8qbafd2rYv9YTCDYqAxCj8HdLM=")},
   280  	}
   281  
   282  	b64e := func(b []byte) string { return base64.StdEncoding.EncodeToString(b) }
   283  
   284  	for _, test := range tests {
   285  		tree := NewCompactMerkleTree(rfc6962.DefaultHasher)
   286  		for i := int64(0); i < test.size; i++ {
   287  			l := []byte{byte(i & 0xff), byte((i >> 8) & 0xff)}
   288  			tree.AddLeaf(l, func(int, int64, []byte) error {
   289  				return nil
   290  			})
   291  		}
   292  		if got, want := tree.CurrentRoot(), test.wantRoot; !bytes.Equal(got, want) {
   293  			t.Errorf("Test (treesize=%v) got root %v, want %v", test.size, b64e(got), b64e(want))
   294  		}
   295  	}
   296  }