TrueCloudLab/neoneo-go
7
0
Fork 2
Code Issues Pull requests Projects Releases Packages Wiki Activity

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

Browse source 
Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>
This commit is contained in:
Dmitrii Stepanov 2024-10-15 16:38:23 +03:00
parent 5481339d69
commit 8aee80dbdc
Signed by: dstepanov-yadro
GPG key ID: 237AF1A763293BC0
2 changed files with 13 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
pkg/rpcclient/client.go
View file

@ -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{
Timeout: opts.DialTimeout,
}).DialContext,
DialContext: dialContext,
MaxConnsPerHost: opts.MaxConnsPerHost,
TLSClientConfig: opts.TLSClientConfig,
}

2
pkg/rpcclient/wsclient.go
View file

@ -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.