go/storage/mkvs: Use cbor.UnmarshalTrusted for internal metadata

.changelog/2800.bugfix.md

@@ -0,0 +1 @@
+go/storage/mkvs: Use cbor.UnmarshalTrusted for internal metadata

.changelog/2800.internal.md

@@ -0,0 +1,5 @@
+go/common/cbor: Add UnmarshalTrusted for trusted inputs
+
+The new method relaxes some decoding restrictions for cases where the inputs
+are trusted (e.g., because they are known to be generated by the local node
+itself).

go/common/cbor/cbor.go

@@ -27,14 +27,23 @@ var (
 		TagsMd:        cbor.TagsForbidden,
 	}
 
+	// decOptions are decoding options for UNTRUSTED inputs (used by default).
 	decOptions = cbor.DecOptions{
 		DupMapKey:   cbor.DupMapKeyEnforcedAPF,
 		IndefLength: cbor.IndefLengthForbidden,
 		TagsMd:      cbor.TagsForbidden,
 	}
 
-	encMode cbor.EncMode
-	decMode cbor.DecMode
+	// decOptionsTrusted are decoding options for TRUSTED inputs. They are only used when explicitly
+	// requested by using the UnmarshalTrusted method.
+	decOptionsTrusted = cbor.DecOptions{
+		MaxArrayElements: 134217728, // Maximum allowed.
+		MaxMapPairs:      134217728, // Maximum allowed.
+	}
+
+	encMode        cbor.EncMode
+	decMode        cbor.DecMode
+	decModeTrusted cbor.DecMode
 )
 
 func init() {
@@ -45,6 +54,9 @@ func init() {
 	if decMode, err = decOptions.DecMode(); err != nil {
 		panic(err)
 	}
+	if decModeTrusted, err = decOptionsTrusted.DecMode(); err != nil {
+		panic(err)
+	}
 }
 
 // FixSliceForSerde will convert `nil` to `[]byte` to work around serde
@@ -74,6 +86,17 @@ func Unmarshal(data []byte, dst interface{}) error {
 	return decMode.Unmarshal(data, dst)
 }
 
+// UnmarshalTrusted deserializes a CBOR byte vector into a given type.
+//
+// This method MUST ONLY BE USED FOR TRUSTED INPUTS as it relaxes some decoding restrictions.
+func UnmarshalTrusted(data []byte, dst interface{}) error {
+	if data == nil {
+		return nil
+	}
+
+	return decModeTrusted.Unmarshal(data, dst)
+}
+
 // MustUnmarshal deserializes a CBOR byte vector into a given type.
 // Panics if unmarshal fails.
 func MustUnmarshal(data []byte, dst interface{}) {

go/storage/mkvs/db/badger/badger.go

@@ -114,7 +114,7 @@ func (d *badgerNodeDB) load() error {
 		// Metadata already exists, just load it and verify that it is
 		// compatible with what we have here.
 		err = item.Value(func(data []byte) error {
-			return cbor.Unmarshal(data, &d.meta.value)
+			return cbor.UnmarshalTrusted(data, &d.meta.value)
 		})
 		if err != nil {
 			return err
@@ -303,7 +303,7 @@ func (d *badgerNodeDB) GetWriteLog(ctx context.Context, startRoot node.Root, end
 
 							var log api.HashedDBWriteLog
 							err = item.Value(func(data []byte) error {
-								return cbor.Unmarshal(data, &log)
+								return cbor.UnmarshalTrusted(data, &log)
 							})
 							if err != nil {
 								return node.Root{}, nil, err
@@ -464,7 +464,7 @@ func (d *badgerNodeDB) Finalize(ctx context.Context, version uint64, roots []has
 
 		var updatedNodes []updatedNode
 		err = item.Value(func(data []byte) error {
-			return cbor.Unmarshal(data, &updatedNodes)
+			return cbor.UnmarshalTrusted(data, &updatedNodes)
 		})
 		if err != nil {
 			panic(fmt.Errorf("mkvs/badger: corrupted root updated nodes index: %w", err))

go/storage/mkvs/tree_test.go

@@ -1889,6 +1889,23 @@ func testSpecialCase5(t *testing.T, ndb db.NodeDB, factory NodeDBFactory) {
 	testSpecialCaseFromJSON(t, ndb, "case-5.json")
 }
 
+func testLargeUpdates(t *testing.T, ndb db.NodeDB, factory NodeDBFactory) {
+	ctx := context.Background()
+	tree := New(nil, ndb)
+
+	// The number of elements is such that it would overflow the maximum number of allowed array
+	// elements in the default (untrusted) CBOR decoder.
+	for i := 0; i < 132_000; i++ {
+		err := tree.Insert(ctx, []byte(fmt.Sprintf("%d", i)), []byte(fmt.Sprintf("%d", i)))
+		require.NoError(t, err, "Insert")
+	}
+
+	_, rootHash, err := tree.Commit(ctx, testNs, 0)
+	require.NoError(t, err, "Commit")
+	err = ndb.Finalize(ctx, 0, []hash.Hash{rootHash})
+	require.NoError(t, err, "Finalize")
+}
+
 func testBackend(
 	t *testing.T,
 	initBackend func(t *testing.T) (NodeDBFactory, func()),
@@ -1929,6 +1946,7 @@ func testBackend(
 		{"SpecialCase3", testSpecialCase3},
 		{"SpecialCase4", testSpecialCase4},
 		{"SpecialCase5", testSpecialCase5},
+		{"LargeUpdates", testLargeUpdates},
 		{"Errors", testErrors},
 	}