From c592bef3d1011ac5126df199fa2185b94ef235e1 Mon Sep 17 00:00:00 2001
From: Andrea Maria Piana <andrea.maria.piana@gmail.com>
Date: Tue, 4 Jun 2024 14:21:04 +0100
Subject: [PATCH] feat_: Add timeout to callprivaterpc endpoint

---
 rpc/call_raw.go      | 45 ++++++++++++++++++++++++++++++++++++++------
 rpc/call_raw_test.go | 29 +++++++++++++++++++++++-----
 2 files changed, 63 insertions(+), 11 deletions(-)

diff --git a/rpc/call_raw.go b/rpc/call_raw.go
index 42873334e..433702d12 100644
--- a/rpc/call_raw.go
+++ b/rpc/call_raw.go
@@ -3,6 +3,7 @@ package rpc
 import (
 	"context"
 	"encoding/json"
+	"time"
 
 	gethrpc "github.com/ethereum/go-ethereum/rpc"
 )
@@ -38,6 +39,7 @@ type jsonrpcRequest struct {
 	ChainID uint64          `json:"chainId"`
 	Method  string          `json:"method"`
 	Params  json.RawMessage `json:"params,omitempty"`
+	Timeout uint64          `json:"timeout,omitempty"`
 }
 
 type jsonrpcSuccessfulResponse struct {
@@ -111,11 +113,28 @@ func (c *Client) callBatchMethods(ctx context.Context, msgs json.RawMessage) str
 // callSingleMethod executes single JSON-RPC message and constructs proper response.
 func (c *Client) callSingleMethod(ctx context.Context, msg json.RawMessage) string {
 	// unmarshal JSON body into json-rpc request
-	chainID, method, params, id, err := methodAndParamsFromBody(msg)
+	jsonrpcParams, err := methodAndParamsFromBody(msg)
 	if err != nil {
-		return newErrorResponse(errInvalidMessageCode, err, id)
+		return newErrorResponse(errInvalidMessageCode, err, nil)
 	}
 
+	chainID := jsonrpcParams.ChainID
+	method := jsonrpcParams.Method
+	params := jsonrpcParams.Params
+	timeout := jsonrpcParams.Timeout
+	id := jsonrpcParams.ID
+
+	if timeout != 0 {
+		// TODO: remove me
+		c.log.Info("setting 50ms timeout", "method", method)
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Millisecond)
+		defer cancel()
+	}
+
+	// TODO: remove me
+	c.log.Info("calling method", "method", method)
+
 	if chainID == 0 {
 		chainID = c.UpstreamChainID
 	}
@@ -139,24 +158,38 @@ func (c *Client) callSingleMethod(ctx context.Context, msg json.RawMessage) stri
 	return newSuccessResponse(result, id)
 }
 
+type jsonrpcParameters struct {
+	ChainID uint64
+	Method  string
+	Params  []interface{}
+	Timeout uint64
+	ID      json.RawMessage
+}
+
 // methodAndParamsFromBody extracts Method and Params of
 // JSON-RPC body into values ready to use with ethereum-go's
 // RPC client Call() function. A lot of empty interface usage is
 // due to the underlying code design :/
-func methodAndParamsFromBody(body json.RawMessage) (uint64, string, []interface{}, json.RawMessage, error) {
+func methodAndParamsFromBody(body json.RawMessage) (*jsonrpcParameters, error) {
 	msg, err := unmarshalMessage(body)
 	if err != nil {
-		return 0, "", nil, nil, err
+		return nil, err
 	}
 	params := []interface{}{}
 	if msg.Params != nil {
 		err = json.Unmarshal(msg.Params, &params)
 		if err != nil {
-			return 0, "", nil, nil, err
+			return nil, err
 		}
 	}
 
-	return msg.ChainID, msg.Method, params, msg.ID, nil
+	return &jsonrpcParameters{
+		ChainID: msg.ChainID,
+		Method:  msg.Method,
+		Params:  params,
+		Timeout: msg.Timeout,
+		ID:      msg.ID,
+	}, nil
 }
 
 // unmarshalMessage tries to unmarshal JSON-RPC message.
diff --git a/rpc/call_raw_test.go b/rpc/call_raw_test.go
index bb75622c6..a776ba55e 100644
--- a/rpc/call_raw_test.go
+++ b/rpc/call_raw_test.go
@@ -59,6 +59,7 @@ func TestMethodAndParamsFromBody(t *testing.T) {
 		id         json.RawMessage
 		chainID    uint64
 		shouldFail bool
+		timeout    uint64
 	}{
 		{
 			"params_array",
@@ -73,6 +74,7 @@ func TestMethodAndParamsFromBody(t *testing.T) {
 			json.RawMessage(`42`),
 			0,
 			false,
+			0,
 		},
 		{
 			"params_empty_array",
@@ -82,6 +84,7 @@ func TestMethodAndParamsFromBody(t *testing.T) {
 			nil,
 			0,
 			false,
+			0,
 		},
 		{
 			"params_none",
@@ -91,6 +94,7 @@ func TestMethodAndParamsFromBody(t *testing.T) {
 			nil,
 			0,
 			false,
+			0,
 		},
 		{
 			"params_chain_id",
@@ -100,6 +104,7 @@ func TestMethodAndParamsFromBody(t *testing.T) {
 			nil,
 			2,
 			false,
+			0,
 		},
 		{
 			"getFilterMessage",
@@ -109,6 +114,7 @@ func TestMethodAndParamsFromBody(t *testing.T) {
 			json.RawMessage(`44`),
 			0,
 			false,
+			0,
 		},
 		{
 			"getFilterMessage_array",
@@ -118,6 +124,7 @@ func TestMethodAndParamsFromBody(t *testing.T) {
 			nil,
 			0,
 			true,
+			0,
 		},
 		{
 			"empty_array",
@@ -127,21 +134,33 @@ func TestMethodAndParamsFromBody(t *testing.T) {
 			nil,
 			0,
 			true,
+			0,
+		},
+		{
+			"timeout",
+			json.RawMessage(`{"jsonrpc": "2.0", "timeout": 2000, "method": "test"}`),
+			[]interface{}{},
+			"test",
+			nil,
+			0,
+			false,
+			2000,
 		},
 	}
 
 	for _, test := range cases {
 		t.Run(test.name, func(t *testing.T) {
-			chainID, method, params, id, err := methodAndParamsFromBody(test.body)
+			response, err := methodAndParamsFromBody(test.body)
 			if test.shouldFail {
 				require.Error(t, err)
 				return
 			}
 			require.NoError(t, err)
-			require.Equal(t, test.chainID, chainID)
-			require.Equal(t, test.method, method)
-			require.Equal(t, test.params, params)
-			require.EqualValues(t, test.id, id)
+			require.Equal(t, test.timeout, response.Timeout)
+			require.Equal(t, test.chainID, response.ChainID)
+			require.Equal(t, test.method, response.Method)
+			require.Equal(t, test.params, response.Params)
+			require.EqualValues(t, test.id, response.ID)
 		})
 	}
 }