proxyclient/client.go at master · chainreactors/proxyclient · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
package proxyclient

import (
	"context"
	"errors"
	"net"
	"net/url"
	"strings"
)

type Dial func(ctx context.Context, network, address string) (net.Conn, error)

type DialFactory func(*url.URL, Dial) (Dial, error)

var DefaultDial = (&net.Dialer{}).DialContext
var schemes = map[string]DialFactory{}
var ErrNilDial = errors.New("proxyclient: dial function is nil")

func NewClient(proxy *url.URL) (Dial, error) {
	return NewClientWithDial(proxy, DefaultDial)
}

func NewClientChain(proxies []*url.URL) (Dial, error) {
	return NewClientChainWithDial(proxies, DefaultDial)
}

func NewClientWithDial(proxy *url.URL, upstreamDial Dial) (_ Dial, err error) {
	if proxy == nil {
		err = errors.New("proxy url is nil")
		return
	}
	if upstreamDial == nil {
		err = errors.New("upstream dial is nil")
		return
	}
	proxy = normalizeLink(*proxy)
	var scheme string
	ss := strings.Split(proxy.Scheme, "+")
	scheme = ss[0]
	if _, ok := schemes[scheme]; !ok {
		err = errors.New("unsupported proxy client.")
		return
	} else {
		dial, dialErr := schemes[scheme](proxy, upstreamDial)
		if dialErr != nil {
			return nil, dialErr
		}
		if dial == nil {
			return nil, ErrNilDial
		}
		return dial, nil
	}
}

func NewClientChainWithDial(proxies []*url.URL, upstreamDial Dial) (dial Dial, err error) {
	dial = upstreamDial
	for _, proxyURL := range proxies {
		dial, err = NewClientWithDial(proxyURL, dial)
		if err != nil {
			return
		}
	}
	return
}

func RegisterScheme(schemeName string, factory DialFactory) {
	schemes[strings.ToUpper(schemeName)] = factory
}

func SupportedSchemes() []string {
	schemeNames := make([]string, 0, len(schemes))
	for schemeName := range schemes {
		schemeNames = append(schemeNames, schemeName)
	}
	return schemeNames
}

func (dial Dial) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	if dial == nil {
		return nil, ErrNilDial
	}
	return dial(ctx, network, address)
}

func (dial Dial) TCPOnly(ctx context.Context, network, address string) (net.Conn, error) {
	if dial == nil {
		return nil, ErrNilDial
	}
	switch strings.ToUpper(network) {
	case "TCP", "TCP4", "TCP6":
		return dial(ctx, network, address)
	default:
		return nil, errors.New("unsupported network type.")
	}
}

func (dial Dial) Dial(network, address string) (net.Conn, error) {
	if dial == nil {
		return nil, ErrNilDial
	}
	return dial(context.Background(), network, address)
}