forked from p4gefau1t/trojan-go
-
Notifications
You must be signed in to change notification settings - Fork 1
/
server.go
147 lines (135 loc) · 3.56 KB
/
server.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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
package dokodemo
import (
"context"
"net"
"sync"
"time"
"github.com/faireal/trojan-go/common"
"github.com/faireal/trojan-go/config"
"github.com/faireal/trojan-go/log"
"github.com/faireal/trojan-go/tunnel"
)
type Server struct {
tunnel.Server
tcpListener net.Listener
udpListener net.PacketConn
packetChan chan tunnel.PacketConn
timeout time.Duration
targetAddr *tunnel.Address
mappingLock sync.Mutex
mapping map[string]*PacketConn
ctx context.Context
cancel context.CancelFunc
}
func (s *Server) dispatchLoop() {
fixedMetadata := &tunnel.Metadata{
Address: s.targetAddr,
}
for {
buf := make([]byte, MaxPacketSize)
n, addr, err := s.udpListener.ReadFrom(buf)
if err != nil {
select {
case <-s.ctx.Done():
default:
log.Fatal(common.NewError("dokodemo failed to read from udp socket").Base(err))
}
return
}
log.Debug("udp packet from", addr)
s.mappingLock.Lock()
if conn, found := s.mapping[addr.String()]; found {
conn.input <- buf[:n]
s.mappingLock.Unlock()
continue
}
ctx, cancel := context.WithCancel(s.ctx)
conn := &PacketConn{
input: make(chan []byte, 16),
output: make(chan []byte, 16),
metadata: fixedMetadata,
src: addr,
PacketConn: s.udpListener,
ctx: ctx,
cancel: cancel,
}
s.mapping[addr.String()] = conn
s.mappingLock.Unlock()
conn.input <- buf[:n]
s.packetChan <- conn
go func(conn *PacketConn) {
for {
select {
case payload := <-conn.output:
// "Multiple goroutines may invoke methods on a Conn simultaneously."
_, err := s.udpListener.WriteTo(payload, conn.src)
if err != nil {
log.Error(common.NewError("dokodemo udp write error").Base(err))
return
}
case <-s.ctx.Done():
return
case <-time.After(s.timeout):
s.mappingLock.Lock()
delete(s.mapping, conn.src.String())
s.mappingLock.Unlock()
conn.Close()
log.Debug("closing timeout packetConn")
return
}
}
}(conn)
}
}
func (s *Server) AcceptConn(tunnel.Tunnel) (tunnel.Conn, error) {
conn, err := s.tcpListener.Accept()
if err != nil {
log.Fatal(common.NewError("dokodemo failed to accept connection").Base(err))
}
return &Conn{
Conn: conn,
targetMetadata: &tunnel.Metadata{
Address: s.targetAddr,
},
}, nil
}
func (s *Server) AcceptPacket(tunnel.Tunnel) (tunnel.PacketConn, error) {
select {
case conn := <-s.packetChan:
return conn, nil
case <-s.ctx.Done():
return nil, common.NewError("dokodemo server closed")
}
}
func (s *Server) Close() error {
s.cancel()
s.tcpListener.Close()
s.udpListener.Close()
return nil
}
func NewServer(ctx context.Context, _ tunnel.Server) (*Server, error) {
cfg := config.FromContext(ctx, Name).(*Config)
targetAddr := tunnel.NewAddressFromHostPort("tcp", cfg.TargetHost, cfg.TargetPort)
listenAddr := tunnel.NewAddressFromHostPort("tcp", cfg.LocalHost, cfg.LocalPort)
tcpListener, err := net.Listen("tcp", listenAddr.String())
if err != nil {
return nil, common.NewError("failed to listen tcp").Base(err)
}
udpListener, err := net.ListenPacket("udp", listenAddr.String())
if err != nil {
return nil, common.NewError("failed to listen udp").Base(err)
}
ctx, cancel := context.WithCancel(ctx)
server := &Server{
tcpListener: tcpListener,
udpListener: udpListener,
targetAddr: targetAddr,
mapping: make(map[string]*PacketConn),
packetChan: make(chan tunnel.PacketConn, 32),
timeout: time.Second * time.Duration(cfg.UDPTimeout),
ctx: ctx,
cancel: cancel,
}
go server.dispatchLoop()
return server, nil
}