Multi-fragment activation: charge gas per fragment (#4312)

* resolve small PR review comments for stylus merge on activate

* Create kolbyml-nit-4406.md

* bump go-ethereum pin

* fix error from git merge

* update go ethereum pin

* Multi-fragment activation: charge gas per fragment

* Add changelog

* Attempt to resolve PR concerns

* Resolve concerns

* pushed file that wasn't saved

* initialize cost properly

* fix lint

* remove evmMemory gas

* update test

* fix lint

* try to make test more robust in ci

* Update programs.go

* Add BurnMultiGas to Burner interface

Close NIT-4455

* feedback from Tsahi

---------

Co-authored-by: Gabriel de Quadros Ligneul <gligneul@offchainlabs.com>
Co-authored-by: Tsahi Zidenberg <65945052+tsahee@users.noreply.github.com>

Author: 3 people

arbos/burn/burn.go

@@ -14,6 +14,7 @@ import (
 
 type Burner interface {
 	Burn(kind multigas.ResourceKind, amount uint64) error
+	BurnMultiGas(amount multigas.MultiGas) error
 	Burned() uint64
 	GasLeft() uint64 // `SystemBurner`s panic (no notion of GasLeft)
 	BurnOut() error
@@ -41,6 +42,11 @@ func (burner *SystemBurner) Burn(kind multigas.ResourceKind, amount uint64) erro
 	return nil
 }
 
+func (burner *SystemBurner) BurnMultiGas(amount multigas.MultiGas) error {
+	burner.gasBurnt.SaturatingAddInto(amount)
+	return nil
+}
+
 func (burner *SystemBurner) Burned() uint64 {
 	return burner.gasBurnt.SingleGas()
 }

arbos/programs/native.go

@@ -277,7 +277,7 @@ func getCompiledProgram(statedb vm.StateDB, moduleHash common.Hash, addressForLo
 	}
 
 	// addressForLogging may be empty or may not correspond to the code, so we need to be careful to use the code passed in separately
-	wasm, err := getWasmFromContractCode(statedb, code, params, false)
+	wasm, err := getWasmFromContractCode(statedb, code, params, nil)
 	if err != nil {
 		log.Error("Failed to reactivate program: getWasm", "address", addressForLogging, "expected moduleHash", moduleHash, "err", err)
 		return nil, fmt.Errorf("failed to reactivate program address: %v err: %w", addressForLogging, err)

arbos/programs/programs.go

@@ -6,10 +6,12 @@ package programs
 import (
 	"errors"
 	"fmt"
+	"math"
 	"math/big"
 
 	"github.com/ethereum/go-ethereum/arbitrum/multigas"
 	"github.com/ethereum/go-ethereum/common"
+	gethMath "github.com/ethereum/go-ethereum/common/math"
 	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/vm"
@@ -19,6 +21,7 @@ import (
 
 	"github.com/offchainlabs/nitro/arbcompress"
 	"github.com/offchainlabs/nitro/arbos/addressSet"
+	"github.com/offchainlabs/nitro/arbos/burn"
 	"github.com/offchainlabs/nitro/arbos/storage"
 	"github.com/offchainlabs/nitro/arbos/util"
 	"github.com/offchainlabs/nitro/arbutil"
@@ -113,7 +116,7 @@ func (p Programs) ActivateProgram(evm *vm.EVM, address common.Address, runCtx *c
 		// already activated and up to date
 		return 0, codeHash, common.Hash{}, nil, false, ProgramUpToDateError()
 	}
-	wasm, err := getWasm(statedb, address, params)
+	wasm, err := getWasm(statedb, address, params, burner)
 	if err != nil {
 		return 0, codeHash, common.Hash{}, nil, false, err
 	}
@@ -290,20 +293,31 @@ func attributeWasmComputation(contract *vm.Contract, startingGas uint64) {
 	}
 }
 
+// toWordSize returns the ceiled word size required for memory expansion.
+func ToWordSize(size uint64) uint64 {
+	if size > math.MaxUint64-31 {
+		return math.MaxUint64/32 + 1
+	}
+
+	return (size + 31) / 32
+}
+
 func evmMemoryCost(size uint64) uint64 {
 	// It would take 100GB to overflow this calculation, so no need to worry about that
-	words := (size + 31) / 32
+	words := ToWordSize(size)
 	linearCost := words * gethParams.MemoryGas
 	squareCost := (words * words) / gethParams.QuadCoeffDiv
 	return linearCost + squareCost
 }
 
-func getWasm(statedb vm.StateDB, program common.Address, params *StylusParams) ([]byte, error) {
+func getWasm(statedb vm.StateDB, program common.Address, params *StylusParams, burner burn.Burner) ([]byte, error) {
 	prefixedWasm := statedb.GetCode(program)
-	return getWasmFromContractCode(statedb, prefixedWasm, params, true)
+	return getWasmFromContractCode(statedb, prefixedWasm, params, burner)
 }
 
-func getWasmFromContractCode(statedb vm.StateDB, prefixedWasm []byte, params *StylusParams, isActivation bool) ([]byte, error) {
+// burner is used to charge gas for reading fragments. If it is present activation is assumed, and activation checks are enforced.
+// Only pass a burner if activating the program.
+func getWasmFromContractCode(statedb vm.StateDB, prefixedWasm []byte, params *StylusParams, burner burn.Burner) ([]byte, error) {
 	if len(prefixedWasm) == 0 {
 		return nil, ProgramNotWasmError()
 	}
@@ -314,7 +328,7 @@ func getWasmFromContractCode(statedb vm.StateDB, prefixedWasm []byte, params *St
 
 	if params.arbosVersion >= gethParams.ArbosVersion_StylusContractLimit {
 		if state.IsStylusRootProgramPrefix(prefixedWasm) {
-			return getWasmFromRootStylus(statedb, prefixedWasm, params.MaxWasmSize, params.MaxFragmentCount, isActivation)
+			return getWasmFromRootStylus(statedb, prefixedWasm, params.MaxWasmSize, params.MaxFragmentCount, burner)
 		}
 
 		if state.IsStylusFragmentPrefix(prefixedWasm) {
@@ -339,13 +353,15 @@ func getWasmFromClassicStylus(data []byte, maxSize uint32) ([]byte, error) {
 	return arbcompress.DecompressWithDictionary(wasm, int(maxSize), dict)
 }
 
-func getWasmFromRootStylus(statedb vm.StateDB, data []byte, maxSize uint32, maxFragments uint8, isActivation bool) ([]byte, error) {
+// burner is used to charge gas for reading fragments. If it is present activation is assumed, and activation checks are enforced.
+// Only pass a burner if activating the program.
+func getWasmFromRootStylus(statedb vm.StateDB, data []byte, maxSize uint32, maxFragments uint8, burner burn.Burner) ([]byte, error) {
 	root, err := state.NewStylusRoot(data)
 	if err != nil {
 		return nil, err
 	}
 
-	if isActivation {
+	if burner != nil {
 		if root.DecompressedLength > maxSize {
 			return nil, fmt.Errorf("invalid wasm: decompressedLength %d is greater then MaxWasmSize %d", root.DecompressedLength, maxSize)
 		}
@@ -361,6 +377,11 @@ func getWasmFromRootStylus(statedb vm.StateDB, data []byte, maxSize uint32, maxF
 	var compressedWasm []byte
 	for _, addr := range root.Addresses {
 		fragCode := statedb.GetCode(addr)
+		if burner != nil {
+			if err := chargeFragmentReadGas(burner, statedb, addr, uint64(len(fragCode))); err != nil {
+				return nil, err
+			}
+		}
 
 		payload, err := state.StripStylusFragmentPrefix(fragCode)
 		if err != nil {
@@ -387,6 +408,33 @@ func getWasmFromRootStylus(statedb vm.StateDB, data []byte, maxSize uint32, maxF
 	return wasm, nil
 }
 
+// chargeFragmentReadGas charges EXTCODECOPY-style gas for reading fragment code.
+func chargeFragmentReadGas(burner burn.Burner, statedb vm.StateDB, addr common.Address, codeSize uint64) error {
+	// charge access gas
+	var cost multigas.MultiGas
+	if statedb.AddressInAccessList(addr) {
+		cost = multigas.ComputationGas(gethParams.WarmStorageReadCostEIP2929)
+	} else {
+		statedb.AddAddressToAccessList(addr)
+		cost = multigas.StorageAccessGas(gethParams.ColdAccountAccessCostEIP2929)
+	}
+	// charge copy gas
+	words := ToWordSize(codeSize)
+	copyGas, overflow := gethMath.SafeMul(words, gethParams.CopyGas)
+	if overflow {
+		log.Trace("fragment copy gas overflow", "address", addr, "codeSize", codeSize, "words", words, "copyGas", gethParams.CopyGas)
+		return vm.ErrGasUintOverflow
+	}
+	if cost, overflow = cost.SafeIncrement(multigas.ResourceKindStorageAccess, copyGas); overflow {
+		log.Trace("fragment copy gas overflow", "address", addr, "codeSize", codeSize, "copyGas", copyGas)
+		return vm.ErrGasUintOverflow
+	}
+	if err := burner.BurnMultiGas(cost); err != nil {
+		return err
+	}
+	return nil
+}
+
 func getStylusCompressionDict(id byte) (arbcompress.Dictionary, error) {
 	switch id {
 	case 0:

arbos/programs/wasm.go

@@ -161,7 +161,9 @@ func handleProgramPrepare(statedb vm.StateDB, moduleHash common.Hash, addressFor
 			debugInt = 1
 		}
 
-		wasm, err := getWasmFromContractCode(statedb, code, params, true)
+		// Not an activation path here, so fragment read gas shouldn't be charged.
+		// Passing nil avoids charging gas through a storage-backed burner here.
+		wasm, err := getWasmFromContractCode(statedb, code, params, nil)
 		if err != nil {
 			panic(fmt.Sprintf("failed to get wasm for program, program address: %v, err: %v", addressForLogging.Hex(), err))
 		}

arbos/programs/wasmstorehelper.go

@@ -51,7 +51,9 @@ func (p Programs) SaveActiveProgramToWasmStore(statedb *state.StateDB, codeHash
 		return nil
 	}
 
-	wasm, err := getWasmFromContractCode(statedb, code, progParams, false)
+	// Not an activation path, so fragment read gas shouldn't be charged.
+	// Passing nil avoids charging gas through a storage-backed burner here.
+	wasm, err := getWasmFromContractCode(statedb, code, progParams, nil)
 	if err != nil {
 		log.Error("Failed to reactivate program while rebuilding wasm store: getWasmFromContractCode", "expected moduleHash", moduleHash, "err", err)
 		return fmt.Errorf("failed to reactivate program while rebuilding wasm store: %w", err)

changelog/gligneul-nit-4455.md

@@ -0,0 +1,2 @@
+### Internal
+- Add BurnMultiGas to Burner interface

changelog/kolbyml-nit-4405.md

@@ -0,0 +1,2 @@
+### Changed
+- Charge EXTCODECOPY-style gas per fragment during multi-fragment Stylus activation.

precompiles/context.go

@@ -49,6 +49,17 @@ func (c *Context) Burn(kind multigas.ResourceKind, amount uint64) error {
 	return nil
 }
 
+func (c *Context) BurnMultiGas(amount multigas.MultiGas) error {
+	if c.free {
+		return nil
+	}
+	if c.GasLeft() < amount.SingleGas() {
+		return c.BurnOut()
+	}
+	c.gasUsed.SaturatingAddInto(amount)
+	return nil
+}
+
 //nolint:unused
 func (c *Context) Burned() uint64 {
 	return c.gasUsed.SingleGas()

precompiles/context_test.go

@@ -68,3 +68,50 @@ func TestContextBurn(t *testing.T) {
 		t.Errorf("wrong computation: got %v, want %v", got, want)
 	}
 }
+
+func TestContextBurnMultiGas(t *testing.T) {
+	ctx := Context{
+		gasSupplied: 1_000,
+		gasUsed:     multigas.ZeroGas(),
+	}
+	if got, want := ctx.GasLeft(), uint64(1000); got != want {
+		t.Errorf("wrong gas left: got %v, want %v", got, want)
+	}
+	if got, want := ctx.Burned(), uint64(0); got != want {
+		t.Errorf("wrong gas burned: got %v, want %v", got, want)
+	}
+
+	gasToBurn := multigas.MultiGasFromPairs(
+		multigas.Pair{Kind: multigas.ResourceKindStorageAccess, Amount: 400},
+		multigas.Pair{Kind: multigas.ResourceKindStorageGrowth, Amount: 200},
+	)
+	if err := ctx.BurnMultiGas(gasToBurn); err != nil {
+		t.Errorf("unexpected error from burn: %v", err)
+	}
+	if got, want := ctx.GasLeft(), uint64(400); got != want {
+		t.Errorf("wrong gas left: got %v, want %v", got, want)
+	}
+	if got, want := ctx.Burned(), uint64(600); got != want {
+		t.Errorf("wrong gas burned: got %v, want %v", got, want)
+	}
+
+	if err := ctx.BurnMultiGas(gasToBurn); !errors.Is(err, vm.ErrOutOfGas) {
+		t.Errorf("wrong erro from burn: got %v, want %v", err, vm.ErrOutOfGas)
+	}
+	if got, want := ctx.GasLeft(), uint64(0); got != want {
+		t.Errorf("wrong gas left: got %v, want %v", got, want)
+	}
+	if got, want := ctx.Burned(), uint64(1000); got != want {
+		t.Errorf("wrong gas burned: got %v, want %v", got, want)
+	}
+
+	if got, want := ctx.gasUsed.Get(multigas.ResourceKindStorageAccess), uint64(400); got != want {
+		t.Errorf("wrong storage access: got %v, want %v", got, want)
+	}
+	if got, want := ctx.gasUsed.Get(multigas.ResourceKindStorageGrowth), uint64(200); got != want {
+		t.Errorf("wrong storage growth: got %v, want %v", got, want)
+	}
+	if got, want := ctx.gasUsed.Get(multigas.ResourceKindComputation), uint64(400); got != want {
+		t.Errorf("wrong computation: got %v, want %v", got, want)
+	}
+}

system_tests/stylus_contract_limit_increase_test.go

@@ -212,6 +212,39 @@ func TestFragmentedContractValidation(t *testing.T) {
 	}
 }
 
+func TestFragmentActivationChargesPerFragmentCodeRead(t *testing.T) {
+	builder, auth, cleanup := setupProgramTest(t, true, func(b *NodeBuilder) {
+		b.WithExtraArchs(allWasmTargets)
+		b.WithArbOSVersion(params.ArbosVersion_StylusContractLimit)
+	})
+	defer cleanup()
+
+	file := rustFile("storage")
+	fragmentsOne, _, _ := readFragmentedContractFile(t, file, 1)
+	fragmentsTwo, _, _ := readFragmentedContractFile(t, file, 2)
+	require.Len(t, fragmentsOne, 1)
+	require.Len(t, fragmentsTwo, 2)
+
+	minDelta := fragmentReadCostWarmOnly(uint64(len(fragmentsTwo[0]))) + fragmentReadCostWarmOnly(uint64(len(fragmentsTwo[1]))) - fragmentReadCostWarmOnly(uint64(len(fragmentsOne[0])))
+	maxDelta := fragmentReadCost(uint64(len(fragmentsTwo[0]))) + fragmentReadCost(uint64(len(fragmentsTwo[1]))) - fragmentReadCost(uint64(len(fragmentsOne[0])))
+	copyDelta := fragmentCopyCost(uint64(len(fragmentsTwo[0]))) + fragmentCopyCost(uint64(len(fragmentsTwo[1]))) - fragmentCopyCost(uint64(len(fragmentsOne[0])))
+	minDelta += copyDelta
+	maxDelta += copyDelta
+
+	_, _, receiptOne := deployAndActivateFragmentedContract(t, builder.ctx, auth, builder.L2.Client, deployConfig{
+		fragmentCount:    1,
+		expectActivation: true,
+	})
+	_, _, receiptTwo := deployAndActivateFragmentedContract(t, builder.ctx, auth, builder.L2.Client, deployConfig{
+		fragmentCount:    2,
+		expectActivation: true,
+	})
+
+	absDelta := absDiffUint64(receiptTwo.GasUsed, receiptOne.GasUsed)
+	require.GreaterOrEqual(t, absDelta, minDelta)
+	require.LessOrEqual(t, absDelta, maxDelta)
+}
+
 // Specific Edge Case Tests
 
 func TestThatWeCantActivateStylusFragmentContract(t *testing.T) {
@@ -636,6 +669,25 @@ func TestArbOwnerSetMaxFragmentCountFailsOnArbOS50(t *testing.T) {
 
 // Utils
 
+func fragmentReadCost(codeSize uint64) uint64 {
+	return params.WarmStorageReadCostEIP2929 + params.ColdAccountAccessCostEIP2929
+}
+
+func fragmentReadCostWarmOnly(codeSize uint64) uint64 {
+	return params.WarmStorageReadCostEIP2929
+}
+
+func fragmentCopyCost(codeSize uint64) uint64 {
+	return programs.ToWordSize(codeSize) * params.CopyGas
+}
+
+func absDiffUint64(left uint64, right uint64) uint64 {
+	if left >= right {
+		return left - right
+	}
+	return right - left
+}
+
 // readFragmentedContractFile reads, compiles, compresses, and fragments a contract.
 func readFragmentedContractFile(t *testing.T, file string, fragmentCount uint16) ([][]byte, []byte, arbcompress.Dictionary) {
 	t.Helper()