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
	"fmt"
20
	"io"
21
	"io/ioutil"
22
	"net/http"
23
	"strconv"
24
	"testing"
25

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

29
func TestServerCrcDecode(t *testing.T) {
30
	router := New()
31
	router.Handle(http.MethodPut, "/put/:size", func(c *Context) {
32
		size, _ := strconv.Atoi(c.Param.ByName("size"))
33

34
		n, err := io.Copy(ioutil.Discard, c.Request.Body)
35
		if err != nil && err != io.EOF {
36
			c.RespondError(err)
37
			return
38
		}
39

40
		if int(n) != size {
41
			c.RespondError(NewError(500, "size", fmt.Errorf("%d != %d", n, size)))
42
			return
43
		}
44
		c.Respond()
45
	})
46

47
	{
48
		w := new(mockResponseWriter)
49
		body := bytes.NewBuffer(make([]byte, 1024))
50
		req, _ := http.NewRequest(http.MethodPut, "/put/1024", body)
51
		router.ServeHTTP(w, req)
52
		require.Equal(t, http.StatusOK, w.status)
53
	}
54
	{
55
		w := new(mockResponseWriter)
56
		body := bytes.NewBuffer(make([]byte, 1024))
57
		req, _ := http.NewRequest(http.MethodPut, "/put/1024", body)
58
		WithCrcEncode()(req)
59
		router.ServeHTTP(w, req)
60
		require.Equal(t, http.StatusOK, w.status)
61
	}
62
	{
63
		w := new(mockResponseWriter)
64
		body := bytes.NewBuffer(make([]byte, 1024))
65
		req, _ := http.NewRequest(http.MethodPut, "/put/1025", body)
66
		WithCrcEncode()(req)
67
		router.ServeHTTP(w, req)
68
		require.Equal(t, http.StatusInternalServerError, w.status)
69
	}
70
}
71

72
func TestServerCrcDecodeWithMiddleware(t *testing.T) {
73
	router := New()
74
	router.Handle(http.MethodPut, "/put/:size", func(c *Context) {
75
		size, _ := strconv.Atoi(c.Param.ByName("size"))
76

77
		n, err := io.Copy(ioutil.Discard, c.Request.Body)
78
		if err != nil && err != io.EOF {
79
			c.RespondError(err)
80
			return
81
		}
82

83
		if int(n) != size {
84
			c.RespondError(NewError(500, "size", fmt.Errorf("%d != %d", n, size)))
85
			return
86
		}
87
		c.Respond()
88
	})
89

90
	handler := MiddlewareHandlerFuncWith(router)
91
	{
92
		w := new(mockResponseWriter)
93
		body := bytes.NewBuffer(make([]byte, 1024))
94
		req, _ := http.NewRequest(http.MethodPut, "/put/1024", body)
95
		handler.ServeHTTP(w, req)
96
		require.Equal(t, http.StatusOK, w.status)
97
	}
98
	{
99
		w := new(mockResponseWriter)
100
		body := bytes.NewBuffer(make([]byte, 1024))
101
		req, _ := http.NewRequest(http.MethodPut, "/put/1024", body)
102
		WithCrcEncode()(req)
103
		handler.ServeHTTP(w, req)
104
		require.Equal(t, http.StatusOK, w.status)
105
	}
106
	{
107
		w := new(mockResponseWriter)
108
		body := bytes.NewBuffer(make([]byte, 1024))
109
		req, _ := http.NewRequest(http.MethodPut, "/put/1025", body)
110
		WithCrcEncode()(req)
111
		handler.ServeHTTP(w, req)
112
		require.Equal(t, http.StatusInternalServerError, w.status)
113
	}
114
}
115