cubefs

1
// Copyright 2022 The CubeFS Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12
// implied. See the License for the specific language governing
13
// permissions and limitations under the License.
14

15
package rpc
16

17
import (
18
	"bytes"
19
	"encoding/json"
20
	"errors"
21
	"io"
22
	"net/http"
23
	"net/url"
24
	"strconv"
25
	"testing"
26

27
	"github.com/stretchr/testify/require"
28
)
29

30
type mockResponseWriter struct {
31
	headers http.Header
32
	status  int
33
	body    []byte
34
}
35

36
var _ http.ResponseWriter = (*mockResponseWriter)(nil)
37

38
func (m *mockResponseWriter) Header() (h http.Header) {
39
	if m.headers == nil {
40
		m.headers = http.Header{}
41
	}
42
	return m.headers
43
}
44

45
func (m *mockResponseWriter) Write(p []byte) (n int, err error) {
46
	m.body = p
47
	return len(p), nil
48
}
49

50
func (m *mockResponseWriter) WriteHeader(status int) {
51
	m.status = status
52
}
53

54
func (m *mockResponseWriter) ToResponse() *http.Response {
55
	return &http.Response{
56
		StatusCode:    m.status,
57
		ContentLength: int64(len(m.body)),
58
		Body:          io.NopCloser(bytes.NewReader(m.body)),
59
	}
60
}
61

62
func TestServerRouterBase(t *testing.T) {
63
	router := New()
64

65
	type A struct {
66
		Bar string
67
	}
68
	args := new(A)
69
	str := ""
70
	router.Use(func(c *Context) {
71
		str += "1"
72
		c.Next()
73
		str += "2"
74
	})
75
	router.Use(func(c *Context) {
76
		str += "3"
77
		c.Next()
78
		str += "4"
79
	})
80
	router.Handle(http.MethodGet, "/foo/:bar", func(c *Context) {
81
		if err := c.ParseArgs(args); err != nil {
82
			panic(err)
83
		}
84
		str += args.Bar
85
	}, OptArgsURI())
86

87
	w := new(mockResponseWriter)
88
	req, _ := http.NewRequest(http.MethodGet, "/foo/you", nil)
89
	router.Router.ServeHTTP(w, req)
90

91
	require.Equal(t, "you", args.Bar)
92
	require.Equal(t, "13you42", str)
93
}
94

95
func TestServerAbort(t *testing.T) {
96
	router := New()
97
	str := ""
98
	router.Use(func(c *Context) {
99
		str += "1"
100
		c.Next()
101

102
		require.True(t, c.IsAborted())
103
		val, ok := c.Get("Abort")
104
		require.True(t, ok)
105
		require.Equal(t, int(1), val.(int))
106

107
		str += "2"
108
	})
109
	router.Use(func(c *Context) {
110
		str += "3"
111
		c.Set("Abort", 1)
112
		c.Abort()
113
		c.Next() // aborted,
114
		str += "4"
115
	})
116
	router.Handle(http.MethodGet, "/foo/:bar", func(c *Context) {
117
		str += "5"
118
	}, OptArgsURI())
119

120
	w := new(mockResponseWriter)
121
	req, _ := http.NewRequest(http.MethodGet, "/foo/you", nil)
122
	router.Router.ServeHTTP(w, req)
123

124
	require.Equal(t, "1342", str)
125
}
126

