Specify address when caching Stylus program

arbos/arbosState/arbosstate.go

@@ -74,8 +74,8 @@ func OpenArbosState(stateDB vm.StateDB, burner burn.Burner) (*ArbosState, error)
 	}
 	return &ArbosState{
 		arbosVersion,
-		30,
-		30,
+		31,
+		31,
 		backingStorage.OpenStorageBackedUint64(uint64(upgradeVersionOffset)),
 		backingStorage.OpenStorageBackedUint64(uint64(upgradeTimestampOffset)),
 		backingStorage.OpenStorageBackedAddress(uint64(networkFeeAccountOffset)),
@@ -318,6 +318,9 @@ func (state *ArbosState) UpgradeArbosVersion(
 		case 30:
 			programs.Initialize(state.backingStorage.OpenSubStorage(programsSubspace))
 
+		case 31:
+			// no state changes needed
+
 		default:
 			return fmt.Errorf(
 				"the chain is upgrading to unsupported ArbOS version %v, %w",

arbos/programs/programs.go

@@ -342,10 +342,13 @@ func (p Programs) ProgramCached(codeHash common.Hash) (bool, error) {
 }
 
 // Sets whether a program is cached. Errors if trying to cache an expired program.
+// `address` must be present if setting cache to true as of ArbOS 31,
+// and if `address` is present it must have the specified codeHash.
 func (p Programs) SetProgramCached(
 	emitEvent func() error,
 	db vm.StateDB,
 	codeHash common.Hash,
+	address common.Address,
 	cache bool,
 	time uint64,
 	params *StylusParams,

precompiles/ArbWasmCache.go

@@ -3,6 +3,8 @@
 
 package precompiles
 
+import "github.com/ethereum/go-ethereum/common"
+
 type ArbWasmCache struct {
 	Address addr // 0x72
 
@@ -20,14 +22,23 @@ func (con ArbWasmCache) AllCacheManagers(c ctx, _ mech) ([]addr, error) {
 	return c.State.Programs().CacheManagers().AllMembers(65536)
 }
 
-// Caches all programs with the given codehash. Caller must be a cache manager or chain owner.
+// Deprecated: replaced with CacheProgram.
 func (con ArbWasmCache) CacheCodehash(c ctx, evm mech, codehash hash) error {
-	return con.setProgramCached(c, evm, codehash, true)
+	return con.setProgramCached(c, evm, common.Address{}, codehash, true)
+}
+
+// Caches all programs with a codehash equal to the given address. Caller must be a cache manager or chain owner.
+func (con ArbWasmCache) CacheProgram(c ctx, evm mech, address addr) error {
+	codehash, err := c.GetCodeHash(address)
+	if err != nil {
+		return err
+	}
+	return con.setProgramCached(c, evm, address, codehash, true)
 }
 
 // Evicts all programs with the given codehash. Caller must be a cache manager or chain owner.
 func (con ArbWasmCache) EvictCodehash(c ctx, evm mech, codehash hash) error {
-	return con.setProgramCached(c, evm, codehash, false)
+	return con.setProgramCached(c, evm, common.Address{}, codehash, false)
 }
 
 // Gets whether a program is cached. Note that the program may be expired.
@@ -36,7 +47,7 @@ func (con ArbWasmCache) CodehashIsCached(c ctx, evm mech, codehash hash) (bool,
 }
 
 // Caches all programs with the given codehash.
-func (con ArbWasmCache) setProgramCached(c ctx, evm mech, codehash hash, cached bool) error {
+func (con ArbWasmCache) setProgramCached(c ctx, evm mech, address addr, codehash hash, cached bool) error {
 	if !con.hasAccess(c) {
 		return c.BurnOut()
 	}
@@ -51,7 +62,7 @@ func (con ArbWasmCache) setProgramCached(c ctx, evm mech, codehash hash, cached
 		return con.UpdateProgramCache(c, evm, c.caller, codehash, cached)
 	}
 	return programs.SetProgramCached(
-		emitEvent, evm.StateDB, codehash, cached, evm.Context.Time, params, txRunMode, debugMode,
+		emitEvent, evm.StateDB, codehash, address, cached, evm.Context.Time, params, txRunMode, debugMode,
 	)
 }
 

precompiles/precompile.go

@@ -72,11 +72,12 @@ type Precompile struct {
 }
 
 type PrecompileMethod struct {
-	name         string
-	template     abi.Method
-	purity       purity
-	handler      reflect.Method
-	arbosVersion uint64
+	name            string
+	template        abi.Method
+	purity          purity
+	handler         reflect.Method
+	arbosVersion    uint64
+	maxArbosVersion uint64
 }
 
 type PrecompileEvent struct {
@@ -226,6 +227,7 @@ func MakePrecompile(metadata *bind.MetaData, implementer interface{}) (addr, *Pr
 			purity,
 			handler,
 			0,
+			0,
 		}
 		methods[id] = &method
 		methodsByName[name] = &method
@@ -575,6 +577,8 @@ func Precompiles() map[addr]ArbosPrecompile {
 	for _, method := range ArbWasmCache.methods {
 		method.arbosVersion = ArbWasmCache.arbosVersion
 	}
+	ArbWasmCache.methodsByName["CacheCodehash"].maxArbosVersion = params.ArbosVersion_Stylus
+	ArbWasmCache.methodsByName["CacheProgram"].arbosVersion = params.ArbosVersion_StylusFixes
 
 	ArbRetryableImpl := &ArbRetryableTx{Address: types.ArbRetryableTxAddress}
 	ArbRetryable := insert(MakePrecompile(pgen.ArbRetryableTxMetaData, ArbRetryableImpl))
@@ -680,7 +684,7 @@ func (p *Precompile) Call(
 	}
 	id := *(*[4]byte)(input)
 	method, ok := p.methods[id]
-	if !ok || arbosVersion < method.arbosVersion {
+	if !ok || arbosVersion < method.arbosVersion || (method.maxArbosVersion > 0 && arbosVersion > method.maxArbosVersion) {
 		// method does not exist or hasn't yet been activated
 		return nil, 0, vm.ErrExecutionReverted
 	}

precompiles/precompile_test.go

@@ -194,6 +194,7 @@ func TestPrecompilesPerArbosVersion(t *testing.T) {
 		11: 4,
 		20: 8,
 		30: 38,
+		31: 1,
 	}
 
 	precompiles := Precompiles()

solgen/gen.go

@@ -68,7 +68,7 @@ func main() {
 	}
 	root := filepath.Dir(filename)
 	parent := filepath.Dir(root)
-	filePaths, err := filepath.Glob(filepath.Join(parent, "contracts", "build", "contracts", "src", "*", "*", "*.json"))
+	filePaths, err := filepath.Glob(filepath.Join(parent, "contracts", "build", "contracts", "src", "*", "*.sol", "*.json"))
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -105,7 +105,7 @@ func main() {
 		modInfo.addArtifact(artifact)
 	}
 
-	yulFilePaths, err := filepath.Glob(filepath.Join(parent, "contracts", "out", "yul", "*", "*.json"))
+	yulFilePaths, err := filepath.Glob(filepath.Join(parent, "contracts", "out", "*", "*.yul", "*.json"))
 	if err != nil {
 		log.Fatal(err)
 	}

system_tests/program_test.go

@@ -1248,7 +1248,7 @@ func TestProgramCacheManager(t *testing.T) {
 	// check ownership
 	assert(arbOwner.IsChainOwner(nil, ownerAuth.From))
 	ensure(arbWasmCache.EvictCodehash(&ownerAuth, codehash))
-	ensure(arbWasmCache.CacheCodehash(&ownerAuth, codehash))
+	ensure(arbWasmCache.CacheProgram(&ownerAuth, program))
 
 	// de-authorize manager
 	ensure(arbOwner.RemoveWasmCacheManager(&ownerAuth, manager))