GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

cloudflare/cloudflared

Public

mirrored from https://github.com/cloudflare/cloudflaredAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
2025.4.2

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

cloudflared/connection

connection/connection_test.go

258lines · modecode

RawDownload
Latest commit unavailable.
1package connection
2
3import (
4 "context"
5 "crypto/rand"
6 "fmt"
7 "io"
8 "math/big"
9 "net/http"
10 "testing"
11 "time"
12
13 pkgerrors "github.com/pkg/errors"
14 "github.com/rs/zerolog"
15 "github.com/stretchr/testify/require"
16
17 cfdflow "github.com/cloudflare/cloudflared/flow"
18
19 "github.com/cloudflare/cloudflared/stream"
20 "github.com/cloudflare/cloudflared/tracing"
21 tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
22 "github.com/cloudflare/cloudflared/websocket"
23)
24
25const (
26 largeFileSize = 2 * 1024 * 1024
27 testGracePeriod = time.Millisecond * 100
28)
29
30var (
31 testOrchestrator = &mockOrchestrator{
32 originProxy: &mockOriginProxy{},
33 }
34 log = zerolog.Nop()
35 testLargeResp = make([]byte, largeFileSize)
36)
37
38var _ ReadWriteAcker = (*HTTPResponseReadWriteAcker)(nil)
39
40type testRequest struct {
41 name string
42 endpoint string
43 expectedStatus int
44 expectedBody []byte
45 isProxyError bool
46}
47
48type mockOrchestrator struct {
49 originProxy OriginProxy
50}
51
52func (mcr *mockOrchestrator) GetConfigJSON() ([]byte, error) {
53 return nil, fmt.Errorf("not implemented")
54}
55
56func (*mockOrchestrator) UpdateConfig(version int32, config []byte) *tunnelpogs.UpdateConfigurationResponse {
57 return &tunnelpogs.UpdateConfigurationResponse{
58 LastAppliedVersion: version,
59 }
60}
61
62func (mcr *mockOrchestrator) GetOriginProxy() (OriginProxy, error) {
63 return mcr.originProxy, nil
64}
65
66func (mcr *mockOrchestrator) WarpRoutingEnabled() (enabled bool) {
67 return true
68}
69
70type mockOriginProxy struct{}
71
72func (moc *mockOriginProxy) ProxyHTTP(
73 w ResponseWriter,
74 tr *tracing.TracedHTTPRequest,
75 isWebsocket bool,
76) error {
77 req := tr.Request
78 if isWebsocket {
79 switch req.URL.Path {
80 case "/ws/echo":
81 return wsEchoEndpoint(w, req)
82 case "/ws/flaky":
83 return wsFlakyEndpoint(w, req)
84 default:
85 originRespEndpoint(w, http.StatusNotFound, []byte("ws endpoint not found"))
86 return fmt.Errorf("unknown websocket endpoint %s", req.URL.Path)
87 }
88 }
89 switch req.URL.Path {
90 case "/ok":
91 originRespEndpoint(w, http.StatusOK, []byte(http.StatusText(http.StatusOK)))
92 case "/large_file":
93 originRespEndpoint(w, http.StatusOK, testLargeResp)
94 case "/400":
95 originRespEndpoint(w, http.StatusBadRequest, []byte(http.StatusText(http.StatusBadRequest)))
96 case "/500":
97 originRespEndpoint(w, http.StatusInternalServerError, []byte(http.StatusText(http.StatusInternalServerError)))
98 case "/error":
99 return fmt.Errorf("Failed to proxy to origin")
100 default:
101 originRespEndpoint(w, http.StatusNotFound, []byte("page not found"))
102 }
103 return nil
104}
105
106func (moc *mockOriginProxy) ProxyTCP(
107 ctx context.Context,
108 rwa ReadWriteAcker,
109 r *TCPRequest,
110) error {
111 if r.CfTraceID == "flow-rate-limited" {
112 return pkgerrors.Wrap(cfdflow.ErrTooManyActiveFlows, "tcp flow rate limited")
113 }
114
115 return nil
116}
117
118type echoPipe struct {
119 reader *io.PipeReader
120 writer *io.PipeWriter
121}
122
123func (ep *echoPipe) Read(p []byte) (int, error) {
124 return ep.reader.Read(p)
125}
126
127func (ep *echoPipe) Write(p []byte) (int, error) {
128 return ep.writer.Write(p)
129}
130
131// A mock origin that echos data by streaming like a tcpOverWSConnection
132// https://github.com/cloudflare/cloudflared/blob/master/ingress/origin_connection.go
133func wsEchoEndpoint(w ResponseWriter, r *http.Request) error {
134 resp := &http.Response{
135 StatusCode: http.StatusSwitchingProtocols,
136 }
137 if err := w.WriteRespHeaders(resp.StatusCode, resp.Header); err != nil {
138 return err
139 }
140 wsCtx, cancel := context.WithCancel(r.Context())
141 readPipe, writePipe := io.Pipe()
142
143 wsConn := websocket.NewConn(wsCtx, NewHTTPResponseReadWriterAcker(w, w.(http.Flusher), r), &log)
144 go func() {
145 select {
146 case <-wsCtx.Done():
147 case <-r.Context().Done():
148 }
149 readPipe.Close()
150 writePipe.Close()
151 }()
152
153 originConn := &echoPipe{reader: readPipe, writer: writePipe}
154 stream.Pipe(wsConn, originConn, &log)
155 cancel()
156 wsConn.Close()
157 return nil
158}
159
160type flakyConn struct {
161 closeAt time.Time
162}
163
164func (fc *flakyConn) Read(p []byte) (int, error) {
165 if time.Now().After(fc.closeAt) {
166 return 0, io.EOF
167 }
168 n := copy(p, "Read from flaky connection")
169 return n, nil
170}
171
172func (fc *flakyConn) Write(p []byte) (int, error) {
173 if time.Now().After(fc.closeAt) {
174 return 0, fmt.Errorf("flaky connection closed")
175 }
176 return len(p), nil
177}
178
179func wsFlakyEndpoint(w ResponseWriter, r *http.Request) error {
180 resp := &http.Response{
181 StatusCode: http.StatusSwitchingProtocols,
182 }
183 if err := w.WriteRespHeaders(resp.StatusCode, resp.Header); err != nil {
184 return err
185 }
186 wsCtx, cancel := context.WithCancel(r.Context())
187
188 wsConn := websocket.NewConn(wsCtx, NewHTTPResponseReadWriterAcker(w, w.(http.Flusher), r), &log)
189
190 rInt, _ := rand.Int(rand.Reader, big.NewInt(50))
191 closedAfter := time.Millisecond * time.Duration(rInt.Int64())
192 originConn := &flakyConn{closeAt: time.Now().Add(closedAfter)}
193 stream.Pipe(wsConn, originConn, &log)
194 cancel()
195 wsConn.Close()
196 return nil
197}
198
199func originRespEndpoint(w ResponseWriter, status int, data []byte) {
200 resp := &http.Response{
201 StatusCode: status,
202 }
203 _ = w.WriteRespHeaders(resp.StatusCode, resp.Header)
204 _, _ = w.Write(data)
205}
206
207type mockConnectedFuse struct{}
208
209func (mcf mockConnectedFuse) Connected() {}
210
211func (mcf mockConnectedFuse) IsConnected() bool {
212 return true
213}
214
215func TestShouldFlushHeaders(t *testing.T) {
216 tests := []struct {
217 headers map[string]string
218 shouldFlush bool
219 }{
220 {
221 headers: map[string]string{contentTypeHeader: "application/json", contentLengthHeader: "1"},
222 shouldFlush: false,
223 },
224 {
225 headers: map[string]string{contentTypeHeader: "text/html", contentLengthHeader: "1"},
226 shouldFlush: false,
227 },
228 {
229 headers: map[string]string{contentTypeHeader: "text/event-stream", contentLengthHeader: "1"},
230 shouldFlush: true,
231 },
232 {
233 headers: map[string]string{contentTypeHeader: "application/grpc", contentLengthHeader: "1"},
234 shouldFlush: true,
235 },
236 {
237 headers: map[string]string{contentTypeHeader: "application/x-ndjson", contentLengthHeader: "1"},
238 shouldFlush: true,
239 },
240 {
241 headers: map[string]string{contentTypeHeader: "application/json"},
242 shouldFlush: true,
243 },
244 {
245 headers: map[string]string{contentTypeHeader: "application/json", contentLengthHeader: "-1", transferEncodingHeader: "chunked"},
246 shouldFlush: true,
247 },
248 }
249
250 for _, test := range tests {
251 headers := http.Header{}
252 for k, v := range test.headers {
253 headers.Add(k, v)
254 }
255
256 require.Equal(t, test.shouldFlush, shouldFlush(headers))
257 }
258}
259