127
func TestServerAbortStatus(t *testing.T) {
128
	resp := func(r *Router) *mockResponseWriter {
129
		w := new(mockResponseWriter)
130
		req, _ := http.NewRequest(http.MethodGet, "/", nil)
131
		r.Router.ServeHTTP(w, req)
132
		return w
133
	}
134

135
	{
136
		router := New()
137
		router.Use(func(c *Context) {
138
			// do nothing
139
		})
140
		router.Handle(http.MethodGet, "/", func(c *Context) { c.Respond() })
141
		w := resp(router)
142
		require.Equal(t, 200, w.status)
143
		require.Equal(t, "0", w.Header().Get(HeaderContentLength))
144
	}
145
	{
146
		router := New()
147
		router.Use(func(c *Context) {
148
			// do nothing
149
		})
150
		router.Handle(http.MethodGet, "/", func(c *Context) { c.RespondStatus(200) })
151
		w := resp(router)
152
		require.Equal(t, 200, w.status)
153
		require.Equal(t, "", w.Header().Get(HeaderContentLength))
154
	}
155
	{
156
		router := New()
157
		router.Use(func(c *Context) {
158
			// do nothing
159
		})
160
		router.Handle(http.MethodGet, "/", func(c *Context) {
161
			// do nothing
162
		})
163
		w := resp(router)
164
		require.Equal(t, 200, w.status)
165
		require.Equal(t, "", w.Header().Get(HeaderContentLength))
166
	}
167
	{
168
		router := New()
169
		router.Use(func(c *Context) { c.Abort() })
170
		router.Handle(http.MethodGet, "/", func(c *Context) { c.RespondStatus(204) })
171
		w := resp(router)
172
		require.Equal(t, 200, w.status)
173
	}
174
	{
175
		router := New()
176
		router.Use(func(c *Context) { c.AbortWithStatus(400) })
177
		router.Handle(http.MethodGet, "/", func(c *Context) { c.RespondStatus(204) })
178
		w := resp(router)
179
		require.Equal(t, 400, w.status)
180
	}
181
	{
182
		router := New()
183
		router.Use(func(c *Context) { c.AbortWithStatusJSON(399, struct{ I int }{10}) })
184
		router.Handle(http.MethodGet, "/", func(c *Context) { c.RespondStatus(204) })
185
		w := resp(router)
186
		require.Equal(t, 399, w.status)
187
	}
188
	{
189
		router := New()
190
		router.Use(func(c *Context) { c.AbortWithError(&Error{Status: 412}) })
191
		router.Handle(http.MethodGet, "/", func(c *Context) { c.RespondStatus(204) })
192
		w := resp(router)
193
		require.Equal(t, 412, w.status)
194
	}
195
}
196

197
func TestServerOptArgs(t *testing.T) {
198
	type A struct {
199
		Bar string
200
	}
201
	body := "{\"bar\": \"bodybar\"}"
202
	{
203
		router := New()
204
		args := new(A)
205
		router.Handle(http.MethodGet, "/:bar", func(c *Context) {
206
			if err := c.ArgsBody(args); err != nil {
207
				panic(err)
208
			}
209
		})
210

211
		w := new(mockResponseWriter)
212
		req, _ := http.NewRequest(http.MethodGet, "/uribar?bar=querybar", bytes.NewBufferString(body))
213
		req.Header.Set(HeaderContentLength, strconv.Itoa(len(body)))
214
		router.Router.ServeHTTP(w, req)
215

216
		require.Equal(t, "bodybar", args.Bar)
217
	}
218
	{
219
		router := New()
220
		args := new(A)
221
		router.Handle(http.MethodGet, "/:bar", func(c *Context) {
222
			if err := c.ArgsURI(args); err != nil {
223
				panic(err)
224
			}
225
		})
226

227
		w := new(mockResponseWriter)
228
		req, _ := http.NewRequest(http.MethodGet, "/uribar?bar=querybar", bytes.NewBufferString(body))
229
		req.Header.Set(HeaderContentLength, strconv.Itoa(len(body)))
230
		router.Router.ServeHTTP(w, req)
231

232
		require.Equal(t, "uribar", args.Bar)
233
	}
234
	{
235
		router := New()
236
		args := new(A)
237
		router.Handle(http.MethodGet, "/:bar", func(c *Context) {
238
			if err := c.ArgsQuery(args); err != nil {
239
				panic(err)
240
			}
241
		})
242

243
		w := new(mockResponseWriter)
244
		req, _ := http.NewRequest(http.MethodGet, "/uribar?bar=querybar", bytes.NewBufferString(body))
245
		req.Header.Set(HeaderContentLength, strconv.Itoa(len(body)))
246
		router.Router.ServeHTTP(w, req)
247

248
		require.Equal(t, "querybar", args.Bar)
249
	}
250
	{
251
		router := New()
252
		args := new(A)
253
		router.Handle(http.MethodPost, "/:bar", func(c *Context) {
254
			if err := c.ArgsForm(args); err != nil {
255
				panic(err)
256
			}
257
		})
258

259
		form := url.Values{}
260
		form.Set("bar", "formbar")
261
		w := new(mockResponseWriter)
262
		req, _ := http.NewRequest(http.MethodPost, "/uribar?bar=querybar", nil)
263
		req.Form = form
264
		router.Router.ServeHTTP(w, req)
265

266
		require.Equal(t, "formbar", args.Bar)
267
	}
268
	{
269
		router := New()
270
		args := new(A)
271
		router.Handle(http.MethodPost, "/:bar", func(c *Context) {
272
			if err := c.ArgsPostForm(args); err != nil {
273
				panic(err)
274
			}
275
		})
276

277
		form := url.Values{}
278
		form.Set("bar", "formbar")
279
		postForm := url.Values{}
280
		postForm.Set("bar", "postformbar")
281
		w := new(mockResponseWriter)
282
		req, _ := http.NewRequest(http.MethodPost, "/uribar?bar=querybar", nil)
283
		req.Form = form
284
		req.PostForm = postForm
285
		router.Router.ServeHTTP(w, req)
286

287
		require.Equal(t, "postformbar", args.Bar)
288
	}
289
}
290

