-
-
Notifications
You must be signed in to change notification settings - Fork 78
/
forward_connection.go
64 lines (57 loc) · 1.24 KB
/
forward_connection.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
package tlsmux
import (
"io"
"net"
"sync"
)
type tcpLike interface {
CloseRead() error
CloseWrite() error
}
func forwardConnection(conn net.Conn, destConn net.Conn) error {
if isTCPTunnel(conn, destConn) {
// each side of the connection can be closed independently so no synchronisation required
go copyAndCloseTCP(conn, destConn)
go copyAndCloseTCP(destConn, conn)
return nil
}
wg := sync.WaitGroup{}
wg.Add(2)
go func() {
io.Copy(conn, destConn)
wg.Done()
}()
go func() {
io.Copy(destConn, conn)
wg.Done()
}()
go func() {
wg.Wait()
conn.Close()
destConn.Close()
}()
return nil
}
func copyAndCloseTCP(dst, src net.Conn) {
io.Copy(dst, src)
dst.(tcpLike).CloseWrite()
src.(tcpLike).CloseRead()
}
// checks if the two connections are "TCP-like"
// i.e. the two connection halves can be close separately
func isTCPTunnel(a, b net.Conn) bool {
_, aTCP := a.(tcpLike)
_, bTCP := b.(tcpLike)
return aTCP && bTCP
}
// The following can be used to debug all traffic forwarded over a connection
//
//type debugReader struct {
// io.Reader
//}
//
//func (d *debugReader) Read(p []byte) (n int, err error) {
// n, err = d.Reader.Read(p)
// fmt.Printf("[%p] Read bytes: %v, %v\n", d, string(p[:n]), err)
// return n, err
//}