Hooking before the first IO instruction in replay.wasm

changelog/pmikolajczyk-nit-4401.md

@@ -0,0 +1,2 @@
+### Added
+ - Added a new hook to `replay.wasm` to enable an action just before the first IO (wavmio) instruction. It is expected that every `wasm` execution environment will provide a module `hooks` with a method `beforeFirstIO`. JIT and Arbitrator provers have noop implementations.

cmd/replay/main.go

@@ -277,7 +277,7 @@ func populateEcdsaCaches() {
 }
 
 func main() {
-	wavmio.StubInit()
+	wavmio.OnInit()
 	gethhook.RequireHookedGeth()
 
 	glogger := log.NewGlogHandler(
@@ -290,6 +290,7 @@ func main() {
 	raw := rawdb.NewDatabase(PreimageDb{})
 	db := state.NewDatabase(triedb.NewDatabase(raw, nil), nil)
 
+	wavmio.OnReady()
 	lastBlockHash := wavmio.GetLastBlockHash()
 
 	var lastBlockHeader *types.Header
@@ -449,5 +450,5 @@ func main() {
 	wavmio.SetLastBlockHash(newBlockHash)
 	wavmio.SetSendRoot(extraInfo.SendRoot)
 
-	wavmio.StubFinal()
+	wavmio.OnFinal()
 }

crates/jit/src/machine.rs

@@ -79,6 +79,9 @@ fn imports(store: &mut Store, func_env: &FunctionEnv<WasmEnv>) -> wasmer::Import
         "arbkeccak" => {
             "keccak256" => func!(arbkeccak::keccak256),
         },
+        "hooks" => {
+            "beforeFirstIO" => func!(|_: WasmEnvMut|{}),
+        },
         "wavmio" => {
             "getGlobalStateBytes32" => func!(wavmio::get_global_state_bytes32),
             "setGlobalStateBytes32" => func!(wavmio::set_global_state_bytes32),

crates/wasm-libraries/host-io/src/lib.rs

@@ -160,3 +160,7 @@ pub unsafe extern "C" fn wavmio__validateCertificate(preimage_type: u8, hash_ptr
     let result = wavm_validate_certificate(our_ptr, preimage_type);
     result
 }
+
+/// A hook called just before the first IO operation.
+#[no_mangle]
+pub unsafe extern "C" fn hooks__beforeFirstIO() {}

wavmio/higher.go

@@ -40,9 +40,13 @@ func readBuffer(f func(uint32, unsafe.Pointer) uint32) []byte {
 	}
 }
 
-func StubInit() {}
+func OnInit() {}
 
-func StubFinal() {}
+func OnReady() {
+	beforeFirstIO()
+}
+
+func OnFinal() {}
 
 func GetLastBlockHash() (hash common.Hash) {
 	hashUnsafe := unsafe.Pointer(&hash[0])

wavmio/stub.go → wavmio/native.go

@@ -75,7 +75,7 @@ func parsePreimageBytes(path string) {
 	}
 }
 
-func StubInit() {
+func OnInit() {
 	preimages = make(map[common.Hash][]byte)
 	var delayedMsgPath arrayFlags
 	seqMsgPosFlag := flag.Uint64("inbox-position", 0, "position for sequencer inbox message")
@@ -110,7 +110,9 @@ func StubInit() {
 	}
 }
 
-func StubFinal() {
+func OnReady() {}
+
+func OnFinal() {
 	log.Info("End state", "lastblockHash", lastBlockHash, "InboxPosition", seqMsgPos+seqAdvanced, "positionWithinMessage", posWithinMsg)
 }
 

wavmio/raw.go

@@ -30,3 +30,6 @@ func resolveTypedPreimage(ty uint32, hash unsafe.Pointer, offset uint32, output
 
 //go:wasmimport wavmio validateCertificate
 func validateCertificate(ty uint32, hash unsafe.Pointer) uint32
+
+//go:wasmimport hooks beforeFirstIO
+func beforeFirstIO()