[#2] rpcclient: Allow to specify custom DialContext func

Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>

pkg/rpcclient/client.go

@@ -70,6 +70,7 @@ type Options struct {
 	// Limit total number of connections per host. No limit by default.
 	MaxConnsPerHost int
 	TLSClientConfig *tls.Config
+	NetDialContext  func(ctx context.Context, network, addr string) (net.Conn, error)
 }
 
 // cache stores cache values for the RPC client methods.
@@ -105,11 +106,19 @@ func initClient(ctx context.Context, cl *Client, endpoint string, opts Options)
 	if opts.RequestTimeout <= 0 {
 		opts.RequestTimeout = defaultRequestTimeout
 	}
+	dialContext := (&net.Dialer{
+		Timeout: opts.DialTimeout,
+	}).DialContext
+	if opts.NetDialContext != nil {
+		dialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
+			ctx, cancel := context.WithTimeout(ctx, opts.DialTimeout)
+			defer cancel()
+			return opts.NetDialContext(ctx, network, addr)
+		}
+	}
 
 	tr := &http.Transport{
-		DialContext: (&net.Dialer{
+		DialContext:     dialContext,
-			Timeout: opts.DialTimeout,
-		}).DialContext,
 		MaxConnsPerHost: opts.MaxConnsPerHost,
 		TLSClientConfig: opts.TLSClientConfig,
 	}

pkg/rpcclient/wsclient.go

@@ -453,7 +453,7 @@ var errConnClosedByUser = errors.New("connection closed by user")
 // You should call Init method to initialize the network magic the client is
 // operating on.
 func NewWS(ctx context.Context, endpoint string, opts WSOptions) (*WSClient, error) {
-	dialer := websocket.Dialer{HandshakeTimeout: opts.DialTimeout, TLSClientConfig: opts.TLSClientConfig}
+	dialer := websocket.Dialer{HandshakeTimeout: opts.DialTimeout, TLSClientConfig: opts.TLSClientConfig, NetDialContext: opts.NetDialContext}
 	ws, resp, err := dialer.DialContext(ctx, endpoint, nil)
 	if resp != nil && resp.Body != nil { // Can be non-nil even with error returned.
 		defer resp.Body.Close() // Not exactly required by websocket, but let's do this for bodyclose checker.