...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15 package tcpproxy
16
17 import (
18 "fmt"
19 "io/ioutil"
20 "net"
21 "net/http"
22 "net/http/httptest"
23 "net/url"
24 "testing"
25 )
26
27 func TestUserspaceProxy(t *testing.T) {
28 l, err := net.Listen("tcp", "127.0.0.1:0")
29 if err != nil {
30 t.Fatal(err)
31 }
32 defer l.Close()
33
34 want := "hello proxy"
35 ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
36 fmt.Fprint(w, want)
37 }))
38 defer ts.Close()
39
40 u, err := url.Parse(ts.URL)
41 if err != nil {
42 t.Fatal(err)
43 }
44
45 var port uint16
46 fmt.Sscanf(u.Port(), "%d", &port)
47 p := TCPProxy{
48 Listener: l,
49 Endpoints: []*net.SRV{{Target: u.Hostname(), Port: port}},
50 }
51 go p.Run()
52 defer p.Stop()
53
54 u.Host = l.Addr().String()
55
56 res, err := http.Get(u.String())
57 if err != nil {
58 t.Fatal(err)
59 }
60 got, gerr := ioutil.ReadAll(res.Body)
61 res.Body.Close()
62 if gerr != nil {
63 t.Fatal(gerr)
64 }
65
66 if string(got) != want {
67 t.Errorf("got = %s, want %s", got, want)
68 }
69 }
70
71 func TestUserspaceProxyPriority(t *testing.T) {
72 l, err := net.Listen("tcp", "127.0.0.1:0")
73 if err != nil {
74 t.Fatal(err)
75 }
76 defer l.Close()
77
78 backends := []struct {
79 Payload string
80 Priority uint16
81 }{
82 {"hello proxy 1", 1},
83 {"hello proxy 2", 2},
84 {"hello proxy 3", 3},
85 }
86
87 var eps []*net.SRV
88 var front *url.URL
89 for _, b := range backends {
90 backend := b
91 ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
92 fmt.Fprint(w, backend.Payload)
93 }))
94 defer ts.Close()
95
96 front, err = url.Parse(ts.URL)
97 if err != nil {
98 t.Fatal(err)
99 }
100
101 var port uint16
102 fmt.Sscanf(front.Port(), "%d", &port)
103
104 ep := &net.SRV{Target: front.Hostname(), Port: port, Priority: backend.Priority}
105 eps = append(eps, ep)
106 }
107
108 p := TCPProxy{
109 Listener: l,
110 Endpoints: eps,
111 }
112 go p.Run()
113 defer p.Stop()
114
115 front.Host = l.Addr().String()
116
117 res, err := http.Get(front.String())
118 if err != nil {
119 t.Fatal(err)
120 }
121 got, gerr := ioutil.ReadAll(res.Body)
122 res.Body.Close()
123 if gerr != nil {
124 t.Fatal(gerr)
125 }
126
127 want := "hello proxy 1"
128 if string(got) != want {
129 t.Errorf("got = %s, want %s", got, want)
130 }
131 }
132
View as plain text