github.com/prysmaticlabs/prysm@v1.4.4/shared/trieutil/sparse_merkle_test.go (about)

     1  package trieutil
     2  
     3  import (
     4  	"strconv"
     5  	"testing"
     6  
     7  	"github.com/ethereum/go-ethereum/accounts/abi/bind"
     8  	contracts "github.com/prysmaticlabs/prysm/contracts/deposit-contract"
     9  	ethpb "github.com/prysmaticlabs/prysm/proto/eth/v1alpha1"
    10  	"github.com/prysmaticlabs/prysm/shared/bytesutil"
    11  	"github.com/prysmaticlabs/prysm/shared/hashutil"
    12  	"github.com/prysmaticlabs/prysm/shared/params"
    13  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    14  )
    15  
    16  func TestMarshalDepositWithProof(t *testing.T) {
    17  	items := [][]byte{
    18  		[]byte("A"),
    19  		[]byte("BB"),
    20  		[]byte("CCC"),
    21  		[]byte("DDDD"),
    22  		[]byte("EEEEE"),
    23  		[]byte("FFFFFF"),
    24  		[]byte("GGGGGGG"),
    25  	}
    26  	m, err := GenerateTrieFromItems(items, params.BeaconConfig().DepositContractTreeDepth)
    27  	require.NoError(t, err)
    28  	proof, err := m.MerkleProof(2)
    29  	require.NoError(t, err)
    30  	require.Equal(t, len(proof), int(params.BeaconConfig().DepositContractTreeDepth)+1)
    31  	someRoot := [32]byte{1, 2, 3, 4}
    32  	someSig := [96]byte{1, 2, 3, 4}
    33  	someKey := [48]byte{1, 2, 3, 4}
    34  	dep := &ethpb.Deposit{
    35  		Proof: proof,
    36  		Data: &ethpb.Deposit_Data{
    37  			PublicKey:             someKey[:],
    38  			WithdrawalCredentials: someRoot[:],
    39  			Amount:                32,
    40  			Signature:             someSig[:],
    41  		},
    42  	}
    43  	enc, err := dep.MarshalSSZ()
    44  	require.NoError(t, err)
    45  	dec := &ethpb.Deposit{}
    46  	require.NoError(t, dec.UnmarshalSSZ(enc))
    47  	require.DeepEqual(t, dec, dep)
    48  }
    49  
    50  func TestMerkleTrie_MerkleProofOutOfRange(t *testing.T) {
    51  	h := hashutil.Hash([]byte("hi"))
    52  	m := &SparseMerkleTrie{
    53  		branches: [][][]byte{
    54  			{
    55  				h[:],
    56  			},
    57  			{
    58  				h[:],
    59  			},
    60  			{
    61  				[]byte{},
    62  			},
    63  		},
    64  		depth: 4,
    65  	}
    66  	if _, err := m.MerkleProof(6); err == nil {
    67  		t.Error("Expected out of range failure, received nil", err)
    68  	}
    69  }
    70  
    71  func TestMerkleTrieRoot_EmptyTrie(t *testing.T) {
    72  	trie, err := NewTrie(params.BeaconConfig().DepositContractTreeDepth)
    73  	require.NoError(t, err)
    74  	testAccount, err := contracts.Setup()
    75  	require.NoError(t, err)
    76  
    77  	depRoot, err := testAccount.Contract.GetDepositRoot(&bind.CallOpts{})
    78  	require.NoError(t, err)
    79  	require.DeepEqual(t, depRoot, trie.HashTreeRoot())
    80  }
    81  
    82  func TestGenerateTrieFromItems_NoItemsProvided(t *testing.T) {
    83  	if _, err := GenerateTrieFromItems(nil, params.BeaconConfig().DepositContractTreeDepth); err == nil {
    84  		t.Error("Expected error when providing nil items received nil")
    85  	}
    86  }
    87  
    88  func TestMerkleTrie_VerifyMerkleProof(t *testing.T) {
    89  	items := [][]byte{
    90  		[]byte("A"),
    91  		[]byte("B"),
    92  		[]byte("C"),
    93  		[]byte("D"),
    94  		[]byte("E"),
    95  		[]byte("F"),
    96  		[]byte("G"),
    97  		[]byte("H"),
    98  	}
    99  	m, err := GenerateTrieFromItems(items, params.BeaconConfig().DepositContractTreeDepth)
   100  	require.NoError(t, err)
   101  	proof, err := m.MerkleProof(0)
   102  	require.NoError(t, err)
   103  	require.Equal(t, int(params.BeaconConfig().DepositContractTreeDepth)+1, len(proof))
   104  	root := m.Root()
   105  	if ok := VerifyMerkleBranch(root[:], items[0], 0, proof, params.BeaconConfig().DepositContractTreeDepth); !ok {
   106  		t.Error("First Merkle proof did not verify")
   107  	}
   108  	proof, err = m.MerkleProof(3)
   109  	require.NoError(t, err)
   110  	require.Equal(t, true, VerifyMerkleBranch(root[:], items[3], 3, proof, params.BeaconConfig().DepositContractTreeDepth))
   111  	require.Equal(t, false, VerifyMerkleBranch(root[:], []byte("buzz"), 3, proof, params.BeaconConfig().DepositContractTreeDepth))
   112  }
   113  
   114  func TestMerkleTrie_VerifyMerkleProof_TrieUpdated(t *testing.T) {
   115  	items := [][]byte{
   116  		{1},
   117  		{2},
   118  		{3},
   119  		{4},
   120  	}
   121  	depth := params.BeaconConfig().DepositContractTreeDepth + 1
   122  	m, err := GenerateTrieFromItems(items, depth)
   123  	require.NoError(t, err)
   124  	proof, err := m.MerkleProof(0)
   125  	require.NoError(t, err)
   126  	root := m.Root()
   127  	require.Equal(t, true, VerifyMerkleBranch(root[:], items[0], 0, proof, depth))
   128  
   129  	// Now we update the trie.
   130  	m.Insert([]byte{5}, 3)
   131  	proof, err = m.MerkleProof(3)
   132  	require.NoError(t, err)
   133  	root = m.Root()
   134  	if ok := VerifyMerkleBranch(root[:], []byte{5}, 3, proof, depth); !ok {
   135  		t.Error("Second Merkle proof did not verify")
   136  	}
   137  	if ok := VerifyMerkleBranch(root[:], []byte{4}, 3, proof, depth); ok {
   138  		t.Error("Old item should not verify")
   139  	}
   140  
   141  	// Now we update the trie at an index larger than the number of items.
   142  	m.Insert([]byte{6}, 15)
   143  }
   144  
   145  func TestRoundtripProto_OK(t *testing.T) {
   146  	items := [][]byte{
   147  		{1},
   148  		{2},
   149  		{3},
   150  		{4},
   151  	}
   152  	m, err := GenerateTrieFromItems(items, params.BeaconConfig().DepositContractTreeDepth+1)
   153  	require.NoError(t, err)
   154  
   155  	protoTrie := m.ToProto()
   156  	depositRoot := m.HashTreeRoot()
   157  
   158  	newTrie := CreateTrieFromProto(protoTrie)
   159  	require.DeepEqual(t, depositRoot, newTrie.HashTreeRoot())
   160  
   161  }
   162  
   163  func TestCopy_OK(t *testing.T) {
   164  	items := [][]byte{
   165  		{1},
   166  		{2},
   167  		{3},
   168  		{4},
   169  	}
   170  	source, err := GenerateTrieFromItems(items, params.BeaconConfig().DepositContractTreeDepth+1)
   171  	require.NoError(t, err)
   172  	copiedTrie := source.Copy()
   173  
   174  	if copiedTrie == source {
   175  		t.Errorf("Original trie returned.")
   176  	}
   177  	copyHash := copiedTrie.HashTreeRoot()
   178  	require.DeepEqual(t, copyHash, copiedTrie.HashTreeRoot())
   179  }
   180  
   181  func BenchmarkGenerateTrieFromItems(b *testing.B) {
   182  	items := [][]byte{
   183  		[]byte("A"),
   184  		[]byte("BB"),
   185  		[]byte("CCC"),
   186  		[]byte("DDDD"),
   187  		[]byte("EEEEE"),
   188  		[]byte("FFFFFF"),
   189  		[]byte("GGGGGGG"),
   190  	}
   191  	for i := 0; i < b.N; i++ {
   192  		_, err := GenerateTrieFromItems(items, params.BeaconConfig().DepositContractTreeDepth)
   193  		require.NoError(b, err, "Could not generate Merkle trie from items")
   194  	}
   195  }
   196  
   197  func BenchmarkInsertTrie_Optimized(b *testing.B) {
   198  	b.StopTimer()
   199  	numDeposits := 16000
   200  	items := make([][]byte, numDeposits)
   201  	for i := 0; i < numDeposits; i++ {
   202  		someRoot := bytesutil.ToBytes32([]byte(strconv.Itoa(i)))
   203  		items[i] = someRoot[:]
   204  	}
   205  	tr, err := GenerateTrieFromItems(items, params.BeaconConfig().DepositContractTreeDepth)
   206  	require.NoError(b, err)
   207  
   208  	someItem := bytesutil.ToBytes32([]byte("hello-world"))
   209  	b.StartTimer()
   210  	for i := 0; i < b.N; i++ {
   211  		tr.Insert(someItem[:], i%numDeposits)
   212  	}
   213  }
   214  
   215  func BenchmarkGenerateProof(b *testing.B) {
   216  	b.StopTimer()
   217  	items := [][]byte{
   218  		[]byte("A"),
   219  		[]byte("BB"),
   220  		[]byte("CCC"),
   221  		[]byte("DDDD"),
   222  		[]byte("EEEEE"),
   223  		[]byte("FFFFFF"),
   224  		[]byte("GGGGGGG"),
   225  	}
   226  	normalTrie, err := GenerateTrieFromItems(items, params.BeaconConfig().DepositContractTreeDepth)
   227  	require.NoError(b, err)
   228  
   229  	b.StartTimer()
   230  	for i := 0; i < b.N; i++ {
   231  		_, err := normalTrie.MerkleProof(3)
   232  		require.NoError(b, err)
   233  	}
   234  }
   235  
   236  func BenchmarkVerifyMerkleBranch(b *testing.B) {
   237  	b.StopTimer()
   238  	items := [][]byte{
   239  		[]byte("A"),
   240  		[]byte("BB"),
   241  		[]byte("CCC"),
   242  		[]byte("DDDD"),
   243  		[]byte("EEEEE"),
   244  		[]byte("FFFFFF"),
   245  		[]byte("GGGGGGG"),
   246  	}
   247  	m, err := GenerateTrieFromItems(items, params.BeaconConfig().DepositContractTreeDepth)
   248  	require.NoError(b, err)
   249  	proof, err := m.MerkleProof(2)
   250  	require.NoError(b, err)
   251  
   252  	root := m.Root()
   253  	b.StartTimer()
   254  	for i := 0; i < b.N; i++ {
   255  		if ok := VerifyMerkleBranch(root[:], items[2], 2, proof, params.BeaconConfig().DepositContractTreeDepth); !ok {
   256  			b.Error("Merkle proof did not verify")
   257  		}
   258  	}
   259  }