1 | // Copyright 2019 The Go Authors. All rights reserved.
|
---|
2 | // Use of this source code is governed by a BSD-style
|
---|
3 | // license that can be found in the LICENSE file.
|
---|
4 |
|
---|
5 | package proxy
|
---|
6 |
|
---|
7 | import (
|
---|
8 | "context"
|
---|
9 | "net"
|
---|
10 | )
|
---|
11 |
|
---|
12 | // A ContextDialer dials using a context.
|
---|
13 | type ContextDialer interface {
|
---|
14 | DialContext(ctx context.Context, network, address string) (net.Conn, error)
|
---|
15 | }
|
---|
16 |
|
---|
17 | // Dial works like DialContext on net.Dialer but using a dialer returned by FromEnvironment.
|
---|
18 | //
|
---|
19 | // The passed ctx is only used for returning the Conn, not the lifetime of the Conn.
|
---|
20 | //
|
---|
21 | // Custom dialers (registered via RegisterDialerType) that do not implement ContextDialer
|
---|
22 | // can leak a goroutine for as long as it takes the underlying Dialer implementation to timeout.
|
---|
23 | //
|
---|
24 | // A Conn returned from a successful Dial after the context has been cancelled will be immediately closed.
|
---|
25 | func Dial(ctx context.Context, network, address string) (net.Conn, error) {
|
---|
26 | d := FromEnvironment()
|
---|
27 | if xd, ok := d.(ContextDialer); ok {
|
---|
28 | return xd.DialContext(ctx, network, address)
|
---|
29 | }
|
---|
30 | return dialContext(ctx, d, network, address)
|
---|
31 | }
|
---|
32 |
|
---|
33 | // WARNING: this can leak a goroutine for as long as the underlying Dialer implementation takes to timeout
|
---|
34 | // A Conn returned from a successful Dial after the context has been cancelled will be immediately closed.
|
---|
35 | func dialContext(ctx context.Context, d Dialer, network, address string) (net.Conn, error) {
|
---|
36 | var (
|
---|
37 | conn net.Conn
|
---|
38 | done = make(chan struct{}, 1)
|
---|
39 | err error
|
---|
40 | )
|
---|
41 | go func() {
|
---|
42 | conn, err = d.Dial(network, address)
|
---|
43 | close(done)
|
---|
44 | if conn != nil && ctx.Err() != nil {
|
---|
45 | conn.Close()
|
---|
46 | }
|
---|
47 | }()
|
---|
48 | select {
|
---|
49 | case <-ctx.Done():
|
---|
50 | err = ctx.Err()
|
---|
51 | case <-done:
|
---|
52 | }
|
---|
53 | return conn, err
|
---|
54 | }
|
---|