Add stylus contract storage to trace

cmd/wazero/wazero.go

@@ -210,13 +210,17 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer) int {
 	}
 
 	// TODO: maybe better usage description?
-	var stylusTracePath string
-	flags.StringVar(&stylusTracePath, "stylus", "",
-		"Imports the EVM hook functions and mocks their IO according the result of debug_traceTransaction in the path provided.")
+	var stylusTxHash string
+	flags.StringVar(&stylusTxHash, "stylus", "",
+		"Imports the EVM hook functions and mocks their IO using debug_traceTransaction for the given transaction hash.")
+
+	var stylusRpcUrl string
+	flags.StringVar(&stylusRpcUrl, "stylus-rpc", "http://localhost:8547",
+		"RPC endpoint used when fetching Stylus traces.")
 
 	var traceDir string
 	flags.StringVar(&traceDir, "trace-dir", "",
-		"Directory where to save the trace record. If empty - no trace is produced. Default \"\".")
+		"Directory where to save the trace record. If empty - no trace is produced. (default \"\")")
 
 	cacheDir := cacheDirFlag(flags)
 
@@ -342,9 +346,9 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer) int {
 		traceRecordPtr = &traceRecord
 	}
 
-	var stylusState *stylus.StylusTrace
-	if stylusTracePath != "" {
-		stylusState, err = stylus.Instantiate(ctx, rt, stylusTracePath, traceRecordPtr)
+	var stylusState *stylus.StylusState
+	if stylusTxHash != "" {
+		stylusState, err = stylus.Instantiate(ctx, rt, stylusRpcUrl, stylusTxHash, traceRecordPtr)
 		if err != nil {
 			fmt.Fprintf(stdErr, "error reading stylus trace: %v\n", err)
 			return 1
@@ -382,7 +386,7 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer) int {
 		return 1
 	}
 
-	if stylusTracePath != "" {
+	if stylusTxHash != "" {
 		arg, err := stylusState.GetEntrypointArg()
 		if err != nil {
 			fmt.Fprintf(stdErr, "error reading stylus entrypoint argument: %v\n", err)

internal/stylus/json_rpc.go

@@ -0,0 +1,66 @@
+package stylus
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/http"
+)
+
+type RpcClient struct {
+	url string
+	id  int
+}
+
+func NewRpcClient(url string) *RpcClient {
+	return &RpcClient{url: url, id: 0}
+}
+
+func (client *RpcClient) Request(method string, params any) (json.RawMessage, error) {
+	data, err := json.Marshal(rpcRequest{
+		Jsonrpc: "2.0",
+		Method:  method,
+		Params:  params,
+		Id:      client.id,
+	})
+	if err != nil {
+		return nil, err
+	}
+	client.id++
+
+	resp, err := http.Post(client.url, "application/json", bytes.NewBuffer(data))
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	var rpcResp rpcResponse
+	if err := json.NewDecoder(resp.Body).Decode(&rpcResp); err != nil {
+		return nil, err
+	}
+
+	if rpcResp.Error != nil {
+		return nil, fmt.Errorf("JSON RPC returned error: %v", *rpcResp.Error)
+	}
+
+	return rpcResp.Result, nil
+}
+
+type rpcRequest struct {
+	Jsonrpc string      `json:"jsonrpc"`
+	Method  string      `json:"method"`
+	Params  interface{} `json:"params,omitempty"`
+	Id      int         `json:"id"`
+}
+
+type rpcResponse struct {
+	Jsonrpc string          `json:"jsonrpc"`
+	Result  json.RawMessage `json:"result,omitempty"`
+	Error   *rpcError       `json:"error,omitempty"`
+	Id      int             `json:"id"`
+}
+
+type rpcError struct {
+	Code    int    `json:"code"`
+	Message string `json:"message"`
+}

internal/stylus/stylus.go

@@ -3,30 +3,125 @@ package stylus
 import (
 	"context"
 	"encoding/json"
-	"os"
+	"fmt"
 
 	"github.com/metacraft-labs/trace_record"
 	"github.com/tetratelabs/wazero"
 )
 
-func Instantiate(ctx context.Context, r wazero.Runtime, stylusTracePath string, record *trace_record.TraceRecord) (*StylusTrace, error) {
-	stylusTraceJson, err := os.ReadFile(stylusTracePath)
+func Instantiate(ctx context.Context, r wazero.Runtime, rpcUrl string, txHash string, record *trace_record.TraceRecord) (result *StylusState, err error) {
+	rpcClient := NewRpcClient(rpcUrl)
+
+	result = &StylusState{}
+
+	result.txHash = txHash
+
+	reciept, err := rpcClient.requestTxReciept(txHash)
 	if err != nil {
-		return nil, err
+		return
+	}
+
+	if reciept.To == nil {
+		err = fmt.Errorf("the \"to\" of the transaction is null")
+		return
 	}
 
-	stylusState := StylusTrace{}
+	result.address = *reciept.To
+	result.blockNumber = reciept.BlockNumber
+	result.blockHash = reciept.BlockHash
+	result.txIndex = reciept.TransactionIndex
 
-	if err := json.Unmarshal(stylusTraceJson, &stylusState.events); err != nil {
-		return nil, err
+	result.evmEvents, err = rpcClient.requestDebugTraceTransaction(txHash)
+	if err != nil {
+		return
 	}
 
+	_, _ = rpcClient.requestAlteredStorageInBlockBeforeTransaction(result.blockHash, result.txIndex, result.address)
+
 	moduleBuilder := r.NewHostModuleBuilder("vm_hooks")
-	moduleBuilder = exportSylusFunctions(moduleBuilder, &stylusState, record)
+	moduleBuilder = exportSylusFunctions(moduleBuilder, result, record)
 
 	if _, err := moduleBuilder.Instantiate(ctx); err != nil {
 		return nil, err
 	}
 
-	return &stylusState, nil
+	return
+}
+
+type TxReceipt struct {
+	BlockNumber      string  `json:"blockNumber"`
+	BlockHash        string  `json:"blockHash"`
+	To               *string `json:"to"`
+	TransactionIndex string  `json:"transactionIndex"`
+}
+
+func (client *RpcClient) requestTxReciept(txHash string) (result TxReceipt, err error) {
+	rawTx, err := client.Request("eth_getTransactionReceipt", []interface{}{txHash})
+	if err != nil {
+		return
+	}
+
+	if err = json.Unmarshal(rawTx, &result); err != nil {
+		return
+	}
+
+	return
+}
+
+func (client *RpcClient) requestDebugTraceTransaction(txHash string) ([]evmEvent, error) {
+	rawJson, err := client.Request(
+		"debug_traceTransaction",
+		[]interface{}{
+			txHash,
+			struct {
+				Tracer string `json:"tracer"`
+			}{
+				Tracer: "stylusTracer",
+			},
+		},
+	)
+
+	if err != nil {
+		return nil, err
+	}
+
+	var res []evmEvent
+	json.Unmarshal(rawJson, &res)
+
+	return res, nil
+}
+
+func (client *RpcClient) requestAlteredStorageInBlockBeforeTransaction(blockHash string, txIndex string, address string) (interface{}, error) {
+	fmt.Printf("%#v\n", []interface{}{
+		blockHash,
+		1,
+		address,
+		"0x0000000000000000000000000000000000000000000000000000000000000000",
+		4096, // TODO: discuss if this is enough
+	})
+	rawJson, err := client.Request(
+		"debug_storageRangeAt",
+		[]interface{}{
+			blockHash,
+			1,
+			address,
+			"0x0000000000000000000000000000000000000000000000000000000000000000",
+			4096, // TODO: discuss if this is enough
+		},
+	)
+
+	if err != nil {
+		fmt.Printf("FAIL: %v\n", err)
+		return nil, err
+	}
+
+	tmp, err := json.Marshal(rawJson)
+
+	fmt.Printf("OK: %s\n", string(tmp))
+
+	if err != nil {
+		return nil, err
+	}
+
+	return rawJson, nil
 }