-
Notifications
You must be signed in to change notification settings - Fork 11
/
transport.go
122 lines (102 loc) · 2.56 KB
/
transport.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
package socks5
import (
"io"
"net"
"strings"
"sync"
)
// Transporter transmit data between client and dest server.
type Transporter interface {
TransportTCP(client *net.TCPConn, remote *net.TCPConn) <-chan error
TransportUDP(server *UDPConn, request *Request) error
}
type transport struct {
}
const maxLenOfDatagram = 65507
var transportPool = &sync.Pool{
New: func() interface{} {
return make([]byte, maxLenOfDatagram)
},
}
// TransportTCP use io.CopyBuffer transmit data.
func (t *transport) TransportTCP(client *net.TCPConn, remote *net.TCPConn) <-chan error {
errCh := make(chan error)
var wg = sync.WaitGroup{}
f := func(dst *net.TCPConn, src *net.TCPConn) {
defer wg.Done()
buf := transportPool.Get().([]byte)
defer transportPool.Put(buf)
_, err := io.CopyBuffer(dst, src, buf)
errCh <- err
}
wg.Add(2)
go func() {
wg.Wait()
defer client.Close()
defer remote.Close()
close(errCh)
}()
go f(remote, client)
go f(client, remote)
return errCh
}
// TransportUDP forwarding UDP packet between client and dest.
func (t *transport) TransportUDP(server *UDPConn, request *Request) error {
// Client udp address, limit access to the association.
clientAddr, err := request.Address.UDPAddr()
if err != nil {
return err
}
// Record dest address, limit access to the association.
forwardAddr := make(map[string]struct{})
buf := transportPool.Get().([]byte)
defer transportPool.Put(buf)
defer server.Close()
for {
select {
default:
// Receive data from remote.
n, addr, err := server.ReadFromUDP(buf)
if err != nil {
return err
}
// Should unpack data when data from client.
if strings.EqualFold(clientAddr.String(), addr.String()) {
destAddr, payload, err := UnpackUDPData(buf[:n])
if err != nil {
return err
}
destUDPAddr, err := destAddr.UDPAddr()
if err != nil {
return err
}
forwardAddr[destUDPAddr.String()] = struct{}{}
// send payload to dest address
_, err = server.WriteToUDP(payload, destUDPAddr)
if err != nil {
return err
}
}
// Should pack data when data from dest host
if _, ok := forwardAddr[addr.String()]; ok {
address, err := ParseAddress(addr.String())
if err != nil {
return err
}
// packed Data
packedData, err := PackUDPData(address, buf[:n])
if err != nil {
return err
}
// send payload to client
_, err = server.WriteToUDP(packedData, clientAddr)
if err != nil {
return err
}
}
case <-server.CloseCh():
return nil
}
}
}
var DefaultTransporter Transporter = &transport{}