fix: Avoid leaking returned bytes outside of Get closure (#3327)

Author: infrmtcs

cmd/juno/dbcmd.go

@@ -100,7 +100,7 @@ func dbInfo(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to get the state update: %v", err)
 	}
 
-	schemaMeta, err := migration.SchemaMetadata(database)
+	schemaMeta, err := migration.SchemaMetadata(utils.NewNopZapLogger(), database)
 	if err != nil {
 		return fmt.Errorf("failed to get schema metadata: %v", err)
 	}

core/accessors.go

@@ -112,16 +112,16 @@ func DeleteStateUpdateByBlockNum(w db.KeyValueWriter, blockNum uint64) error {
 }
 
 func GetStateUpdateByHash(r db.KeyValueReader, hash *felt.Felt) (*StateUpdate, error) {
-	var val []byte
+	var blockNum uint64
 	err := r.Get(db.BlockHeaderNumbersByHashKey(hash), func(data []byte) error {
-		val = data
+		blockNum = binary.BigEndian.Uint64(data)
 		return nil
 	})
 	if err != nil {
 		return nil, err
 	}
 
-	return GetStateUpdateByBlockNum(r, binary.BigEndian.Uint64(val))
+	return GetStateUpdateByBlockNum(r, blockNum)
 }
 
 // WriteContractStorageHistory writes the old value of a storage location

core/state.go

@@ -160,24 +160,18 @@ func (s *State) globalTrie(bucket db.Bucket, newTrie trie.NewTrieFunc) (*trie.Tr
 
 	// fetch root key
 	rootKeyDBKey := dbPrefix
-	var val []byte
+	rootKey := new(trie.BitArray)
 	err := s.txn.Get(rootKeyDBKey, func(data []byte) error {
-		val = data
+		if len(data) > 0 {
+			return rootKey.UnmarshalBinary(data)
+		}
 		return nil
 	})
 	// if some error other than "not found"
 	if err != nil && !errors.Is(err, db.ErrKeyNotFound) {
 		return nil, nil, err
 	}
 
-	rootKey := new(trie.BitArray)
-	if len(val) > 0 {
-		err = rootKey.UnmarshalBinary(val)
-		if err != nil {
-			return nil, nil, err
-		}
-	}
-
 	gTrie, err := newTrie(s.txn, dbPrefix, globalTrieHeight)
 	if err != nil {
 		return nil, nil, err

db/remote/db_test.go

@@ -45,17 +45,15 @@ func TestRemote(t *testing.T) {
 	t.Run("Get", func(t *testing.T) {
 		assert.NoError(t, remoteDB.View(func(txn db.Snapshot) error {
 			for i := byte(0); i < 3; i++ {
-				var val []byte
 				err := txn.Get([]byte{i}, func(data []byte) error {
-					val = data
+					if !bytes.Equal(data, []byte{i}) {
+						return errors.New("wrong value")
+					}
 					return nil
 				})
 				if err != nil {
 					return err
 				}
-				if !bytes.Equal(val, []byte{i}) {
-					return errors.New("wrong value")
-				}
 
 				assert.Equal(t, db.ErrKeyNotFound, txn.Get([]byte{0xDE, 0xAD}, func(b []byte) error { return nil }))
 			}

grpc/handlers.go

@@ -76,18 +76,16 @@ func (h Handler) handleTxCursor(
 		responsePair.CursorId = cursorID
 		return server.Send(responsePair)
 	} else if cur.Op == gen.Op_GET {
-		var val []byte
 		err := tx.dbTx.Get(cur.K, func(data []byte) error {
-			val = data
+			if data != nil {
+				responsePair.V = data
+				responsePair.K = cur.K
+			}
 			return nil
 		})
 		if err != nil && !errors.Is(err, db.ErrKeyNotFound) {
 			return err
 		}
-		if val != nil {
-			responsePair.V = val
-			responsePair.K = cur.K
-		}
 		return server.Send(responsePair)
 	}
 

migration/bucket_migrator_test.go

@@ -45,30 +45,26 @@ func TestBucketMover(t *testing.T) {
 	require.NoError(t, err)
 
 	err = testDB.View(func(txn db.Snapshot) error {
-		var val []byte
 		err := txn.Get(sourceBucket.Key(), func(data []byte) error {
-			val = data
+			if !bytes.Equal(data, []byte{44}) {
+				return errors.New("shouldnt have changed")
+			}
 			return nil
 		})
 		if err != nil {
 			return err
 		}
-		if !bytes.Equal(val, []byte{44}) {
-			return errors.New("shouldnt have changed")
-		}
 
 		for i := byte(0); i < 3; i++ {
-			var val []byte
 			err := txn.Get(destBucket.Key([]byte{i}), func(data []byte) error {
-				val = data
+				if !bytes.Equal(data, []byte{i}) {
+					return errors.New("shouldve moved")
+				}
 				return nil
 			})
 			if err != nil {
 				return err
 			}
-			if !bytes.Equal(val, []byte{i}) {
-				return errors.New("shouldve moved")
-			}
 
 			err = txn.Get(sourceBucket.Key([]byte{i}), func([]byte) error { return nil })
 			require.ErrorIs(t, db.ErrKeyNotFound, err)

migration/migration.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/binary"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -117,7 +118,7 @@ func migrateIfNeeded(
 		new ones. It will be able to do this since the schema version it reads from the database will be
 		non-zero and that is what we use to initialise the i loop variable.
 	*/
-	metadata, err := SchemaMetadata(targetDB)
+	metadata, err := SchemaMetadata(log, targetDB)
 	if err != nil {
 		return err
 	}
@@ -182,36 +183,34 @@ func closeMigrationServer(srv *http.Server, log utils.SimpleLogger) {
 }
 
 // SchemaMetadata retrieves metadata about a database schema from the given database.
-func SchemaMetadata(targetDB db.KeyValueStore) (schemaMetadata, error) {
+func SchemaMetadata(log utils.SimpleLogger, targetDB db.KeyValueStore) (schemaMetadata, error) {
 	metadata := schemaMetadata{}
 	txn := targetDB
 
-	var sv []byte
 	err := txn.Get(db.SchemaVersion.Key(), func(data []byte) error {
-		sv = data
+		metadata.Version = binary.BigEndian.Uint64(data)
 		return nil
 	})
 	if err != nil && !errors.Is(err, db.ErrKeyNotFound) {
 		return metadata, err
 	}
 
-	if sv != nil {
-		metadata.Version = binary.BigEndian.Uint64(sv)
-	}
-
-	var is []byte
 	err = txn.Get(db.SchemaIntermediateState.Key(), func(data []byte) error {
-		is = data
+		err := cbor.Unmarshal(data, &metadata.IntermediateState)
+		if err != nil {
+			// TODO: Instead of returning nil, we log the error for now to debug the issue
+			log.Errorw(
+				"Failed to unmarshal intermediate state",
+				"version", metadata.Version,
+				"state", hex.EncodeToString(data),
+				"err", err,
+			)
+		}
 		return nil
 	})
 	if err != nil && !errors.Is(err, db.ErrKeyNotFound) {
 		return metadata, err
 	}
-	if is != nil {
-		if err := cbor.Unmarshal(is, &metadata.IntermediateState); err != nil {
-			return metadata, err
-		}
-	}
 
 	return metadata, nil
 }

migration/migration_pkg_test.go

@@ -65,13 +65,16 @@ func TestRelocateContractStorageRootKeys(t *testing.T) {
 		exampleBytes := new(felt.Felt).SetUint64(uint64(i)).Bytes()
 
 		// New entry exists.
-		var val []byte
 		err := txn.Get(db.ContractStorage.Key(exampleBytes[:]), func(data []byte) error {
-			val = data
+			require.Equal(
+				t,
+				exampleBytes[:],
+				data,
+				"the correct value was not transferred to the new location",
+			)
 			return nil
 		})
 		require.NoError(t, err)
-		require.Equal(t, exampleBytes[:], val, "the correct value was not transferred to the new location")
 
 		// Old entry does not exist.
 		oldKey := db.Peer.Key(exampleBytes[:])
@@ -400,7 +403,7 @@ func TestSchemaMetadata(t *testing.T) {
 	t.Run("conversion", func(t *testing.T) {
 		t.Run("version not set", func(t *testing.T) {
 			testDB := memory.New()
-			metadata, err := SchemaMetadata(testDB)
+			metadata, err := SchemaMetadata(utils.NewNopZapLogger(), testDB)
 			require.NoError(t, err)
 			require.Equal(t, uint64(0), metadata.Version)
 			require.Nil(t, metadata.IntermediateState)
@@ -414,7 +417,7 @@ func TestSchemaMetadata(t *testing.T) {
 				return txn.Put(db.SchemaVersion.Key(), version[:])
 			}))
 
-			metadata, err := SchemaMetadata(testDB)
+			metadata, err := SchemaMetadata(utils.NewNopZapLogger(), testDB)
 			require.NoError(t, err)
 			require.Equal(t, uint64(1), metadata.Version)
 			require.Nil(t, metadata.IntermediateState)
@@ -430,7 +433,7 @@ func TestSchemaMetadata(t *testing.T) {
 					IntermediateState: nil,
 				})
 			}))
-			metadata, err := SchemaMetadata(testDB)
+			metadata, err := SchemaMetadata(utils.NewNopZapLogger(), testDB)
 			require.NoError(t, err)
 			require.Equal(t, version, metadata.Version)
 			require.Nil(t, metadata.IntermediateState)
@@ -448,7 +451,7 @@ func TestSchemaMetadata(t *testing.T) {
 					IntermediateState: intermediateState,
 				})
 			}))
-			metadata, err := SchemaMetadata(testDB)
+			metadata, err := SchemaMetadata(utils.NewNopZapLogger(), testDB)
 			require.NoError(t, err)
 			require.Equal(t, version, metadata.Version)
 			require.Equal(t, intermediateState, metadata.IntermediateState)
@@ -466,7 +469,7 @@ func TestSchemaMetadata(t *testing.T) {
 					IntermediateState: intermediateState,
 				})
 			}))
-			metadata, err := SchemaMetadata(testDB)
+			metadata, err := SchemaMetadata(utils.NewNopZapLogger(), testDB)
 			require.NoError(t, err)
 			require.Equal(t, version, metadata.Version)
 			require.Equal(t, intermediateState, metadata.IntermediateState)

migration/migration_test.go

@@ -19,7 +19,7 @@ func TestMigrateIfNeeded(t *testing.T) {
 		require.ErrorIs(t, migration.MigrateIfNeeded(ctx, testDB, &utils.Mainnet, utils.NewNopZapLogger(), &migration.HTTPConfig{}), ctx.Err())
 	})
 
-	meta, err := migration.SchemaMetadata(testDB)
+	meta, err := migration.SchemaMetadata(utils.NewNopZapLogger(), testDB)
 	require.NoError(t, err)
 	require.Equal(t, uint64(0), meta.Version)
 	require.Nil(t, meta.IntermediateState)
@@ -28,14 +28,14 @@ func TestMigrateIfNeeded(t *testing.T) {
 		require.NoError(t, migration.MigrateIfNeeded(t.Context(), testDB, &utils.Mainnet, utils.NewNopZapLogger(), &migration.HTTPConfig{}))
 	})
 
-	meta, err = migration.SchemaMetadata(testDB)
+	meta, err = migration.SchemaMetadata(utils.NewNopZapLogger(), testDB)
 	require.NoError(t, err)
 	require.NotEqual(t, uint64(0), meta.Version)
 	require.Nil(t, meta.IntermediateState)
 
 	t.Run("subsequent calls to MigrateIfNeeded should not change the DB version", func(t *testing.T) {
 		require.NoError(t, migration.MigrateIfNeeded(t.Context(), testDB, &utils.Mainnet, utils.NewNopZapLogger(), &migration.HTTPConfig{}))
-		postVersion, postErr := migration.SchemaMetadata(testDB)
+		postVersion, postErr := migration.SchemaMetadata(utils.NewNopZapLogger(), testDB)
 		require.NoError(t, postErr)
 		require.Equal(t, meta, postVersion)
 	})