-
Notifications
You must be signed in to change notification settings - Fork 236
/
server.go
86 lines (69 loc) · 1.69 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
package proxy
import (
"context"
"io"
"net"
"os"
"sync"
"time"
"github.com/superfly/flyctl/terminal"
)
type Server struct {
LocalAddr string
Addr string
Listener net.Listener
Dial func(ctx context.Context, network, addr string) (net.Conn, error)
}
func (srv *Server) ProxyServer(ctx context.Context) error {
defer srv.Listener.Close() //skipcq: GO-S2307
for {
select {
case <-ctx.Done():
return nil
default:
if ls, ok := srv.Listener.(*net.TCPListener); ok {
if err := ls.SetDeadline(time.Now().Add(time.Second)); err != nil {
return err
}
} else if ls, ok := srv.Listener.(*net.UnixListener); ok {
if err := ls.SetDeadline(time.Now().Add(time.Second)); err != nil {
return err
}
}
source, err := srv.Listener.Accept()
if err != nil {
if os.IsTimeout(err) {
continue
}
terminal.Debug("Error accepting connection: ", err)
}
terminal.Debug("accepted new connection from: ", source.RemoteAddr())
go func() {
defer source.Close() //skipcq: GO-S2307
target, err := srv.Dial(ctx, "tcp", srv.Addr)
if err != nil {
terminal.Debug("failed to connect to target: ", err)
return
}
defer target.Close() //skipcq: GO-S2307
wg := &sync.WaitGroup{}
wg.Add(2)
copyFunc := func(dst net.Conn, src net.Conn) {
defer wg.Done()
io.Copy(dst, src)
// close the write half if it exports a CloseWrite() method
if conn, ok := dst.(ClosableWrite); ok {
conn.CloseWrite()
}
}
go copyFunc(target, source)
go copyFunc(source, target)
wg.Wait()
terminal.Debug("connection closed")
}()
}
}
}
type ClosableWrite interface {
CloseWrite() error
}