From 2eeb8dd27188f6500ae28576aa258ff0b4e4d3c0 Mon Sep 17 00:00:00 2001
From: Chris Pacia <ctpacia@gmail.com>
Date: Mon, 6 Jan 2020 06:06:29 -0500
Subject: [PATCH] rpc: add DialWebsocketWithDialer (#20471)

This commit intents to replicate the DialHTTPWithClient function which allows
creating a RPC Client using a custom dialer but for websockets.

We introduce a new DialWebsocketWithDialer function which allows the caller
to instantiate a new websocket client using a custom dialer.
---
 rpc/websocket.go | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/rpc/websocket.go b/rpc/websocket.go
index e34241e47..b7ec56c6a 100644
--- a/rpc/websocket.go
+++ b/rpc/websocket.go
@@ -124,21 +124,13 @@ func (e wsHandshakeError) Error() string {
 	return s
 }
 
-// DialWebsocket creates a new RPC client that communicates with a JSON-RPC server
-// that is listening on the given endpoint.
-//
-// The context is used for the initial connection establishment. It does not
-// affect subsequent interactions with the client.
-func DialWebsocket(ctx context.Context, endpoint, origin string) (*Client, error) {
+// DialWebsocketWithDialer creates a new RPC client that communicates with a JSON-RPC server
+// that is listening on the given endpoint using the provided dialer.
+func DialWebsocketWithDialer(ctx context.Context, endpoint, origin string, dialer websocket.Dialer) (*Client, error) {
 	endpoint, header, err := wsClientHeaders(endpoint, origin)
 	if err != nil {
 		return nil, err
 	}
-	dialer := websocket.Dialer{
-		ReadBufferSize:  wsReadBuffer,
-		WriteBufferSize: wsWriteBuffer,
-		WriteBufferPool: wsBufferPool,
-	}
 	return newClient(ctx, func(ctx context.Context) (ServerCodec, error) {
 		conn, resp, err := dialer.DialContext(ctx, endpoint, header)
 		if err != nil {
@@ -152,6 +144,20 @@ func DialWebsocket(ctx context.Context, endpoint, origin string) (*Client, error
 	})
 }
 
+// DialWebsocket creates a new RPC client that communicates with a JSON-RPC server
+// that is listening on the given endpoint.
+//
+// The context is used for the initial connection establishment. It does not
+// affect subsequent interactions with the client.
+func DialWebsocket(ctx context.Context, endpoint, origin string) (*Client, error) {
+	dialer := websocket.Dialer{
+		ReadBufferSize:  wsReadBuffer,
+		WriteBufferSize: wsWriteBuffer,
+		WriteBufferPool: wsBufferPool,
+	}
+	return DialWebsocketWithDialer(ctx, endpoint, origin, dialer)
+}
+
 func wsClientHeaders(endpoint, origin string) (string, http.Header, error) {
 	endpointURL, err := url.Parse(endpoint)
 	if err != nil {