refactor: Add key length limit for trie kvstore

iotaledger · Jun 6, 2023 · 88f01d6 · 88f01d6
1 parent a27183a
commit 88f01d6
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/packages/trie/test/trie_test.go b/packages/trie/test/trie_test.go
@@ -117,6 +117,20 @@ func TestBasic3(t *testing.T) {
 	require.Equal(t, []byte{1}, tr.Get([]byte{0xb2}))
 }
 
+func TestKeyTooLong(t *testing.T) {
+	store := NewInMemoryKVStore()
+
+	tooLongKey := make([]byte, trie.KeyMaxLength+1)
+	root0 := trie.MustInitRoot(store)
+	{
+		tr, err := trie.NewTrieUpdatable(store, root0)
+		require.NoError(t, err)
+		require.Panics(t, func() {
+			tr.Update(tooLongKey, []byte{0})
+		})
+	}
+}
+
 func TestCreateTrie(t *testing.T) {
 	t.Run("ok init-"+"", func(t *testing.T) {
 		rootC1 := trie.MustInitRoot(NewInMemoryKVStore())

diff --git a/packages/trie/trie_kvstore.go b/packages/trie/trie_kvstore.go
@@ -5,6 +5,8 @@ import (
 	"encoding/hex"
 )
 
+var KeyMaxLength = 256
+
 // Update updates TrieUpdatable with the unpackedKey/value. Reorganizes and re-calculates trie, keeps cache consistent
 func (tr *TrieUpdatable) Update(key []byte, value []byte) {
 	assertf(len(key) > 0, "len(key) must be > 0")
@@ -105,7 +107,8 @@ func (tr *TrieReader) Iterator(prefix []byte) KVIterator {
 }
 
 func (tr *TrieUpdatable) update(triePath []byte, value []byte) {
-	assertf(len(value) > 0, "len(value)>0")
+	assertf(len(value) > 0, "len(value) > 0")
+	assertf(len(triePath) < KeyMaxLength, "len(key) = %d, must under KeyMaxLength", len(triePath))
 
 	nodes := make([]*bufferedNode, 0)
 	var ends pathEndingCode