Add OverwriteContractCode method to ArbDebug precompile

arbos/tx_processor.go

@@ -814,3 +814,7 @@ func (p *TxProcessor) IsCalldataPricingIncreaseEnabled() bool {
 	}
 	return enabled
 }
+
+func (p *TxProcessor) EVM() *vm.EVM {
+	return p.evm
+}

precompiles/ArbDebug.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/tracing"
 )
 
 // All calls to this precompile are authorized by the DebugPrecompile wrapper,
@@ -59,6 +60,11 @@ func (con ArbDebug) BecomeChainOwner(c ctx, evm mech) error {
 	return c.State.ChainOwners().Add(c.caller)
 }
 
+// Overwrite an existing contract's code
+func (con ArbDebug) OverwriteContractCode(c ctx, evm mech, addr addr, code []byte) ([]byte, error) {
+	return c.txProcessor.EVM().StateDB.SetCode(addr, code, tracing.CodeChangeUnspecified), nil
+}
+
 // Halts the chain by panicking in the STF
 func (con ArbDebug) Panic(c ctx, evm mech) error {
 	panic("called ArbDebug's debug-only Panic method")

system_tests/precompile_test.go

@@ -4,6 +4,7 @@
 package arbtest
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"math/big"
@@ -1312,3 +1313,55 @@ func TestArbAggregatorGetPreferredAggregator(t *testing.T) {
 		Fatal(t, "expected default preferred aggregator to be", l1pricing.BatchPosterAddress, "got", prefAgg)
 	}
 }
+
+func TestArbDebugOverwriteContractCode(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	builder := NewNodeBuilder(ctx).DefaultConfig(t, false)
+	cleanup := builder.Build(t)
+	defer cleanup()
+
+	// Become chain owner
+	auth := builder.L2Info.GetDefaultTransactOpts("Owner", ctx)
+	arbDebug, err := precompilesgen.NewArbDebug(types.ArbDebugAddress, builder.L2.Client)
+	Require(t, err)
+	tx, err := arbDebug.BecomeChainOwner(&auth)
+	Require(t, err)
+	_, err = builder.L2.EnsureTxSucceeded(tx)
+	Require(t, err)
+
+	// create EOA to test against
+	addr := common.BytesToAddress(crypto.Keccak256([]byte{})[:20])
+
+	// test that code is empty
+	code, err := builder.L2.Client.CodeAt(ctx, addr, nil)
+	Require(t, err)
+	if len(code) != 0 {
+		t.Fatal("expected code to be empty")
+	}
+
+	// overwrite with some code
+	testCodeA := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
+	tx, err = arbDebug.OverwriteContractCode(&auth, addr, testCodeA)
+	Require(t, err)
+	_, err = builder.L2.EnsureTxSucceeded(tx)
+	Require(t, err)
+	code, err = builder.L2.Client.CodeAt(ctx, addr, nil)
+	Require(t, err)
+	if !bytes.Equal(code, testCodeA) {
+		t.Fatal("expected code A to be", testCodeA, "got", code)
+	}
+
+	// overwrite with some other code
+	testCodeB := []byte{9, 8, 7, 6, 5, 4, 3, 2, 1, 0}
+	tx, err = arbDebug.OverwriteContractCode(&auth, addr, testCodeB)
+	Require(t, err)
+	_, err = builder.L2.EnsureTxSucceeded(tx)
+	Require(t, err)
+	code, err = builder.L2.Client.CodeAt(ctx, addr, nil)
+	Require(t, err)
+	if !bytes.Equal(code, testCodeB) {
+		t.Fatal("expected code B to be", testCodeB, "got", code)
+	}
+}