291
func TestServerOptMetaCapacity(t *testing.T) {
292
	router := New()
293
	router.Handle(http.MethodGet, "/none", func(c *Context) {
294
		if c.opts.metaCapacity != 0 {
295
			panic(c.opts.metaCapacity)
296
		}
297
	}, OptMetaCapacity(-199))
298
	router.Handle(http.MethodGet, "/1024", func(c *Context) {
299
		if c.opts.metaCapacity != 1024 {
300
			panic(c.opts.metaCapacity)
301
		}
302
	}, OptMetaCapacity(1024))
303

304
	{
305
		w := new(mockResponseWriter)
306
		req, _ := http.NewRequest(http.MethodGet, "/none", nil)
307
		router.ServeHTTP(w, req)
308
	}
309
	{
310
		w := new(mockResponseWriter)
311
		req, _ := http.NewRequest(http.MethodGet, "/1024", nil)
312
		router.ServeHTTP(w, req)
313
	}
314
}
315

316
type marshalData struct {
317
	I int
318
	S string
319
}
320

321
func (d *marshalData) Marshal() ([]byte, string, error) {
322
	if d.I > 0 {
323
		if d.I == 100 {
324
			return nil, "", &Error{Status: 204, Err: errors.New("none")}
325
		}
326
		return []byte(d.S), MIMEStream, nil
327
	}
328
	return nil, "", &Error{Status: 400, Err: errors.New("fake error")}
329
}
330

331
func (d *marshalData) Unmarshal(buff []byte) error {
332
	for i, j := 0, len(buff)-1; i < j; i, j = i+1, j-1 {
333
		buff[i], buff[j] = buff[j], buff[i]
334
	}
335
	d.S = string(buff)
336
	return nil
337
}
338

339
type marshalToData struct {
340
	I int
341
	S string
342
}
343

344
func (d *marshalToData) MarshalTo(w io.Writer) (string, error) {
345
	if d.I > 0 {
346
		w.Write([]byte(d.S))
347
		return MIMEPlain, nil
348
	}
349
	return "", &Error{Status: 400, Err: errors.New("fake error")}
350
}
351

352
func (d *marshalToData) UnmarshalFrom(r io.Reader) error {
353
	buff, _ := io.ReadAll(r)
354
	if string(buff) == "error" {
355
		return errors.New("unmarshaler from error")
356
	}
357
	for i, j := 0, len(buff)-1; i < j; i, j = i+1, j-1 {
358
		buff[i], buff[j] = buff[j], buff[i]
359
	}
360
	d.S = string(buff)
361
	return nil
362
}
363

