Browse Source

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.
Chris Pacia 5 years ago
parent
commit
2eeb8dd271
1 changed files with 17 additions and 11 deletions
  1. 17 11
      rpc/websocket.go

+ 17 - 11
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 {