Integrate CaptureArbitrumStorageGet/Set into the prestate tracer

arbos/storage/storage.go

@@ -130,7 +130,7 @@ func (s *Storage) Get(key common.Hash) (common.Hash, error) {
 		return common.Hash{}, err
 	}
 	if info := s.burner.TracingInfo(); info != nil {
-		info.RecordStorageGet(key)
+		info.RecordStorageGet(key, s.mapAddress(key))
 	}
 	return s.GetFree(key), nil
 }
@@ -167,7 +167,7 @@ func (s *Storage) Set(key common.Hash, value common.Hash) error {
 		return err
 	}
 	if info := s.burner.TracingInfo(); info != nil {
-		info.RecordStorageSet(key, value)
+		info.RecordStorageSet(key, s.mapAddress(key), value)
 	}
 	s.db.SetState(s.account, s.mapAddress(key), value)
 	return nil
@@ -377,7 +377,7 @@ func (ss *StorageSlot) Get() (common.Hash, error) {
 		return common.Hash{}, err
 	}
 	if info := ss.burner.TracingInfo(); info != nil {
-		info.RecordStorageGet(ss.slot)
+		info.RecordStorageGet(ss.slot, ss.slot)
 	}
 	return ss.db.GetState(ss.account, ss.slot), nil
 }
@@ -392,7 +392,7 @@ func (ss *StorageSlot) Set(value common.Hash) error {
 		return err
 	}
 	if info := ss.burner.TracingInfo(); info != nil {
-		info.RecordStorageSet(ss.slot, value)
+		info.RecordStorageSet(ss.slot, ss.slot, value)
 	}
 	ss.db.SetState(ss.account, ss.slot, value)
 	return nil

arbos/util/tracing.go

@@ -53,8 +53,20 @@ func NewTracingInfo(evm *vm.EVM, from, to common.Address, scenario TracingScenar
 	}
 }
 
