witness: only include elements if there's enough gas for it

cmd/evm/internal/t8ntool/execution.go

@@ -343,7 +343,7 @@ func (pre *Prestate) Apply(vmConfig vm.Config, chainConfig *params.ChainConfig,
 		// Amount is in gwei, turn into wei
 		amount := new(big.Int).Mul(new(big.Int).SetUint64(w.Amount), big.NewInt(params.GWei))
 		statedb.AddBalance(w.Address, amount)
-		statedb.Witness().TouchFullAccount(w.Address[:], true)
+		statedb.Witness().TouchFullAccount(w.Address[:], true, nil)
 	}
 	if chainConfig.IsVerkle(big.NewInt(int64(pre.Env.Number)), pre.Env.Timestamp) {
 		if err := overlay.OverlayVerkleTransition(statedb, common.Hash{}, chainConfig.OverlayStride); err != nil {

consensus/beacon/consensus.go

@@ -358,7 +358,7 @@ func (beacon *Beacon) Finalize(chain consensus.ChainHeaderReader, header *types.
 		state.AddBalance(w.Address, amount)
 
 		// The returned gas is not charged
-		state.Witness().TouchFullAccount(w.Address[:], true)
+		state.Witness().TouchFullAccount(w.Address[:], true, nil)
 	}
 
 	if chain.Config().IsVerkle(header.Number, header.Time) {

core/state/access_witness.go

@@ -18,7 +18,6 @@
 
 import (
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/common/math"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/holiman/uint256"
@@ -30,6 +29,9 @@
 // * the second bit is set if the branch has been read
 type mode byte
 
+// UseGasFn is a function that can be used to charge gas for a given amount.
+type UseGasFn func(uint64) bool
+
 const (
 	AccessWitnessReadFlag  = mode(1)
 	AccessWitnessWriteFlag = mode(2)
@@ -89,132 +91,130 @@
 	return naw
 }
 
-func (aw *AccessWitness) TouchFullAccount(addr []byte, isWrite bool) uint64 {
-	var gas uint64
+func (aw *AccessWitness) TouchFullAccount(addr []byte, isWrite bool, useGasFn UseGasFn) bool {
 	for i := utils.BasicDataLeafKey; i <= utils.CodeHashLeafKey; i++ {
-		gas += aw.touchAddressAndChargeGas(addr, zeroTreeIndex, byte(i), isWrite)
+		if _, ok := aw.touchAddressAndChargeGas(addr, zeroTreeIndex, byte(i), isWrite, useGasFn); !ok {
+			return false
+		}
 	}
-	return gas
+	return true
 }
 
-func (aw *AccessWitness) TouchAndChargeMessageCall(addr []byte) uint64 {
-	var gas uint64
-	gas += aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.BasicDataLeafKey, false)
-	return gas
+func (aw *AccessWitness) TouchAndChargeMessageCall(addr []byte, useGasFn UseGasFn) (uint64, bool) {
+	return aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.BasicDataLeafKey, false, useGasFn)
 }
 
-func (aw *AccessWitness) TouchAndChargeValueTransfer(callerAddr, targetAddr []byte) uint64 {
-	var gas uint64
-	gas += aw.touchAddressAndChargeGas(callerAddr, zeroTreeIndex, utils.BasicDataLeafKey, true)
-	gas += aw.touchAddressAndChargeGas(targetAddr, zeroTreeIndex, utils.BasicDataLeafKey, true)
-	return gas
+func (aw *AccessWitness) TouchAndChargeValueTransfer(callerAddr, targetAddr []byte, useGasFn UseGasFn) bool {
+	_, ok := aw.touchAddressAndChargeGas(callerAddr, zeroTreeIndex, utils.BasicDataLeafKey, true, useGasFn)
+	if !ok {
+		return false
+	}
+	_, ok = aw.touchAddressAndChargeGas(targetAddr, zeroTreeIndex, utils.BasicDataLeafKey, true, useGasFn)
+	return ok
 }
 
 // TouchAndChargeContractCreateCheck charges access costs before
 // a contract creation is initiated. It is just reads, because the
 // address collision is done before the transfer, and so no write
 // are guaranteed to happen at this point.
-func (aw *AccessWitness) TouchAndChargeContractCreateCheck(addr []byte) uint64 {
-	var gas uint64
-	gas += aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.BasicDataLeafKey, false)
-	gas += aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.CodeHashLeafKey, false)
-	return gas
+func (aw *AccessWitness) TouchAndChargeContractCreateCheck(addr []byte, useGasFn UseGasFn) bool {
+	if _, ok := aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.BasicDataLeafKey, false, useGasFn); !ok {
+		return false
+	}
+	_, ok := aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.CodeHashLeafKey, false, useGasFn)
+	return ok
 }
 
 // TouchAndChargeContractCreateInit charges access costs to initiate
 // a contract creation.
-func (aw *AccessWitness) TouchAndChargeContractCreateInit(addr []byte) uint64 {
-	var gas uint64
-	gas += aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.BasicDataLeafKey, true)
-	gas += aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.CodeHashLeafKey, true)
-	return gas
+func (aw *AccessWitness) TouchAndChargeContractCreateInit(addr []byte, useGasFn UseGasFn) bool {
+	if _, ok := aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.BasicDataLeafKey, true, useGasFn); !ok {
+		return false
+	}
+	_, ok := aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.CodeHashLeafKey, true, useGasFn)
+	return ok
 }
 
