fix(lib/runtime): ext_default_child_storage_next_key_version_1 retu…

…rn `None` correctly (#3473)
ChainSafe · Sep 11, 2023 · c7d574b · c7d574b
1 parent 80558c4
commit c7d574b
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 26 deletions.
diff --git a/lib/runtime/wazero/imports.go b/lib/runtime/wazero/imports.go
@@ -363,7 +363,7 @@ func ext_crypto_secp256k1_ecdsa_recover_version_1(ctx context.Context, m api.Mod
 	return ret
 }
 
-func ext_crypto_secp256k1_ecdsa_recover_version_2(ctx context.Context, m api.Module, sig uint32, msg uint32) uint64 {
+func ext_crypto_secp256k1_ecdsa_recover_version_2(ctx context.Context, m api.Module, sig, msg uint32) uint64 {
 	return ext_crypto_secp256k1_ecdsa_recover_version_1(ctx, m, sig, msg)
 }
 
@@ -1172,17 +1172,17 @@ func ext_default_child_storage_next_key_version_1(
 
 	keyToChild := read(m, childStorageKey)
 	keyBytes := read(m, key)
-	child, err := storage.GetChildNextKey(keyToChild, keyBytes)
+	childNextKey, err := storage.GetChildNextKey(keyToChild, keyBytes)
 	if err != nil {
 		logger.Errorf("failed to get child's next key: %s", err)
-		return 0
+		return mustWrite(m, rtCtx.Allocator, noneEncoded)
 	}
 
-	ret, err := write(m, rtCtx.Allocator, scale.MustMarshal(&child))
-	if err != nil {
-		panic(err)
+	if childNextKey == nil {
+		return mustWrite(m, rtCtx.Allocator, noneEncoded)
 	}
-	return ret
+
+	return mustWrite(m, rtCtx.Allocator, scale.MustMarshal(&childNextKey))
 }
 
 func ext_default_child_storage_root_version_1(

diff --git a/lib/runtime/wazero/imports_test.go b/lib/runtime/wazero/imports_test.go
@@ -1040,8 +1040,6 @@ func Test_ext_default_child_storage_get_version_1(t *testing.T) {
 }
 
 func Test_ext_default_child_storage_next_key_version_1(t *testing.T) {
-	inst := NewTestInstance(t, runtime.HOST_API_TEST_RUNTIME)
-
 	testKeyValuePair := []struct {
 		key   []byte
 		value []byte
@@ -1050,30 +1048,73 @@ func Test_ext_default_child_storage_next_key_version_1(t *testing.T) {
 		{[]byte("key"), []byte("value2")},
 	}
 
-	key := testKeyValuePair[0].key
+	testcases := map[string]struct {
+		setupInstance func(t *testing.T) *Instance
+		expected      *[]byte
+	}{
+		"next_key_exists": {
+			setupInstance: func(t *testing.T) *Instance {
+				inst := NewTestInstance(t, runtime.HOST_API_TEST_RUNTIME)
+
+				err := inst.Context.Storage.SetChild(testChildKey, trie.NewEmptyTrie())
+				require.NoError(t, err)
 
-	err := inst.Context.Storage.SetChild(testChildKey, trie.NewEmptyTrie())
-	require.NoError(t, err)
+				for _, kv := range testKeyValuePair {
+					err = inst.Context.Storage.SetChildStorage(testChildKey, kv.key, kv.value)
+					require.NoError(t, err)
+				}
 
-	for _, kv := range testKeyValuePair {
-		err = inst.Context.Storage.SetChildStorage(testChildKey, kv.key, kv.value)
-		require.NoError(t, err)
+				return inst
+			},
+			expected: &testKeyValuePair[1].key,
+		},
+		"child_tree_not_exists": {
+			setupInstance: func(t *testing.T) *Instance {
+				return NewTestInstance(t, runtime.HOST_API_TEST_RUNTIME)
+			},
+			expected: nil,
+		},
+		"with_only_one_key": {
+			setupInstance: func(t *testing.T) *Instance {
+				inst := NewTestInstance(t, runtime.HOST_API_TEST_RUNTIME)
+
+				err := inst.Context.Storage.SetChild(testChildKey, trie.NewEmptyTrie())
+				require.NoError(t, err)
+
+				kv := testKeyValuePair[0]
+				err = inst.Context.Storage.SetChildStorage(testChildKey, kv.key, kv.value)
+				require.NoError(t, err)
+
+				return inst
+			},
+			expected: nil,
+		},
 	}
 
-	encChildKey, err := scale.Marshal(testChildKey)
-	require.NoError(t, err)
+	for tname, tt := range testcases {
+		tt := tt
 
-	encKey, err := scale.Marshal(key)
-	require.NoError(t, err)
+		t.Run(tname, func(t *testing.T) {
+			key := testKeyValuePair[0].key
 
-	ret, err := inst.Exec("rtm_ext_default_child_storage_next_key_version_1", append(encChildKey, encKey...))
-	require.NoError(t, err)
+			encChildKey, err := scale.Marshal(testChildKey)
+			require.NoError(t, err)
+
+			encKey, err := scale.Marshal(key)
+			require.NoError(t, err)
+
+			inst := tt.setupInstance(t)
+			ret, err := inst.Exec("rtm_ext_default_child_storage_next_key_version_1", append(encChildKey, encKey...))
+			require.NoError(t, err)
+
+			var read *[]byte
+			err = scale.Unmarshal(ret, &read)
+			require.NoError(t, err)
+
+			require.Equal(t, tt.expected, read)
+		})
+	}
 
-	var read *[]byte
-	err = scale.Unmarshal(ret, &read)
-	require.NoError(t, err)
-	require.NotNil(t, read)
-	require.Equal(t, testKeyValuePair[1].key, *read)
 }
 
 func Test_ext_default_child_storage_root_version_1(t *testing.T) {