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
5package proxy
6
7import (
8	"context"
9	"net"
10)
11
12// A ContextDialer dials using a context.
13type 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.
25func 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.
35func 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}
55