/
proxy_test.go
133 lines (111 loc) · 3.57 KB
/
proxy_test.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
package proxy_test
import (
"net"
"testing"
"time"
"github.com/bbokorney/proxy"
)
var proxyAddr = "127.0.0.1:54321"
var remoteAddr = "127.0.0.1:54322"
func checkUnexpected(err error, msg string, t *testing.T) {
if err != nil {
t.Fatalf("Unexpected failure: %s: %s", msg, err)
}
}
func TestProxiedTCPConn(t *testing.T) {
// start the proxy server
proxyServer, err := net.Listen("tcp", proxyAddr)
checkUnexpected(err, "Error opening proxy server port", t)
defer proxyServer.Close()
// start the remote server
remoteServer, err := net.Listen("tcp", remoteAddr)
checkUnexpected(err, "Error opening remote server port", t)
defer remoteServer.Close()
// open connection to proxy
proxyClient, err := net.Dial("tcp", proxyAddr)
checkUnexpected(err, "Error opening proxy client connection", t)
defer proxyClient.Close()
// accept connection to proxyj
clientConn, err := proxyServer.Accept()
checkUnexpected(err, "Error accepting proxy client connection", t)
defer clientConn.Close()
// proxy the connection
p := proxy.ProxiedTCPConn{
LocalConn: clientConn,
RemoteAddr: remoteAddr,
ConnectTimeout: 10 * time.Second,
IOTimeout: 10 * time.Second,
}
go func() {
err := p.Proxy()
checkUnexpected(err, "Error proxying connection", t)
}()
// accept the proxied connection
remoteConn, err := remoteServer.Accept()
checkUnexpected(err, "Error accepting proxied connection", t)
defer remoteConn.Close()
// write to the client
testData := "test data"
_, err = proxyClient.Write([]byte(testData))
checkUnexpected(err, "Error writing to client", t)
// check the data on the remote server
buf := make([]byte, 128)
n, err := remoteConn.Read(buf)
checkUnexpected(err, "Error reading from client", t)
data := string(buf[0:n])
if data != testData {
t.Fatalf("Expected %s but got %s", testData, data)
}
// now send a response from the server
testResp := "test response"
_, err = remoteConn.Write([]byte(testResp))
checkUnexpected(err, "Error responding to client", t)
// receive the response on the client
n, err = proxyClient.Read(buf)
checkUnexpected(err, "Error reading response", t)
respData := string(buf[0:n])
if respData != testResp {
t.Fatalf("Expected %s but got %s", testResp, respData)
}
}
func TestProxiedTCPConnIOTimeout(t *testing.T) {
// start the proxy server
proxyServer, err := net.Listen("tcp", proxyAddr)
checkUnexpected(err, "Error opening proxy server port", t)
defer proxyServer.Close()
// start the remote server
remoteServer, err := net.Listen("tcp", remoteAddr)
checkUnexpected(err, "Error opening remote server port", t)
defer remoteServer.Close()
// open connection to proxy
proxyClient, err := net.Dial("tcp", proxyAddr)
checkUnexpected(err, "Error opening proxy client connection", t)
defer proxyClient.Close()
// accept connection to proxy
clientConn, err := proxyServer.Accept()
checkUnexpected(err, "Error accepting proxy client connection", t)
defer clientConn.Close()
// proxy the connection
p := proxy.ProxiedTCPConn{
LocalConn: clientConn,
RemoteAddr: remoteAddr,
ConnectTimeout: 10 * time.Second,
IOTimeout: 1 * time.Second,
}
go func() {
err := p.Proxy()
if err == nil {
t.Fatalf("Timeout error not caught")
}
}()
// accept the proxied connection
remoteConn, err := remoteServer.Accept()
checkUnexpected(err, "Error accepting proxied connection", t)
defer remoteConn.Close()
// write to the client
testData := "test data"
_, err = proxyClient.Write([]byte(testData))
checkUnexpected(err, "Error writing to client", t)
// now we wait for the timeout
time.Sleep(2 * time.Second)
}