fix: Implement correct balance updating logic

params/extras/state_upgrade.go

@@ -6,6 +6,7 @@ package extras
 import (
 	"errors"
 	"fmt"
+	"math/big"
 	"reflect"
 
 	"github.com/ava-labs/libevm/common"
@@ -44,6 +45,7 @@ func (s *StateUpgrade) Equal(other *StateUpgrade) bool {
 
 // verifyStateUpgrades checks [c.StateUpgrades] is well formed:
 // - the specified blockTimestamps must monotonically increase
+// - all BalanceChange values must fit within uint256
 func (c *ChainConfig) verifyStateUpgrades() error {
 	var previousUpgradeTimestamp *uint64
 	for i, upgrade := range c.StateUpgrades {
@@ -61,6 +63,19 @@ func (c *ChainConfig) verifyStateUpgrades() error {
 			return fmt.Errorf("%w: StateUpgrade[%d] has timestamp %v, previous timestamp %v", errStateUpgradeTimestampNotMonotonic, i, *upgradeTimestamp, *previousUpgradeTimestamp)
 		}
 		previousUpgradeTimestamp = upgradeTimestamp
+
+		// Verify all BalanceChange values fit within uint256
+		for account, accountUpgrade := range upgrade.StateUpgradeAccounts {
+			if accountUpgrade.BalanceChange != nil {
+				bigChange := (*big.Int)(accountUpgrade.BalanceChange)
+				// Check if the absolute value fits in uint256
+				absChange := new(big.Int).Abs(bigChange)
+				if absChange.BitLen() > 256 {
+					return fmt.Errorf("StateUpgrade[%d]: account %s has BalanceChange %s that exceeds uint256 bit length",
+						i, account.Hex(), bigChange.String())
+				}
+			}
+		}
 	}
 	return nil
 }

params/extras/state_upgrade_test.go

@@ -22,6 +22,14 @@
 			BalanceChange: (*math.HexOrDecimal256)(common.Big1),
 		},
 	}
+	// Create a value that exceeds uint256 (2^256 + 1)
+	maxUint256 := new(big.Int).Lsh(big.NewInt(1), 256)
+	maxUint256.Sub(maxUint256, big.NewInt(1)) // 2^256 - 1
+	overflowValue := new(big.Int).Add(maxUint256, big.NewInt(2))
+
+	// Create a negative value that exceeds uint256 in absolute value
+	negativeOverflowValue := new(big.Int).Neg(overflowValue)
+
 	tests := []struct {
 		name          string
 		upgrades      []StateUpgrade
@@ -58,6 +66,60 @@
 			},
 			expectedError: errStateUpgradeTimestampZero,
 		},