-func (info *TracingInfo) RecordStorageGet(key common.Hash) {
+func (info *TracingInfo) RecordStorageGet(key, mappedKey common.Hash) {
 	tracer := info.Tracer
+	// RecordStorageGet is only called from arbos storage object, meaning the SLOAD opcode tracing in the scenario TracingDuringEVM
+	// has no impact at all as the caller address (scope.Contract) and the key being passed dont associate with each other, instead
+	// the key corresponds to types.ArbosStateAddress. For this exact reason, when RecordStorageGet is called from inside arbos storage,
+	// other tracers should implement their own CaptureArbitrumStorageGet functions irrespective of tracing scenario to remove dependency
+	// on opcode tracing.
+	// ...
+	// Since CaptureArbitrumStorageGet is only implemented for prestateTracer there's no harm in calling it along with opcode tracing
+	// during TracingDuringEVM scenario (as prestate tracer fails to record any meaningful change due to the reason given above),
+	// prestateTracer now supports in removing opcode tracing from this function but other tracers don't, and that should be changed
+	if tracer.CaptureArbitrumStorageGet != nil {
+		tracer.CaptureArbitrumStorageGet(info.Contract.Address(), key, mappedKey, info.Depth, info.Scenario == TracingBeforeEVM)
+	}
 	if info.Scenario == TracingDuringEVM {
 		scope := &vm.ScopeContext{
 			Memory:   vm.NewMemory(),
@@ -64,13 +76,23 @@ func (info *TracingInfo) RecordStorageGet(key common.Hash) {
 		if tracer.OnOpcode != nil {
 			tracer.OnOpcode(0, byte(vm.SLOAD), 0, 0, scope, []byte{}, info.Depth, nil)
 		}
-	} else if tracer.CaptureArbitrumStorageGet != nil {
-		tracer.CaptureArbitrumStorageGet(key, info.Depth, info.Scenario == TracingBeforeEVM)
 	}
 }
 
-func (info *TracingInfo) RecordStorageSet(key, value common.Hash) {
+func (info *TracingInfo) RecordStorageSet(key, mappedKey, value common.Hash) {
 	tracer := info.Tracer
+	// RecordStorageSet is only called from arbos storage object, meaning the SSTORE opcode tracing in the scenario TracingDuringEVM
+	// has no impact at all as the caller address (scope.Contract) and the key being passed dont associate with each other, instead
+	// the key corresponds to types.ArbosStateAddress. For this exact reason, as RecordStorageSet is called from inside arbos storage,
+	// other tracers should implement their own CaptureArbitrumStorageSet functions irrespective of tracing scenario to remove dependency
+	// on opcode tracing.
+	// ...
+	// Since CaptureArbitrumStorageSet is only implemented for prestateTracer there's no harm in calling it along with opcode tracing
+	// during TracingDuringEVM scenario (as prestate tracer fails to record any meaningful change due to the reason given above),
+	// prestateTracer now supports in removing opcode tracing from this function but other tracers don't, and that should be changed
+	if tracer.CaptureArbitrumStorageSet != nil {
+		tracer.CaptureArbitrumStorageSet(info.Contract.Address(), key, mappedKey, value, info.Depth, info.Scenario == TracingBeforeEVM)
+	}
 	if info.Scenario == TracingDuringEVM {
 		scope := &vm.ScopeContext{
 			Memory:   vm.NewMemory(),
@@ -80,8 +102,6 @@ func (info *TracingInfo) RecordStorageSet(key, value common.Hash) {
 		if tracer.OnOpcode != nil {
 			tracer.OnOpcode(0, byte(vm.SSTORE), 0, 0, scope, []byte{}, info.Depth, nil)
 		}
-	} else if tracer.CaptureArbitrumStorageSet != nil {
-		tracer.CaptureArbitrumStorageSet(key, value, info.Depth, info.Scenario == TracingBeforeEVM)
 	}
 }
 

system_tests/debugapi_test.go

@@ -1,21 +1,178 @@
 package arbtest
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
+	"math"
 	"testing"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/eth"
 	"github.com/ethereum/go-ethereum/eth/tracers"
 	"github.com/ethereum/go-ethereum/rpc"
 
+	"github.com/offchainlabs/nitro/arbos/programs"
+	"github.com/offchainlabs/nitro/solgen/go/mocksgen"
 	"github.com/offchainlabs/nitro/solgen/go/precompilesgen"
+	pgen "github.com/offchainlabs/nitro/solgen/go/precompilesgen"
+	"github.com/offchainlabs/nitro/util/arbmath"
 )
 
+type account struct {
+	Balance         *hexutil.Big                `json:"balance,omitempty"`
+	Code            hexutil.Bytes               `json:"code,omitempty"`
+	Nonce           uint64                      `json:"nonce,omitempty"`
+	Storage         map[common.Hash]common.Hash `json:"storage,omitempty"`
+	ArbitrumStorage map[common.Hash]common.Hash `json:"arbitrumStorage,omitempty"`
+}
+type prestateTrace struct {
+	Post map[common.Address]*account `json:"post"`
+	Pre  map[common.Address]*account `json:"pre"`
+}
+
+func TestPrestateTracerArbitrumStorage(t *testing.T) {
+	builder, ownerAuth, cleanup := setupProgramTest(t, true)
+	ctx := builder.ctx
+	l2client := builder.L2.Client
+	l2info := builder.L2Info
+	defer cleanup()
+
+	ensure := func(tx *types.Transaction, err error) *types.Receipt {
+		t.Helper()
+		Require(t, err)
+		receipt, err := EnsureTxSucceeded(ctx, l2client, tx)
+		Require(t, err)
+		return receipt
+	}
+	assert := func(cond bool, err error, msg ...interface{}) {
+		t.Helper()
+		Require(t, err)
+		if !cond {
+			Fatal(t, msg...)
+		}
+	}
+
+	// precompiles we plan to use
+	arbWasm, err := pgen.NewArbWasm(types.ArbWasmAddress, builder.L2.Client)
+	Require(t, err)
+	arbWasmCache, err := pgen.NewArbWasmCache(types.ArbWasmCacheAddress, builder.L2.Client)
+	Require(t, err)
+	arbOwner, err := pgen.NewArbOwner(types.ArbOwnerAddress, builder.L2.Client)
+	Require(t, err)
+	ensure(arbOwner.SetInkPrice(&ownerAuth, 10_000))
+	parseLog := logParser[pgen.ArbWasmCacheUpdateProgramCache](t, pgen.ArbWasmCacheABI, "UpdateProgramCache")
+
+	// fund a user account we'll use to probe access-restricted methods
+	l2info.GenerateAccount("Anyone")
+	userAuth := l2info.GetDefaultTransactOpts("Anyone", ctx)
+	userAuth.GasLimit = 3e6
+	TransferBalance(t, "Owner", "Anyone", arbmath.BigMulByUint(oneEth, 32), l2info, l2client, ctx)
+
+	// deploy without activating a wasm
+	wasm, _ := readWasmFile(t, rustFile("keccak"))
+	program := deployContract(t, ctx, userAuth, l2client, wasm)
+	codehash := crypto.Keccak256Hash(wasm)
+
+	// athorize the manager
+	manager, tx, mock, err := mocksgen.DeploySimpleCacheManager(&ownerAuth, l2client)
+	ensure(tx, err)
+	isManager, err := arbWasmCache.IsCacheManager(nil, manager)
+	assert(!isManager, err)
+	ensure(arbOwner.AddWasmCacheManager(&ownerAuth, manager))
+	assert(arbWasmCache.IsCacheManager(nil, manager))
+	all, err := arbWasmCache.AllCacheManagers(nil)
+	assert(len(all) == 1 && all[0] == manager, err)
+
+	// cache the active program
+	activateWasm(t, ctx, userAuth, l2client, program, "keccak")
+	cacheTx, err := mock.CacheProgram(&userAuth, program)
+	ensure(cacheTx, err)
+	assert(arbWasmCache.CodehashIsCached(nil, codehash))
+
+	l2rpc := builder.L2.Stack.Attach()
+
+	var result prestateTrace
+	traceConfig := map[string]interface{}{
+		"tracer": "prestateTracer",
+		"tracerConfig": map[string]interface{}{
+			"diffMode": true,
+		},
+	}
+	err = l2rpc.CallContext(ctx, &result, "debug_traceTransaction", cacheTx.Hash(), traceConfig)
+	Require(t, err)
+
+	// Validate trace result
+	validate := func(traceMap map[common.Address]*account, kind string) common.Hash {
+		_, ok := traceMap[manager]
+		assert(ok, nil, fmt.Sprintf("manager address not found in %s section of trace", kind))
+		assert(traceMap[manager].ArbitrumStorage != nil, nil, "changes to arbitrum storage not picked up by prestate tracer")
+		_, ok = traceMap[manager].ArbitrumStorage[codehash]
+		assert(ok, nil, fmt.Sprintf("activated program's codehash key not found in the arbitrum storage trace entry for manager address in %s", kind))
+		return traceMap[manager].ArbitrumStorage[codehash]
+	}
+	preData := validate(result.Pre, "pre")
+	postData := validate(result.Post, "post")
+
+	// since we are just caching the program the only thing that should differ between the pre and post values is the cached byte
+	assert(!(preData == postData), nil, "preData and postData shouldnt be equal")
+	assert(bytes.Equal(preData[:14], postData[:14]), nil, "preData and postData should only differ in cached byte")
+	assert(bytes.Equal(preData[15:], postData[15:]), nil, "preData and postData should only differ in cached byte")
+	assert(!arbmath.BytesToBool(preData[14:15]), nil, "cached byte of preData should be false")
+	assert(arbmath.BytesToBool(postData[14:15]), nil, "cached byte of postData should be true")
+
+	version, err := arbWasm.StylusVersion(nil)
+	assert(arbmath.BytesToUint16(postData[:2]) == version, err, "stylus version mismatch")
+
+	programMemoryFootprint, err := arbWasm.ProgramMemoryFootprint(nil, program)
+	assert(arbmath.BytesToUint16(postData[6:8]) == programMemoryFootprint, err, "programMemoryFootprint mismatch")
+
+	codehashAsmSize, err := arbWasm.CodehashAsmSize(nil, codehash)
+	codehashAsmSizeFromTrace := arbmath.SaturatingUMul(arbmath.BytesToUint24(postData[11:14]).ToUint32(), 1024)
+	assert(codehashAsmSizeFromTrace == codehashAsmSize, err, "codehashAsmSize mismatch")
+
+	hourNow := (time.Now().Unix() - programs.ArbitrumStartTime) / 3600
+	hourActivatedFromTrace := arbmath.BytesToUint24(postData[8:11])
+	// #nosec G115
+	if !(uint64(hourActivatedFromTrace) == uint64(hourNow)) {
+		// Although very low but there's a chance that this assert might fail with hourNow being off by one
+		// from hourActivatedFromTrace considering the time that passed between the program activation and now.
+		// #nosec G115
+		if !(math.Abs(float64(hourActivatedFromTrace)-float64(hourNow)) != 1) {
+			Fatal(t, "wrong activated time in trace")
+		}
+	}
+
+	// Compare gas costs
+	keccak := func() uint64 {
+		tx := l2info.PrepareTxTo("Owner", &program, 1e9, nil, []byte{0x00})
+		return ensure(tx, l2client.SendTransaction(ctx, tx)).GasUsedForL2()
+	}
+	ensure(mock.EvictProgram(&userAuth, program))
+	miss := keccak()
+	ensure(mock.CacheProgram(&userAuth, program))
+	hits := keccak()
+	cost, err := arbWasm.ProgramInitGas(nil, program)
+	// We check that GasUsedForL2 contains correct portion of gas usage wrt when a program is cached vs when it isn't
+	assert(hits-cost.GasWhenCached == miss-cost.Gas, err)
+
+	// When a program is already cached, no logs are returned upon caching it again
+	empty := len(ensure(mock.CacheProgram(&userAuth, program)).Logs)
+	assert(empty == 0, nil)
+
+	evict := parseLog(ensure(mock.EvictProgram(&userAuth, program)).Logs[0])
+	assert(evict.Manager == manager && !evict.Cached, nil)
+
+	cache := parseLog(ensure(mock.CacheProgram(&userAuth, program)).Logs[0])
+	assert(cache.Codehash == codehash && cache.Cached, nil)
+}
+
 func TestDebugAPI(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

system_tests/precompile_test.go

@@ -434,12 +434,15 @@ func TestGasAccountingParams(t *testing.T) {
 	Require(t, err)
 	arbGasInfoSpeedLimit, arbGasInfoPoolSize, arbGasInfoTxGasLimit, err := arbGasInfo.GetGasAccountingParams(&bind.CallOpts{Context: ctx})
 	Require(t, err)
+	// #nosec G115
 	if arbGasInfoSpeedLimit.Cmp(big.NewInt(int64(speedLimit))) != 0 {
 		Fatal(t, "expected speed limit to be", speedLimit, "got", arbGasInfoSpeedLimit)
 	}
+	// #nosec G115
 	if arbGasInfoPoolSize.Cmp(big.NewInt(int64(txGasLimit))) != 0 {
 		Fatal(t, "expected pool size to be", txGasLimit, "got", arbGasInfoPoolSize)
 	}
+	// #nosec G115
 	if arbGasInfoTxGasLimit.Cmp(big.NewInt(int64(txGasLimit))) != 0 {
 		Fatal(t, "expected tx gas limit to be", txGasLimit, "got", arbGasInfoTxGasLimit)
 	}