-
Notifications
You must be signed in to change notification settings - Fork 1.6k
/
redirector.go
96 lines (87 loc) · 2.25 KB
/
redirector.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
package redirector
import (
"context"
"io"
"net"
"reflect"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/log"
)
type Dial func(net.Addr) (net.Conn, error)
func defaultDial(addr net.Addr) (net.Conn, error) {
return net.Dial("tcp", addr.String())
}
type Redirection struct {
Dial
RedirectTo net.Addr
InboundConn net.Conn
}
type Redirector struct {
ctx context.Context
redirectionChan chan *Redirection
}
func (r *Redirector) Redirect(redirection *Redirection) {
select {
case r.redirectionChan <- redirection:
log.Debug("redirect request")
case <-r.ctx.Done():
log.Debug("exiting")
}
}
func (r *Redirector) worker() {
for {
select {
case redirection := <-r.redirectionChan:
handle := func(redirection *Redirection) {
if redirection.InboundConn == nil || reflect.ValueOf(redirection.InboundConn).IsNil() {
log.Error("nil inbound conn")
return
}
defer redirection.InboundConn.Close()
if redirection.RedirectTo == nil || reflect.ValueOf(redirection.RedirectTo).IsNil() {
log.Error("nil redirection addr")
return
}
if redirection.Dial == nil {
redirection.Dial = defaultDial
}
log.Warn("redirecting connection from", redirection.InboundConn.RemoteAddr(), "to", redirection.RedirectTo.String())
outboundConn, err := redirection.Dial(redirection.RedirectTo)
if err != nil {
log.Error(common.NewError("failed to redirect to target address").Base(err))
return
}
defer outboundConn.Close()
errChan := make(chan error, 2)
copyConn := func(a, b net.Conn) {
_, err := io.Copy(a, b)
errChan <- err
}
go copyConn(outboundConn, redirection.InboundConn)
go copyConn(redirection.InboundConn, outboundConn)
select {
case err := <-errChan:
if err != nil {
log.Error(common.NewError("failed to redirect").Base(err))
}
log.Info("redirection done")
case <-r.ctx.Done():
log.Debug("exiting")
return
}
}
go handle(redirection)
case <-r.ctx.Done():
log.Debug("shutting down redirector")
return
}
}
}
func NewRedirector(ctx context.Context) *Redirector {
r := &Redirector{
ctx: ctx,
redirectionChan: make(chan *Redirection, 64),
}
go r.worker()
return r
}