Integrate CaptureArbitrumStorageGet/Set into the prestate tracer

arbos/storage/storage.go

@@ -129,7 +129,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
 }
@@ -166,7 +166,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
@@ -376,7 +376,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
 }
@@ -391,7 +391,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

@@ -52,8 +52,18 @@ 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
+	tracer.CaptureArbitrumStorageGet(info.Contract.Address(), key, mappedKey, info.Depth, info.Scenario == TracingBeforeEVM)
 	if info.Scenario == TracingDuringEVM {
 		scope := &vm.ScopeContext{
 			Memory:   vm.NewMemory(),
@@ -63,13 +73,21 @@ 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 {
-		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
+	tracer.CaptureArbitrumStorageSet(info.Contract.Address(), key, mappedKey, value, info.Depth, info.Scenario == TracingBeforeEVM)
 	if info.Scenario == TracingDuringEVM {
 		scope := &vm.ScopeContext{
 			Memory:   vm.NewMemory(),
@@ -79,8 +97,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 {
-		tracer.CaptureArbitrumStorageSet(key, value, info.Depth, info.Scenario == TracingBeforeEVM)
 	}
 }
 

execution/gethexec/stylus_tracer.go

@@ -204,8 +204,10 @@ func (t *stylusTracer) Stop(err error) {
 
 func (t *stylusTracer) CaptureArbitrumTransfer(from, to *common.Address, value *big.Int, before bool, purpose string) {
 }
-func (t *stylusTracer) CaptureArbitrumStorageGet(key common.Hash, depth int, before bool)        {}
-func (t *stylusTracer) CaptureArbitrumStorageSet(key, value common.Hash, depth int, before bool) {}
+func (t *stylusTracer) CaptureArbitrumStorageGet(addr common.Address, key, mappedKey common.Hash, depth int, before bool) {
+}
+func (t *stylusTracer) CaptureArbitrumStorageSet(addr common.Address, key, mappedKey, value common.Hash, depth int, before bool) {
+}
 func (t *stylusTracer) OnOpcode(pc uint64, opcode byte, gas, cost uint64, scope tracing.OpContext, rData []byte, depth int, err error) {
 }
 func (t *stylusTracer) OnFault(pc uint64, op byte, gas, cost uint64, _ tracing.OpContext, depth int, err error) {

system_tests/debugapi_test.go

@@ -1,20 +1,163 @@
 package arbtest
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"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
+	_, ok := result.Pre[manager]
+	assert(ok, nil, "manager address not found in pre section of trace")
+	assert(result.Pre[manager].ArbitrumStorage != nil, nil, "changes to arbitrum storage not picked up by prestate tracer")
+	_, ok = result.Pre[manager].ArbitrumStorage[codehash]
+	assert(ok, nil, "activated program's codehash key not found in the arbitrum storage trace entry for manager address in Pre")
+	preData := result.Pre[manager].ArbitrumStorage[codehash]
+
+	_, ok = result.Post[manager]
+	assert(ok, nil, "manager address not found in post section oftrace")
+	assert(result.Post[manager].ArbitrumStorage != nil, nil, "changes to arbitrum storage not picked up by prestate tracer")
+	_, ok = result.Post[manager].ArbitrumStorage[codehash]
+	assert(ok, nil, "activated program's codehash key not found in the arbitrum storage trace entry for manager address in Post")
+	postData := result.Post[manager].ArbitrumStorage[codehash]
+
+	// 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])
+	assert(uint64(hourActivatedFromTrace) == uint64(hourNow), nil, "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)
+	assert(hits-cost.GasWhenCached == miss-cost.Gas, err)
+	empty := len(ensure(mock.CacheProgram(&userAuth, program)).Logs)
+	evict := parseLog(ensure(mock.EvictProgram(&userAuth, program)).Logs[0])
+	cache := parseLog(ensure(mock.CacheProgram(&userAuth, program)).Logs[0])
+	assert(empty == 0 && evict.Manager == manager && !evict.Cached && cache.Codehash == codehash && cache.Cached, nil)
+}
+
 func TestDebugAPI(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()