feat: add a state version guard

Author: MaksymMalicki

blockchain/state_version.go

@@ -0,0 +1,57 @@
+package blockchain
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/NethermindEth/juno/core"
+	"github.com/NethermindEth/juno/db"
+)
+
+// ValidateStateVersion checks that the state implementation selected by useNewState
+// is consistent with the data already stored in disk. It must be called before any
+// state reads or writes so that incompatible databases are rejected early.
+//
+// It returns an error if:
+//   - the database was built with the deprecated state and --new-state is set, or
+//   - the database was built with the new state and --new-state is not set.
+//
+// An empty database is always accepted.
+func ValidateStateVersion(disk db.KeyValueStore, useNewState bool) error {
+	// A fresh database has no conflict.
+	if _, err := core.GetChainHeight(disk); errors.Is(err, db.ErrKeyNotFound) {
+		return nil
+	}
+
+	hasNewState, err := bucketHasData(disk, db.ClassTrie)
+	if err != nil {
+		return fmt.Errorf("probe new-state bucket: %w", err)
+	}
+	hasOldState, err := bucketHasData(disk, db.ClassesTrie)
+	if err != nil {
+		return fmt.Errorf("probe state bucket: %w", err)
+	}
+
+	switch {
+	case useNewState && hasOldState:
+		return errors.New(
+			"database was built with the existing state implementation, " +
+				"but --new-state is set; remove --new-state to use this database",
+		)
+	case !useNewState && hasNewState:
+		return errors.New(
+			"database was built with the new state implementation, " +
+				"but --new-state is not set; add --new-state to use this database",
+		)
+	}
+	return nil
+}
+
+func bucketHasData(disk db.KeyValueStore, bucket db.Bucket) (bool, error) {
+	it, err := disk.NewIterator(bucket.Key(), false)
+	if err != nil {
+		return false, err
+	}
+	defer it.Close()
+	return it.Next(), nil
+}

blockchain/state_version_test.go

@@ -0,0 +1,78 @@
+package blockchain_test
+
+import (
+	"testing"
+
+	"github.com/NethermindEth/juno/blockchain"
+	"github.com/NethermindEth/juno/core"
+	"github.com/NethermindEth/juno/db"
+	"github.com/NethermindEth/juno/db/memory"
+	"github.com/stretchr/testify/require"
+)
+
+func TestValidateStateVersion(t *testing.T) {
+	t.Run("fresh database is always accepted", func(t *testing.T) {
+		testDB := memory.New()
+		require.NoError(t, blockchain.ValidateStateVersion(testDB, false))
+		require.NoError(t, blockchain.ValidateStateVersion(testDB, true))
+	})
+
+	t.Run("deprecated-state DB with new-state flag returns error", func(t *testing.T) {
+		testDB := memory.New()
+		seedChainHeight(t, testDB)
+		seedBucket(t, testDB, db.ClassesTrie)
+
+		err := blockchain.ValidateStateVersion(testDB, true)
+		require.ErrorContains(t, err, "deprecated state")
+		require.ErrorContains(t, err, "--new-state")
+	})
+
+	t.Run("new-state DB without new-state flag returns error", func(t *testing.T) {
+		testDB := memory.New()
+		seedChainHeight(t, testDB)
+		seedBucket(t, testDB, db.ClassTrie)
+
+		err := blockchain.ValidateStateVersion(testDB, false)
+		require.ErrorContains(t, err, "new state")
+		require.ErrorContains(t, err, "--new-state")
+	})
+
+	t.Run("deprecated-state DB without new-state flag is accepted", func(t *testing.T) {
+		testDB := memory.New()
+		seedChainHeight(t, testDB)
+		seedBucket(t, testDB, db.ClassesTrie)
+
+		require.NoError(t, blockchain.ValidateStateVersion(testDB, false))
+	})
+
+	t.Run("new-state DB with new-state flag is accepted", func(t *testing.T) {
+		testDB := memory.New()
+		seedChainHeight(t, testDB)
+		seedBucket(t, testDB, db.ClassTrie)
+
+		require.NoError(t, blockchain.ValidateStateVersion(testDB, true))
+	})
+
+	t.Run("DB with blocks but no trie data is accepted regardless of flag", func(t *testing.T) {
+		testDB := memory.New()
+		seedChainHeight(t, testDB)
+
+		require.NoError(t, blockchain.ValidateStateVersion(testDB, false))
+		require.NoError(t, blockchain.ValidateStateVersion(testDB, true))
+	})
+}
+
+func seedChainHeight(t *testing.T, testDB db.KeyValueStore) {
+	t.Helper()
+	batch := testDB.NewBatch()
+	require.NoError(t, core.WriteChainHeight(batch, 1))
+	require.NoError(t, batch.Write())
+}
+
+func seedBucket(t *testing.T, testDB db.KeyValueStore, bucket db.Bucket) {
+	t.Helper()
+	batch := testDB.NewBatch()
+	key := append(bucket.Key(), []byte("sentinel")...)
+	require.NoError(t, batch.Put(key, []byte("1")))
+	require.NoError(t, batch.Write())
+}

node/node.go

@@ -190,6 +190,10 @@ func New(cfg *Config, version string, logLevel *utils.LogLevel) (*Node, error) {
 	if err != nil {
 		return nil, fmt.Errorf("open DB: %w", err)
 	}
+	if err = blockchain.ValidateStateVersion(database, cfg.NewState); err != nil {
+		return nil, err
+	}
+
 	ua := fmt.Sprintf("Juno/%s Starknet Client", version)
 
 	services := make([]service.Service, 0)