-
Notifications
You must be signed in to change notification settings - Fork 15
/
command.go
126 lines (105 loc) · 2.66 KB
/
command.go
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
package netutils
import (
"context"
"net"
"sync"
"sync/atomic"
"github.com/wzshiming/bridge"
"github.com/wzshiming/cmux"
)
func ConnWithCloser(conn net.Conn, closer func() error) net.Conn {
return &connCloser{Conn: conn, closer: closer}
}
type connCloser struct {
net.Conn
closer func() error
}
func (w *connCloser) Close() error {
return w.closer()
}
func ConnWithAddr(conn net.Conn, localAddr, remoteAddr net.Addr) net.Conn {
return &connAddr{Conn: conn, localAddr: localAddr, remoteAddr: remoteAddr}
}
type connAddr struct {
net.Conn
localAddr net.Addr
remoteAddr net.Addr
}
func (w *connAddr) LocalAddr() net.Addr {
if w.localAddr == nil {
return w.Conn.LocalAddr()
}
return w.localAddr
}
func (w *connAddr) RemoteAddr() net.Addr {
if w.remoteAddr == nil {
return w.Conn.RemoteAddr()
}
return w.remoteAddr
}
func NewNetAddr(network, address string) net.Addr {
return &addr{network: network, address: address}
}
type addr struct {
network string
address string
}
func (a *addr) Network() string {
return a.network
}
func (a *addr) String() string {
return a.address
}
func NewCommandDialContext(ctx context.Context, commandDialer bridge.CommandDialer, localAddr, remoteAddr net.Addr, proxy []string) (net.Conn, error) {
conn, err := commandDialer.CommandDialContext(ctx, proxy[0], proxy[1:]...)
if err != nil {
return nil, err
}
conn = ConnWithAddr(conn, localAddr, remoteAddr)
return conn, nil
}
func NewCommandListener(ctx context.Context, commandDialer bridge.CommandDialer, localAddr net.Addr, remoteAddr net.Addr, proxy []string) (net.Listener, error) {
return &listener{
ctx: ctx,
commandDialer: commandDialer,
localAddr: localAddr,
remoteAddr: remoteAddr,
proxy: proxy,
}, nil
}
type listener struct {
ctx context.Context
commandDialer bridge.CommandDialer
proxy []string
localAddr net.Addr
remoteAddr net.Addr
isClose uint32
mux sync.Mutex
}
func (l *listener) Accept() (net.Conn, error) {
l.mux.Lock()
defer l.mux.Unlock()
if atomic.LoadUint32(&l.isClose) == 1 {
return nil, ErrClosedConn
}
n, err := NewCommandDialContext(l.ctx, l.commandDialer, l.localAddr, l.remoteAddr, l.proxy)
if err != nil {
return nil, err
}
// Because there is no way to tell if there is a connection coming in from the command line,
// the next listen can only be performed if the data is read or closed
var tmp [1]byte
_, err = n.Read(tmp[:])
if err != nil {
return nil, err
}
n = cmux.UnreadConn(n, tmp[:])
return n, nil
}
func (l *listener) Close() error {
atomic.StoreUint32(&l.isClose, 1)
return nil
}
func (l *listener) Addr() net.Addr {
return l.localAddr
}