+		{
+			name: "balance change exceeds uint256",
+			upgrades: []StateUpgrade{
+				{
+					BlockTimestamp: utils.NewUint64(1),
+					StateUpgradeAccounts: map[common.Address]StateUpgradeAccount{
+						{1}: {
+							BalanceChange: (*math.HexOrDecimal256)(overflowValue),
+						},
+					},
+				},
+			},
+			expectedError: "exceeds uint256 bit length",
+		},
+		{
+			name: "negative balance change exceeds uint256",
+			upgrades: []StateUpgrade{
+				{
+					BlockTimestamp: utils.NewUint64(1),
+					StateUpgradeAccounts: map[common.Address]StateUpgradeAccount{
+						{1}: {
+							BalanceChange: (*math.HexOrDecimal256)(negativeOverflowValue),
+						},
+					},
+				},
+			},
+			expectedError: "exceeds uint256 bit length",
+		},
+		{
+			name: "max uint256 balance change is valid",
+			upgrades: []StateUpgrade{
+				{
+					BlockTimestamp: utils.NewUint64(1),
+					StateUpgradeAccounts: map[common.Address]StateUpgradeAccount{
+						{1}: {
+							BalanceChange: (*math.HexOrDecimal256)(maxUint256),
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "negative max uint256 balance change is valid",
+			upgrades: []StateUpgrade{
+				{
+					BlockTimestamp: utils.NewUint64(1),
+					StateUpgradeAccounts: map[common.Address]StateUpgradeAccount{
+						{1}: {
+							BalanceChange: (*math.HexOrDecimal256)(new(big.Int).Neg(maxUint256)),
+						},
+					},
+				},
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

stateupgrade/interfaces.go

@@ -15,7 +15,10 @@ import (
 type StateDB interface {
 	SetState(common.Address, common.Hash, common.Hash, ...stateconf.StateDBStateOption)
 	SetCode(common.Address, []byte)
+
+	GetBalance(common.Address) *uint256.Int
 	AddBalance(common.Address, *uint256.Int)
+	SubBalance(common.Address, *uint256.Int)
 
 	GetNonce(common.Address) uint64
 	SetNonce(common.Address, uint64)

stateupgrade/state_upgrade.go

@@ -7,33 +7,68 @@ import (
 	"math/big"
 
 	"github.com/ava-labs/libevm/common"
+	"github.com/ava-labs/libevm/log"
 	"github.com/holiman/uint256"
 
 	"github.com/ava-labs/subnet-evm/params/extras"
 )
 
 // Configure applies the state upgrade to the state.
+// Note: This function should not return errors as any error would break the chain.
+// All validation should be done in verifyStateUpgrades at config load time.
 func Configure(stateUpgrade *extras.StateUpgrade, chainConfig ChainContext, state StateDB, blockContext BlockContext) error {
 	isEIP158 := chainConfig.IsEIP158(blockContext.Number())
 	for account, upgrade := range stateUpgrade.StateUpgradeAccounts {
-		if err := upgradeAccount(account, upgrade, state, isEIP158); err != nil {
-			return err
-		}
+		upgradeAccount(account, upgrade, state, isEIP158)
 	}
 	return nil
 }
 
 // upgradeAccount applies the state upgrade to the given account.
-func upgradeAccount(account common.Address, upgrade extras.StateUpgradeAccount, state StateDB, isEIP158 bool) error {
+func upgradeAccount(account common.Address, upgrade extras.StateUpgradeAccount, state StateDB, isEIP158 bool) {
 	// Create the account if it does not exist
 	if !state.Exist(account) {
 		state.CreateAccount(account)
 	}
 
+	// Balance change detected - update the balance of the account
 	if upgrade.BalanceChange != nil {
-		balanceChange, _ := uint256.FromBig((*big.Int)(upgrade.BalanceChange))
-		state.AddBalance(account, balanceChange)
+		// BalanceChange is a HexOrDecimal256, which is just a typed big.Int
+		// Note: We validate at config load time that the change itself fits in uint256,
+		// but we still need to check if the resulting balance would overflow/underflow.
+		bigChange := (*big.Int)(upgrade.BalanceChange)
+		currentBalance := state.GetBalance(account)
+
+		switch bigChange.Sign() {
+		case 1: // Positive - check for overflow
+			balanceChange, _ := uint256.FromBig(bigChange)
+			_, overflow := new(uint256.Int).AddOverflow(currentBalance, balanceChange)
+			if overflow {
+				log.Warn("State upgrade balance change would overflow, clamping to max uint256",
+					"account", account.Hex(),
+					"currentBalance", currentBalance.ToBig().String(),
+					"balanceChange", balanceChange.ToBig().String())
+				state.SubBalance(account, currentBalance)
+				state.AddBalance(account, new(uint256.Int).SetAllOne())
+			} else {
+				state.AddBalance(account, balanceChange)
+			}
+		case -1: // Negative - check for underflow
+			absChange := new(big.Int).Abs(bigChange)
+			balanceChange, _ := uint256.FromBig(absChange)
+			if currentBalance.Cmp(balanceChange) < 0 {
+				log.Warn("State upgrade balance change would underflow, clamping to zero",
+					"account", account.Hex(),
+					"currentBalance", currentBalance.ToBig().String(),
+					"balanceChange", balanceChange.ToBig().String())
+				state.SubBalance(account, currentBalance)
+			} else {
+				state.SubBalance(account, balanceChange)
+			}
+		}
+		// If zero (Sign() == 0), do nothing
 	}
+
 	if len(upgrade.Code) != 0 {
 		// if the nonce is 0, set the nonce to 1 as we would when deploying a contract at
 		// the address.
@@ -45,5 +80,4 @@ func upgradeAccount(account common.Address, upgrade extras.StateUpgradeAccount,
 	for key, value := range upgrade.Storage {
 		state.SetState(account, key, value)
 	}
-	return nil
 }

stateupgrade/state_upgrade_test.go

@@ -0,0 +1,173 @@
+// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package stateupgrade
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/ava-labs/libevm/common"
+	"github.com/ava-labs/libevm/common/math"
+	"github.com/ava-labs/libevm/core/rawdb"
+	"github.com/ava-labs/libevm/core/state"
+	"github.com/ava-labs/libevm/core/types"
+	"github.com/holiman/uint256"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ava-labs/subnet-evm/core/extstate"
+	"github.com/ava-labs/subnet-evm/params/extras"
+)
+
+func TestUpgradeAccount_BalanceChanges(t *testing.T) {
+	testAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
+
+	tests := []struct {
+		name           string
+		initialBalance *uint256.Int
+		balanceChange  *math.HexOrDecimal256
+		accountExists  bool
+		wantBalance    *uint256.Int
+	}{
+		{
+			name:           "positive balance change on existing account",
+			initialBalance: uint256.NewInt(100),
+			balanceChange:  hexOrDecimal256FromInt64(50),
+			accountExists:  true,
+			wantBalance:    uint256.NewInt(150),
+		},
+		{
+			name:           "negative balance change with sufficient funds",
+			initialBalance: uint256.NewInt(100),
+			balanceChange:  hexOrDecimal256FromInt64(-30),
+			accountExists:  true,
+			wantBalance:    uint256.NewInt(70),
+		},
+		{
+			name:           "zero balance change",
+			initialBalance: uint256.NewInt(100),
+			balanceChange:  hexOrDecimal256FromInt64(0),
+			accountExists:  true,
+			wantBalance:    uint256.NewInt(100),
+		},
+		{
+			name:           "nil balance change",
+			initialBalance: uint256.NewInt(100),
+			accountExists:  true,
+			wantBalance:    uint256.NewInt(100),
+		},
+		{
+			name:           "new account with positive balance",
+			initialBalance: uint256.NewInt(0),
+			balanceChange:  hexOrDecimal256FromInt64(1000),
+			wantBalance:    uint256.NewInt(1000),
+		},
+		{
+			name:           "exact balance subtraction",
+			initialBalance: uint256.NewInt(100),
+			balanceChange:  hexOrDecimal256FromInt64(-100),
+			accountExists:  true,
+			wantBalance:    uint256.NewInt(0),
+		},
+		{
+			name: "large positive balance change",
+			// Initial balance: 2^255 - 1000 (a very large number near half of max uint256)
+			initialBalance: uint256.MustFromBig(new(big.Int).Sub(new(big.Int).Lsh(big.NewInt(1), 255), big.NewInt(1000))),
+			balanceChange:  hexOrDecimal256FromInt64(500),
+			accountExists:  true,
+			// Expected balance: 2^255 - 500 (initial + 500)
+			wantBalance: uint256.MustFromBig(new(big.Int).Sub(new(big.Int).Lsh(big.NewInt(1), 255), big.NewInt(500))),
+		},
+		{
+			name: "balance overflow clamped to max uint256",
+			// Set initial balance to max uint256 - 100
+			initialBalance: new(uint256.Int).Sub(new(uint256.Int).SetAllOne(), uint256.NewInt(100)),
+			balanceChange:  hexOrDecimal256FromInt64(200), // This would overflow
+			accountExists:  true,
+			wantBalance:    new(uint256.Int).SetAllOne(), // max uint256
+		},
+		{
+			name:           "balance underflow clamped to zero",
+			initialBalance: uint256.NewInt(50),
+			balanceChange:  hexOrDecimal256FromInt64(-100), // More than current balance
+			accountExists:  true,
+			wantBalance:    uint256.NewInt(0),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create a fresh state database for each test
+			statedb := createTestStateDB(t)
+
+			// Set up initial account state
+			if tt.accountExists {
+				setAccountBalance(statedb, testAddr, tt.initialBalance)
+			}
+
+			upgrade := extras.StateUpgradeAccount{
+				BalanceChange: tt.balanceChange,
+			}
+			upgradeAccount(testAddr, upgrade, statedb, false)
+
+			// Check final balance
+			actualBalance := statedb.GetBalance(testAddr)
+			require.Equal(t, tt.wantBalance, actualBalance, "final balance mismatch")
+
+			// Verify account was created if it didn't exist
+			require.True(t, statedb.Exist(testAddr), "account should exist after upgrade")
+		})
+	}
+}
+
+func TestUpgradeAccount_CompleteUpgrade(t *testing.T) {
+	statedb := createTestStateDB(t)
+	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
+
+	// Create a complete state upgrade with balance change, code, and storage
+	code := []byte{0xde, 0xad, 0xbe, 0xef}
+	storageKey := common.HexToHash("0x1234")
+	storageValue := common.HexToHash("0x5678")
+
+	upgrade := extras.StateUpgradeAccount{
+		BalanceChange: hexOrDecimal256FromInt64(1000),
+		Code:          code,
+		Storage:       map[common.Hash]common.Hash{storageKey: storageValue},
+	}
+	upgradeAccount(addr, upgrade, statedb, true) // Test with EIP158 = true
+
+	// Verify all changes were applied
+	require.Equal(t, uint256.NewInt(1000), statedb.GetBalance(addr))
+	require.True(t, statedb.Exist(addr))
+	require.Equal(t, uint64(1), statedb.GetNonce(addr)) // Should be set to 1 due to EIP158 and code deployment
+
+	// Cast to access code and storage (these aren't in our StateDB interface)
+	extStateDB := statedb.(*extstate.StateDB)
+	require.Equal(t, code, extStateDB.GetCode(addr))
+	require.Equal(t, storageValue, extStateDB.GetState(addr, storageKey))
+}
+
+// createTestStateDB creates a real StateDB instance for testing
+func createTestStateDB(t *testing.T) StateDB {
+	t.Helper()
+
+	db := rawdb.NewMemoryDatabase()
+	statedb, err := state.New(types.EmptyRootHash, state.NewDatabase(db), nil)
+	require.NoError(t, err)
+
+	// Wrap it with extstate for predicate support (matching the interface)
+	return extstate.New(statedb)
+}
+
+// setAccountBalance is a helper to set an account's initial balance for testing
+func setAccountBalance(statedb StateDB, addr common.Address, balance *uint256.Int) {
+	// Cast to extstate.StateDB to access the underlying state
+	extStateDB := statedb.(*extstate.StateDB)
+	extStateDB.CreateAccount(addr)
+	extStateDB.SetBalance(addr, balance)
+}
+
+// Helper function to create HexOrDecimal256 from int64
+func hexOrDecimal256FromInt64(i int64) *math.HexOrDecimal256 {
+	return (*math.HexOrDecimal256)(big.NewInt(i))
+}