From 113e2d3a1d47cf34b7fa6508fc7303f8963eed13 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ste=CC=81phane=20Duchesneau?=
 <stephane.duchesneau@streamingfast.io>
Date: Tue, 11 Feb 2025 14:35:30 -0500
Subject: [PATCH] add firehose handling of SetBalance, add compat for SetNonce
 for systemtransactions

---
 core/state/statedb.go        |  4 +++-
 core/state/statedb_hooked.go | 12 ++++++++++--
 core/vm/interface.go         |  2 +-
 eth/tracers/firehose.go      | 26 ++++++++++++++++++++++++++
 4 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/core/state/statedb.go b/core/state/statedb.go
index 89132f828452..6b59930ffc6e 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -535,11 +535,13 @@ func (s *StateDB) SubBalance(addr common.Address, amount *uint256.Int, reason tr
 	return stateObject.SetBalance(new(uint256.Int).Sub(stateObject.Balance(), amount))
 }
 
-func (s *StateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) {
+func (s *StateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) (changed bool) {
 	stateObject := s.getOrNewStateObject(addr)
 	if stateObject != nil {
 		stateObject.SetBalance(amount)
+		return true
 	}
+	return false
 }
 
 func (s *StateDB) SetNonce(addr common.Address, nonce uint64, reason tracing.NonceChangeReason) {
diff --git a/core/state/statedb_hooked.go b/core/state/statedb_hooked.go
index 0aa4e885fdbf..91ad36f86621 100644
--- a/core/state/statedb_hooked.go
+++ b/core/state/statedb_hooked.go
@@ -61,8 +61,16 @@ func (s *hookedStateDB) GetBalance(addr common.Address) *uint256.Int {
 	return s.inner.GetBalance(addr)
 }
 
-func (s *hookedStateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) {
-	s.inner.SetBalance(addr, amount, reason)
+func (s *hookedStateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) (changed bool) {
+	var prev *uint256.Int
+	if s.hooks.OnBalanceChange != nil {
+		prev = s.inner.GetBalance(addr)
+	}
+	changed = s.inner.SetBalance(addr, amount, reason)
+	if changed && s.hooks.OnBalanceChange != nil {
+		s.hooks.OnBalanceChange(addr, prev.ToBig(), amount.ToBig(), reason)
+	}
+	return
 }
 
 func (s *hookedStateDB) GetNonce(addr common.Address) uint64 {
diff --git a/core/vm/interface.go b/core/vm/interface.go
index 038604b682ff..ae747a2a9766 100644
--- a/core/vm/interface.go
+++ b/core/vm/interface.go
@@ -37,7 +37,7 @@ type StateDB interface {
 	SubBalance(common.Address, *uint256.Int, tracing.BalanceChangeReason) uint256.Int
 	AddBalance(common.Address, *uint256.Int, tracing.BalanceChangeReason) uint256.Int
 	GetBalance(common.Address) *uint256.Int
-	SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason)
+	SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) bool
 
 	GetNonce(common.Address) uint64
 	SetNonce(common.Address, uint64, tracing.NonceChangeReason)
diff --git a/eth/tracers/firehose.go b/eth/tracers/firehose.go
index 65dcee0d3f84..4244fd4d3855 100644
--- a/eth/tracers/firehose.go
+++ b/eth/tracers/firehose.go
@@ -164,6 +164,7 @@ type Firehose struct {
 	transaction          *pbeth.TransactionTrace
 	transactionLogIndex  uint32
 	inSystemCall         bool
+	inSystemTx           bool
 	transactionIsolated  bool
 	transactionTransient *pbeth.TransactionTrace
 
@@ -518,6 +519,20 @@ func (f *Firehose) reorderCallOrdinals(call *pbeth.Call, ordinalBase uint64) (or
 	return call.EndOrdinal
 }
 
+func (f *Firehose) OnSystemTxStart() {
+	firehoseInfo("system tx start")
+	f.ensureInBlockAndNotInTrx()
+
+	f.inSystemTx = true
+}
+
+func (f *Firehose) OnSystemTxEnd() {
+	firehoseInfo("system tx end")
+	f.ensureInSystemTx()
+
+	f.inSystemTx = false
+}
+
 func (f *Firehose) OnSystemCallStart() {
 	firehoseInfo("system call start")
 	f.ensureInBlockAndNotInTrx()
@@ -1278,6 +1293,11 @@ func (f *Firehose) newBalanceChange(tag string, address common.Address, oldValue
 func (f *Firehose) OnNonceChange(a common.Address, prev, new uint64) {
 	f.ensureInBlockAndInTrx()
 
+	if *f.applyBackwardCompatibility && f.inSystemTx {
+		// Known Firehose issue: The nonce changes for system transactions are not recorded in the old Firehose instrumentation
+		return
+	}
+
 	activeCall := f.callStack.Peek()
 	change := &pbeth.NonceChange{
 		Address:  a.Bytes(),
@@ -1559,6 +1579,12 @@ func (f *Firehose) ensureInSystemCall() {
 	}
 }
 
+func (f *Firehose) ensureInSystemTx() {
+	if !f.inSystemTx {
+		f.panicInvalidState("caller expected to be in system transaction state but we were not, this is a bug", 2)
+	}
+}
+
 func (f *Firehose) panicInvalidState(msg string, callerSkip int) string {
 	caller := "N/A"
 	if _, file, line, ok := runtime.Caller(callerSkip); ok {