-func (aw *AccessWitness) TouchTxOriginAndComputeGas(originAddr []byte) uint64 {
+func (aw *AccessWitness) TouchTxOriginAndComputeGas(originAddr []byte) {
 	for i := utils.BasicDataLeafKey; i <= utils.CodeHashLeafKey; i++ {
-		aw.touchAddressAndChargeGas(originAddr, zeroTreeIndex, byte(i), i == utils.BasicDataLeafKey)
+		aw.touchAddressAndChargeGas(originAddr, zeroTreeIndex, byte(i), i == utils.BasicDataLeafKey, nil)
 	}
-
-	// Kaustinen note: we're currently experimenting with stop chargin gas for the origin address
-	// so simple transfer still take 21000 gas. This is to potentially avoid breaking existing tooling.
-	// This is the reason why we return 0 instead of `gas`.
-	// Note that we still have to touch the addresses to make sure the witness is correct.
-	return 0
 }
 
-func (aw *AccessWitness) TouchTxExistingAndComputeGas(targetAddr []byte, sendsValue bool) uint64 {
-	aw.touchAddressAndChargeGas(targetAddr, zeroTreeIndex, utils.BasicDataLeafKey, sendsValue)
-	aw.touchAddressAndChargeGas(targetAddr, zeroTreeIndex, utils.CodeHashLeafKey, false)
-
-	// Kaustinen note: we're currently experimenting with stop chargin gas for the origin address
-	// so simple transfer still take 21000 gas. This is to potentially avoid breaking existing tooling.
-	// This is the reason why we return 0 instead of `gas`.
-	// Note that we still have to touch the addresses to make sure the witness is correct.
-	return 0
+func (aw *AccessWitness) TouchTxExistingAndComputeGas(targetAddr []byte, sendsValue bool) {
+	aw.touchAddressAndChargeGas(targetAddr, zeroTreeIndex, utils.BasicDataLeafKey, sendsValue, nil)
+	aw.touchAddressAndChargeGas(targetAddr, zeroTreeIndex, utils.CodeHashLeafKey, false, nil)
 }
 
-func (aw *AccessWitness) TouchSlotAndChargeGas(addr []byte, slot common.Hash, isWrite bool) uint64 {
+func (aw *AccessWitness) TouchSlotAndChargeGas(addr []byte, slot common.Hash, isWrite bool, useGasFn UseGasFn) (uint64, bool) {
 	treeIndex, subIndex := utils.GetTreeKeyStorageSlotTreeIndexes(slot.Bytes())
-	return aw.touchAddressAndChargeGas(addr, *treeIndex, subIndex, isWrite)
-}
-
-func (aw *AccessWitness) touchAddressAndChargeGas(addr []byte, treeIndex uint256.Int, subIndex byte, isWrite bool) uint64 {
-	stemRead, selectorRead, stemWrite, selectorWrite, selectorFill := aw.touchAddress(addr, treeIndex, subIndex, isWrite)
-
-	var gas uint64
-	if stemRead {
-		gas += params.WitnessBranchReadCost
-	}
-	if selectorRead {
-		gas += params.WitnessChunkReadCost
-	}
-	if stemWrite {
-		gas += params.WitnessBranchWriteCost
-	}
-	if selectorWrite {
-		gas += params.WitnessChunkWriteCost
-	}
-	if selectorFill {
-		gas += params.WitnessChunkFillCost
-	}
-
-	return gas
+	return aw.touchAddressAndChargeGas(addr, *treeIndex, subIndex, isWrite, useGasFn)
 }
 
-// touchAddress adds any missing access event to the witness.
-func (aw *AccessWitness) touchAddress(addr []byte, treeIndex uint256.Int, subIndex byte, isWrite bool) (bool, bool, bool, bool, bool) {
+func (aw *AccessWitness) touchAddressAndChargeGas(addr []byte, treeIndex uint256.Int, subIndex byte, isWrite bool, useGasFn UseGasFn) (uint64, bool) {
 	branchKey := newBranchAccessKey(addr, treeIndex)
 	chunkKey := newChunkAccessKey(branchKey, subIndex)
 
 	// Read access.
 	var branchRead, chunkRead bool
 	if _, hasStem := aw.branches[branchKey]; !hasStem {
 		branchRead = true
-		aw.branches[branchKey] = AccessWitnessReadFlag
 	}
 	if _, hasSelector := aw.chunks[chunkKey]; !hasSelector {
 		chunkRead = true
-		aw.chunks[chunkKey] = AccessWitnessReadFlag
 	}
 
 	// Write access.
 	var branchWrite, chunkWrite, chunkFill bool
 	if isWrite {
 		if (aw.branches[branchKey] & AccessWitnessWriteFlag) == 0 {
 			branchWrite = true
-			aw.branches[branchKey] |= AccessWitnessWriteFlag
 		}
 
 		chunkValue := aw.chunks[chunkKey]
 		if (chunkValue & AccessWitnessWriteFlag) == 0 {
 			chunkWrite = true
-			aw.chunks[chunkKey] |= AccessWitnessWriteFlag
 		}
+	}
 
-		// TODO: charge chunk filling costs if the leaf was previously empty in the state
+	var gas uint64
+	if branchRead {
+		gas += params.WitnessBranchReadCost
+	}
+	if chunkRead {
+		gas += params.WitnessChunkReadCost
+	}
+	if branchWrite {
+		gas += params.WitnessBranchWriteCost
+	}
+	if chunkWrite {
+		gas += params.WitnessChunkWriteCost
+	}
+	if chunkFill {
+		gas += params.WitnessChunkFillCost
+	}
+
+	if useGasFn != nil {
+		if ok := useGasFn(gas); !ok {
+			return 0, false
+		}
+	}
+
+	if branchRead {
+		aw.branches[branchKey] = AccessWitnessReadFlag
+	}
+	if branchWrite {
+		aw.branches[branchKey] |= AccessWitnessWriteFlag
+	}
+	if chunkRead {
+		aw.chunks[chunkKey] = AccessWitnessReadFlag
+	}
+	if chunkWrite {
+		chunkWrite = true
+		aw.chunks[chunkKey] |= AccessWitnessWriteFlag
 	}
 
-	return branchRead, chunkRead, branchWrite, chunkWrite, chunkFill
+	return gas, true
 }
 
 type branchAccessKey struct {
@@ -242,15 +242,15 @@
 }
 
 // touchCodeChunksRangeOnReadAndChargeGas is a helper function to touch every chunk in a code range and charge witness gas costs
-func (aw *AccessWitness) TouchCodeChunksRangeAndChargeGas(contractAddr []byte, startPC, size uint64, codeLen uint64, isWrite bool) uint64 {
+func (aw *AccessWitness) TouchCodeChunksRangeAndChargeGas(contractAddr []byte, startPC, size uint64, codeLen uint64, isWrite bool, useGasFn UseGasFn) bool {
 	// note that in the case where the copied code is outside the range of the
 	// contract code but touches the last leaf with contract code in it,
 	// we don't include the last leaf of code in the AccessWitness.  The
 	// reason that we do not need the last leaf is the account's code size
 	// is already in the AccessWitness so a stateless verifier can see that
 	// the code from the last leaf is not needed.
 	if size == 0 || startPC >= codeLen {
-		return 0
+		return true
 	}
 
 	endPC := startPC + size
@@ -261,25 +261,21 @@
 		endPC -= 1 // endPC is the last bytecode that will be touched.
 	}
 
-	var statelessGasCharged uint64
 	for chunkNumber := startPC / 31; chunkNumber <= endPC/31; chunkNumber++ {
 		treeIndex := *uint256.NewInt((chunkNumber + 128) / 256)
 		subIndex := byte((chunkNumber + 128) % 256)
-		gas := aw.touchAddressAndChargeGas(contractAddr, treeIndex, subIndex, isWrite)
-		var overflow bool
-		statelessGasCharged, overflow = math.SafeAdd(statelessGasCharged, gas)
-		if overflow {
-			panic("overflow when adding gas")
+		if _, ok := aw.touchAddressAndChargeGas(contractAddr, treeIndex, subIndex, isWrite, useGasFn); !ok {
+			return false
 		}
 	}
 
-	return statelessGasCharged
+	return true
 }
 
-func (aw *AccessWitness) TouchBasicData(addr []byte, isWrite bool) uint64 {
-	return aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.BasicDataLeafKey, isWrite)
+func (aw *AccessWitness) TouchBasicData(addr []byte, isWrite bool, useGasFn UseGasFn) (uint64, bool) {
+	return aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.BasicDataLeafKey, isWrite, useGasFn)
 }
 
-func (aw *AccessWitness) TouchCodeHash(addr []byte, isWrite bool) uint64 {
-	return aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.CodeHashLeafKey, isWrite)
+func (aw *AccessWitness) TouchCodeHash(addr []byte, isWrite bool, useGasFn UseGasFn) (uint64, bool) {
+	return aw.touchAddressAndChargeGas(addr, zeroTreeIndex, utils.CodeHashLeafKey, isWrite, useGasFn)
 }

core/state_processor.go

@@ -177,7 +177,7 @@ func ApplyTransaction(config *params.ChainConfig, bc ChainContext, author *commo
 
 func InsertBlockHashHistoryAtEip2935Fork(statedb *state.StateDB, prevNumber uint64, prevHash common.Hash, chain consensus.ChainHeaderReader) {
 	// Make sure that the historical contract is added to the witness
-	statedb.Witness().TouchFullAccount(params.HistoryStorageAddress[:], true)
+	statedb.Witness().TouchFullAccount(params.HistoryStorageAddress[:], true, nil)
 
 	ancestor := chain.GetHeader(prevHash, prevNumber)
 	for i := prevNumber; i > 0 && i >= prevNumber-params.Eip2935BlockHashHistorySize; i-- {
@@ -191,5 +191,5 @@ func ProcessParentBlockHash(statedb *state.StateDB, prevNumber uint64, prevHash
 	var key common.Hash
 	binary.BigEndian.PutUint64(key[24:], ringIndex)
 	statedb.SetState(params.HistoryStorageAddress, key, prevHash)
-	statedb.Witness().TouchSlotAndChargeGas(params.HistoryStorageAddress[:], key, true)
+	statedb.Witness().TouchSlotAndChargeGas(params.HistoryStorageAddress[:], key, true, nil)
 }

core/state_transition.go

@@ -339,19 +339,6 @@ func (st *StateTransition) preCheck() error {
 	return st.buyGas()
 }
 
-// tryConsumeGas tries to subtract gas from gasPool, setting the result in gasPool
-// if subtracting more gas than remains in gasPool, set gasPool = 0 and return false
-// otherwise, do the subtraction setting the result in gasPool and return true
-func tryConsumeGas(gasPool *uint64, gas uint64) bool {
-	if *gasPool < gas {
-		*gasPool = 0
-		return false
-	}
-
-	*gasPool -= gas
-	return true
-}
-
 // TransitionDb will transition the state by applying the current message and
 // returning the evm execution result with following fields.
 //
@@ -406,16 +393,10 @@ func (st *StateTransition) TransitionDb() (*ExecutionResult, error) {
 		targetAddr := msg.To
 		originAddr := msg.From
 
-		statelessGasOrigin := st.evm.Accesses.TouchTxOriginAndComputeGas(originAddr.Bytes())
-		if !tryConsumeGas(&st.gasRemaining, statelessGasOrigin) {
-			return nil, fmt.Errorf("%w: Insufficient funds to cover witness access costs for transaction: have %d, want %d", ErrInsufficientBalanceWitness, st.gasRemaining, gas)
-		}
+		st.evm.Accesses.TouchTxOriginAndComputeGas(originAddr.Bytes())
 
 		if msg.To != nil {
-			statelessGasDest := st.evm.Accesses.TouchTxExistingAndComputeGas(targetAddr.Bytes(), msg.Value.Sign() != 0)
-			if !tryConsumeGas(&st.gasRemaining, statelessGasDest) {
-				return nil, fmt.Errorf("%w: Insufficient funds to cover witness access costs for transaction: have %d, want %d", ErrInsufficientBalanceWitness, st.gasRemaining, gas)
-			}
+			st.evm.Accesses.TouchTxExistingAndComputeGas(targetAddr.Bytes(), msg.Value.Sign() != 0)
 
 			// ensure the code size ends up in the access witness
 			st.evm.StateDB.GetCodeSize(*targetAddr)
@@ -472,7 +453,7 @@ func (st *StateTransition) TransitionDb() (*ExecutionResult, error) {
 
 		// add the coinbase to the witness iff the fee is greater than 0
 		if rules.IsEIP4762 && fee.Sign() != 0 {
-			st.evm.Accesses.TouchFullAccount(st.evm.Context.Coinbase[:], true)
+			st.evm.Accesses.TouchFullAccount(st.evm.Context.Coinbase[:], true, nil)
 		}
 	}
 

core/vm/common.go

@@ -92,3 +92,15 @@ func allZero(b []byte) bool {
 	}
 	return true
 }
+
+type gasConsumer struct {
+	availableGas uint64
+}
+
+func (gc *gasConsumer) consumeGas(gas uint64) bool {
+	if gc.availableGas < gas {
+		return false
+	}
+	gc.availableGas -= gas
+	return true
+}