364
func TestServerResponseWith(t *testing.T) {
365
	resp := func(r *Router) *mockResponseWriter {
366
		w := new(mockResponseWriter)
367
		req, _ := http.NewRequest(http.MethodGet, "/", nil)
368
		r.Router.ServeHTTP(w, req)
369
		return w
370
	}
371

372
	{
373
		router := New()
374
		router.Use(func(c *Context) {
375
			// do nothing
376
		})
377
		router.Handle(http.MethodGet, "/", func(c *Context) { c.Respond() })
378
		w := resp(router)
379
		require.Equal(t, 0, len(w.body))
380
	}
381
	{
382
		router := New()
383
		router.Use(func(c *Context) { c.AbortWithStatusJSON(400, "bad") })
384
		router.Handle(http.MethodGet, "/", func(c *Context) { c.RespondStatus(204) })
385
		w := resp(router)
386
		require.Equal(t, "\"bad\"", string(w.body))
387
	}
388
	{
389
		router := New()
390
		router.Use(func(c *Context) {
391
			// do nothing
392
		})
393
		router.Handle(http.MethodGet, "/", func(c *Context) { c.RespondStatusData(255, nil) })
394
		w := resp(router)
395
		require.Equal(t, 255, w.status)
396
		require.Equal(t, "null", string(w.body))
397
	}
398
	{
399
		router := New()
400
		router.Use(func(c *Context) {
401
			// do nothing
402
		})
403
		router.Handle(http.MethodGet, "/", func(c *Context) { c.RespondStatusData(255, NoneBody) })
404
		w := resp(router)
405
		require.Equal(t, 255, w.status)
406
		require.Equal(t, "", string(w.body))
407
	}
408
	{
409
		router := New()
410
		router.Handle(http.MethodGet, "/", func(c *Context) {
411
			c.RespondJSON(&marshalData{I: 11, S: "foo bar"})
412
		})
413
		w := resp(router)
414
		require.Equal(t, 200, w.status)
415
		require.Equal(t, "foo bar", string(w.body))
416
	}
417
	{
418
		router := New()
419
		router.Handle(http.MethodGet, "/", func(c *Context) {
420
			c.RespondJSON(&marshalData{I: -11, S: "foo bar"})
421
		})
422
		w := resp(router)
423
		require.Equal(t, 400, w.status)
424

425
		ret := new(errorResponse)
426
		err := json.Unmarshal(w.body, ret)
427
		require.NoError(t, err)
428
		require.Equal(t, "", ret.Code)
429
		require.Equal(t, "fake error", ret.Error)
430
	}
431
	{
432
		router := New()
433
		router.Handle(http.MethodGet, "/", func(c *Context) {
434
			c.RespondJSON(marshalData{I: 11, S: "foo bar"})
435
		})
436
		w := resp(router)
437
		require.Equal(t, 200, w.status)
438

439
		ret := new(marshalData)
440
		err := json.Unmarshal(w.body, ret)
441
		require.NoError(t, err)
442
		require.Equal(t, 11, ret.I)
443
		require.Equal(t, "foo bar", ret.S)
444
	}
445

446
	{
447
		router := New()
448
		router.Handle(http.MethodGet, "/", func(c *Context) {
449
			c.RespondJSON(&marshalToData{I: 11, S: "bar foo"})
450
		})
451
		w := resp(router)
452
		require.Equal(t, 200, w.status)
453
		require.Equal(t, MIMEPlain, w.headers.Get(HeaderContentType))
454
		require.Equal(t, "bar foo", string(w.body))
455
	}
456
	{
457
		router := New()
458
		router.Handle(http.MethodGet, "/", func(c *Context) {
459
			c.RespondJSON(&marshalToData{I: -11, S: "bar foo"})
460
		})
461
		w := resp(router)
462
		require.Equal(t, 400, w.status)
463

464
		ret := new(errorResponse)
465
		err := json.Unmarshal(w.body, ret)
466
		require.NoError(t, err)
467
		require.Equal(t, "", ret.Code)
468
		require.Equal(t, "fake error", ret.Error)
469
	}
470

471
	{
472
		router := New()
473
		router.Handle(http.MethodGet, "/", func(c *Context) {
474
			c.RespondJSON(&marshalData{I: 100, S: "foo bar"})
475
		})
476
		resp := resp(router).ToResponse()
477
		require.Equal(t, 204, resp.StatusCode)
478

479
		err := parseData(resp, nil)
480
		require.Error(t, err)
481
	}
482
	{
483
		router := New()
484
		router.Handle(http.MethodGet, "/", func(c *Context) {
485
			c.RespondJSON(&marshalData{I: 11, S: "foo bar"})
486
		})
487
		resp := resp(router).ToResponse()
488
		require.Equal(t, 200, resp.StatusCode)
489

490
		var r marshalData
491
		err := parseData(resp, &r)
492
		require.NoError(t, err)
493
		require.Equal(t, "rab oof", r.S)
494

495
		resp.Body = io.NopCloser(bytes.NewBuffer(nil))
496
		err = parseData(resp, &r)
497
		require.Error(t, err)
498
	}
499
	{
500
		router := New()
501
		router.Handle(http.MethodGet, "/", func(c *Context) {
502
			c.RespondJSON(&marshalToData{I: 1, S: "bar foo"})
503
		})
504
		resp := resp(router).ToResponse()
505
		require.Equal(t, 200, resp.StatusCode)
506

507
		var r marshalToData
508
		err := parseData(resp, &r)
509
		require.NoError(t, err)
510
		require.Equal(t, "oof rab", r.S)
511
	}
512
	{
513
		router := New()
514
		router.Handle(http.MethodGet, "/", func(c *Context) {
515
			c.RespondJSON(&marshalToData{I: 1, S: "error"})
516
		})
517
		resp := resp(router).ToResponse()
518
		require.Equal(t, 200, resp.StatusCode)
519

520
		var r marshalToData
521
		err := parseData(resp, &r)
522
		require.Error(t, err)
523
	}
524
}
525

526
func BenchmarkServerBase(b *testing.B) {
527
	router := New()
528
	router.Use(func(c *Context) { c.Next() })
529
	router.Use(func(c *Context) { c.Next() })
530
	router.Handle(http.MethodGet, "/", func(c *Context) { c.Respond() })
531

532
	b.Helper()
533
	b.ResetTimer()
534
	for ii := 0; ii < b.N; ii++ {
535
		w := new(mockResponseWriter)
536
		req, _ := http.NewRequest(http.MethodGet, "/", nil)
537
		router.ServeHTTP(w, req)
538
	}
539
}
540