Merge pull request ethereum#192 from OffchainLabs/rpc-client-hooks

Add request and response hooks to rpc.Client

Author: PlasmaPower

rpc/client.go

@@ -74,6 +74,8 @@ type BatchElem struct {
 
 // Client represents a connection to an RPC server.
 type Client struct {
+	requestHook RequestHook
+
 	idgen    func() ID // for subscriptions
 	isHTTP   bool      // connection type: http, ws or ipc
 	services *serviceRegistry
@@ -300,25 +302,32 @@ func (c *Client) CallContext(ctx context.Context, result interface{}, method str
 	}
 	op := &requestOp{ids: []json.RawMessage{msg.ID}, resp: make(chan *jsonrpcMessage, 1)}
 
+	resultHook := c.onRequest(msg)
 	if c.isHTTP {
 		err = c.sendHTTP(ctx, op, msg)
 	} else {
 		err = c.send(ctx, op, msg)
 	}
 	if err != nil {
+		resultHook.OnResult(nil, err)
 		return err
 	}
 
 	// dispatch has accepted the request and will close the channel when it quits.
 	switch resp, err := op.wait(ctx, c); {
 	case err != nil:
+		resultHook.OnResult(resp, err)
 		return err
 	case resp.Error != nil:
+		resultHook.OnResult(resp, resp.Error)
 		return resp.Error
 	case len(resp.Result) == 0:
+		resultHook.OnResult(resp, ErrNoResult)
 		return ErrNoResult
 	default:
-		return json.Unmarshal(resp.Result, &result)
+		err := json.Unmarshal(resp.Result, &result)
+		resultHook.OnResult(resp, err)
+		return err
 	}
 }
 
@@ -362,6 +371,12 @@ func (c *Client) BatchCallContext(ctx context.Context, b []BatchElem) error {
 		byID[string(msg.ID)] = i
 	}
 
+	resultHooks := make([]ResultHook, len(msgs))
+	responsesForHooks := make([]interface{}, len(msgs))
+	for i, msg := range msgs {
+		resultHooks[i] = c.onRequest(msg)
+	}
+
 	var err error
 	if c.isHTTP {
 		err = c.sendBatchHTTP(ctx, op, msgs)
@@ -379,7 +394,9 @@ func (c *Client) BatchCallContext(ctx context.Context, b []BatchElem) error {
 		// Find the element corresponding to this response.
 		// The element is guaranteed to be present because dispatch
 		// only sends valid IDs to our channel.
-		elem := &b[byID[string(resp.ID)]]
+		idx := byID[string(resp.ID)]
+		responsesForHooks[idx] = resp
+		elem := &b[idx]
 		if resp.Error != nil {
 			elem.Error = resp.Error
 			continue
@@ -390,6 +407,9 @@ func (c *Client) BatchCallContext(ctx context.Context, b []BatchElem) error {
 		}
 		elem.Error = json.Unmarshal(resp.Result, elem.Result)
 	}
+	for i, hook := range resultHooks {
+		hook.OnResult(responsesForHooks[i], err)
+	}
 	return err
 }
 

rpc/client_arbitrum.go

@@ -1,3 +1,19 @@
+// Copyright 2022 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
 package rpc
 
 import (
@@ -34,3 +50,37 @@ func DialTransport(ctx context.Context, rawUrl string, transport *http.Transport
 	}
 	return rpcClient, nil
 }
+
+func DialContextWithRequestHook(ctx context.Context, url string, hook RequestHook) (*Client, error) {
+	client, err := DialContext(ctx, url)
+	if err != nil {
+		return nil, err
+	}
+	client.requestHook = hook
+	return client, nil
+}
+
+type RequestHook interface {
+	OnRequest(request interface{}) ResultHook
+}
+
+type ResultHook interface {
+	OnResult(response interface{}, err error)
+}
+
+type noopResultHook struct{}
+
+func (h noopResultHook) OnResult(interface{}, error) {
+}
+
+func (c *Client) onRequest(request interface{}) ResultHook {
+	hooks := c.requestHook
+	var respHooks ResultHook
+	if hooks != nil {
+		respHooks = hooks.OnRequest(request)
+	}
+	if respHooks == nil {
+		respHooks = noopResultHook{}
+	}
+	return respHooks
+}