From f506350ddb98e9568f7856a68b1e2a2100ea4cdf Mon Sep 17 00:00:00 2001
From: Anshal Shukla <53994948+anshalshukla@users.noreply.github.com>
Date: Fri, 1 Dec 2023 11:39:58 +0530
Subject: [PATCH] consensus/bor: handle blockalloc balance changes (#1074)

* fix: set balance in blockalloc

* chg: don't update balance if not zero

* fix: lint

* fix logic, add test cases
---
 consensus/bor/bor.go      |  4 ++++
 consensus/bor/bor_test.go | 31 ++++++++++++++++++++++++-------
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/consensus/bor/bor.go b/consensus/bor/bor.go
index 40999f73c4..0314cef8af 100644
--- a/consensus/bor/bor.go
+++ b/consensus/bor/bor.go
@@ -883,6 +883,10 @@ func (c *Bor) changeContractCodeIfNeeded(headerNumber uint64, state *state.State
 			for addr, account := range allocs {
 				log.Info("change contract code", "address", addr)
 				state.SetCode(addr, account.Code)
+
+				if state.GetBalance(addr).Cmp(big.NewInt(0)) == 0 {
+					state.SetBalance(addr, account.Balance)
+				}
 			}
 		}
 	}
diff --git a/consensus/bor/bor_test.go b/consensus/bor/bor_test.go
index 020289c4dc..9713ce7516 100644
--- a/consensus/bor/bor_test.go
+++ b/consensus/bor/bor_test.go
@@ -41,6 +41,12 @@ func TestGenesisContractChange(t *testing.T) {
 						"balance": "0x1000",
 					},
 				},
+				"6": map[string]interface{}{
+					addr0.Hex(): map[string]interface{}{
+						"code":    hexutil.Bytes{0x1, 0x4},
+						"balance": "0x2000",
+					},
+				},
 			},
 		},
 	}
@@ -87,24 +93,35 @@ func TestGenesisContractChange(t *testing.T) {
 
 	root := genesis.Root()
 
-	// code does not change
+	// code does not change, balance remains 0
 	root, statedb = addBlock(root, 1)
 	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x1})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(0))
 
-	// code changes 1st time
+	// code changes 1st time, balance remains 0
 	root, statedb = addBlock(root, 2)
 	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x2})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(0))
 
-	// code same as 1st change
+	// code same as 1st change, balance remains 0
 	root, statedb = addBlock(root, 3)
 	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x2})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(0))
 
-	// code changes 2nd time
-	_, statedb = addBlock(root, 4)
+	// code changes 2nd time, balance updates to 4096
+	root, statedb = addBlock(root, 4)
 	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x3})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(4096))
 
-	// make sure balance change DOES NOT take effect
-	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(0))
+	// code same as 2nd change, balance remains 4096
+	root, statedb = addBlock(root, 5)
+	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x3})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(4096))
+
+	// code changes 3rd time, balance remains 4096
+	_, statedb = addBlock(root, 6)
+	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x4})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(4096))
 }
 
 func TestEncodeSigHeaderJaipur(t *testing.T) {