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(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(s.mapAddress(key), value)
 	}
 	s.db.SetState(s.account, s.mapAddress(key), value)
 	return nil

system_tests/debugapi_test.go

@@ -3,6 +3,7 @@ package arbtest
 import (
 	"context"
 	"encoding/json"
+	"math/big"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -13,9 +14,18 @@ import (
 	"github.com/ethereum/go-ethereum/eth/tracers"
 	"github.com/ethereum/go-ethereum/rpc"
 
+	"github.com/offchainlabs/nitro/arbos/arbosState"
+	"github.com/offchainlabs/nitro/arbos/burn"
 	"github.com/offchainlabs/nitro/solgen/go/precompilesgen"
 )
 
+type account struct {
+	Balance *hexutil.Big                `json:"balance,omitempty"`
+	Code    []byte                      `json:"code,omitempty"`
+	Nonce   uint64                      `json:"nonce,omitempty"`
+	Storage map[common.Hash]common.Hash `json:"storage,omitempty"`
+}
+
 func TestDebugAPI(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -57,3 +67,54 @@ func TestDebugAPI(t *testing.T) {
 	err = l2rpc.CallContext(ctx, &result, "debug_traceTransaction", tx.Hash(), &tracers.TraceConfig{Tracer: &flatCallTracer})
 	Require(t, err)
 }
+
+func TestPrestateTracerRegistersArbitrumStorage(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	builder := NewNodeBuilder(ctx).DefaultConfig(t, false)
+	cleanup := builder.Build(t)
+	defer cleanup()
+
+	auth := builder.L2Info.GetDefaultTransactOpts("Owner", ctx)
+	arbOwner, err := precompilesgen.NewArbOwner(common.HexToAddress("0x70"), builder.L2.Client)
+	Require(t, err, "could not bind ArbOwner contract")
+
+	// Schedule a noop upgrade
+	tx, err := arbOwner.ScheduleArbOSUpgrade(&auth, 1, 1)
+	Require(t, err)
+	_, err = builder.L2.EnsureTxSucceeded(tx)
+	Require(t, err)
+
+	statedb, err := builder.L2.ExecNode.Backend.ArbInterface().BlockChain().State()
+	Require(t, err)
+	burner := burn.NewSystemBurner(nil, false)
+	arbosSt, err := arbosState.OpenArbosState(statedb, burner)
+	Require(t, err)
+
+	l2rpc := builder.L2.Stack.Attach()
+	var result map[common.Address]*account
+	prestateTracer := "prestateTracer"
+	err = l2rpc.CallContext(ctx, &result, "debug_traceTransaction", tx.Hash(), &tracers.TraceConfig{Tracer: &prestateTracer})
+	Require(t, err)
+
+	// ArbOSVersion and BrotliCompressionLevel storage slots should be accessed by arbos so we check if the current values of these appear in the prestateTrace
+	arbOSVersionHash := common.BigToHash(new(big.Int).SetUint64(arbosSt.ArbOSVersion()))
+	bcl, err := arbosSt.BrotliCompressionLevel()
+	Require(t, err)
+	bclHash := common.BigToHash(new(big.Int).SetUint64(bcl))
+
+	if _, ok := result[types.ArbosStateAddress]; !ok {
+		t.Fatal("ArbosStateAddress storage accesses not logged in the prestateTracer's trace")
+	}
+
+	found := 0
+	for _, val := range result[types.ArbosStateAddress].Storage {
+		if val == arbOSVersionHash || val == bclHash {
+			found++
+		}
+	}
+	if found != 2 {
+		t.Fatal("ArbosStateAddress storage accesses for ArbOSVersion and BrotliCompressionLevel not logged in the prestateTracer's trace")
+	}
+}

system_tests/precompile_test.go

@@ -5,6 +5,7 @@ package arbtest
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"math/big"
 	"sort"
@@ -539,6 +540,17 @@ func TestScheduleArbosUpgrade(t *testing.T) {
 		t.Errorf("expected completed scheduled upgrade to be ignored, got version %v timestamp %v", scheduled.ArbosVersion, scheduled.ScheduledForTimestamp)
 	}
 
+	l2rpc := builder.L2.Stack.Attach()
+	var result json.RawMessage
+	traceConfig := map[string]interface{}{
+		"tracer": "prestateTracer",
+		"tracerConfig": map[string]interface{}{
+			"diffMode": true,
+		},
+	}
+	err = l2rpc.CallContext(ctx, &result, "debug_traceTransaction", tx.Hash(), traceConfig)
+	Require(t, err)
+
 	// TODO: Once we have an ArbOS 30, test a real upgrade with it
 	// We can't test 11 -> 20 because 11 doesn't have the GetScheduledUpgrade method we want to test
 	var testVersion uint64 = 100