github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/core/mpt/proof_test.go (about)

     1  package mpt
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/require"
     7  )
     8  
     9  func newProofTrie(t *testing.T, missingHashNode bool) *Trie {
    10  	l := NewLeafNode([]byte("somevalue"))
    11  	e := NewExtensionNode([]byte{0x05, 0x06, 0x07}, l)
    12  	l2 := NewLeafNode([]byte("invalid"))
    13  	e2 := NewExtensionNode([]byte{0x05}, NewHashNode(l2.Hash()))
    14  	b := NewBranchNode()
    15  	b.Children[4] = NewHashNode(e.Hash())
    16  	b.Children[5] = e2
    17  
    18  	tr := NewTrie(b, ModeAll, newTestStore())
    19  	require.NoError(t, tr.Put([]byte{0x12, 0x31}, []byte("value1")))
    20  	require.NoError(t, tr.Put([]byte{0x12, 0x32}, []byte("value2")))
    21  	tr.putToStore(l)
    22  	tr.putToStore(e)
    23  	if !missingHashNode {
    24  		tr.putToStore(l2)
    25  	}
    26  	return tr
    27  }
    28  
    29  func TestTrie_GetProof(t *testing.T) {
    30  	tr := newProofTrie(t, true)
    31  
    32  	t.Run("MissingKey", func(t *testing.T) {
    33  		_, err := tr.GetProof([]byte{0x12})
    34  		require.Error(t, err)
    35  	})
    36  
    37  	t.Run("Valid", func(t *testing.T) {
    38  		_, err := tr.GetProof([]byte{0x12, 0x31})
    39  		require.NoError(t, err)
    40  	})
    41  
    42  	t.Run("MissingHashNode", func(t *testing.T) {
    43  		_, err := tr.GetProof([]byte{0x55})
    44  		require.Error(t, err)
    45  	})
    46  }
    47  
    48  func TestVerifyProof(t *testing.T) {
    49  	tr := newProofTrie(t, true)
    50  
    51  	t.Run("Simple", func(t *testing.T) {
    52  		proof, err := tr.GetProof([]byte{0x12, 0x32})
    53  		require.NoError(t, err)
    54  
    55  		t.Run("Good", func(t *testing.T) {
    56  			v, ok := VerifyProof(tr.root.Hash(), []byte{0x12, 0x32}, proof)
    57  			require.True(t, ok)
    58  			require.Equal(t, []byte("value2"), v)
    59  		})
    60  
    61  		t.Run("Bad", func(t *testing.T) {
    62  			_, ok := VerifyProof(tr.root.Hash(), []byte{0x12, 0x31}, proof)
    63  			require.False(t, ok)
    64  		})
    65  	})
    66  
    67  	t.Run("InsideHash", func(t *testing.T) {
    68  		key := []byte{0x45, 0x67}
    69  		proof, err := tr.GetProof(key)
    70  		require.NoError(t, err)
    71  
    72  		v, ok := VerifyProof(tr.root.Hash(), key, proof)
    73  		require.True(t, ok)
    74  		require.Equal(t, []byte("somevalue"), v)
    75  	})
    76  }