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
	"bufio"
19
	"fmt"
20
	"io"
21
	"math"
22
	"net"
23
	"net/http"
24
	"strconv"
25
	"strings"
26
	"sync"
27

28
	"github.com/julienschmidt/httprouter"
29
)
30

31
const (
32
	abortIndex int8 = math.MaxInt8 >> 1
33
)
34

35
var jsonNull = [4]byte{'n', 'u', 'l', 'l'}
36

37
// Context handler context with http variables
38
type Context struct {
39
	opts  *serverOptions
40
	Param httprouter.Params
41

42
	Request *http.Request
43
	Writer  http.ResponseWriter
44

45
	// pass key/value in whole request
46
	mu   sync.RWMutex
47
	Meta map[string]interface{}
48

49
	wroteHeader bool
50

51
	// interceptors control
52
	index    int8
53
	handlers []HandlerFunc
54
}
55

56
// ArgsBody args in body
57
func (c *Context) ArgsBody(args interface{}) error {
58
	return c.ParseArgs(args, OptArgsBody())
59
}
60

61
// ArgsURI args in uri
62
func (c *Context) ArgsURI(args interface{}) error {
63
	return c.ParseArgs(args, OptArgsURI())
64
}
65

66
// ArgsQuery args in query
67
func (c *Context) ArgsQuery(args interface{}) error {
68
	return c.ParseArgs(args, OptArgsQuery())
69
}
70

71
// ArgsForm args in form
72
func (c *Context) ArgsForm(args interface{}) error {
73
	return c.ParseArgs(args, OptArgsForm())
74
}
75

76
// ArgsPostForm args in post form
77
func (c *Context) ArgsPostForm(args interface{}) error {
78
	return c.ParseArgs(args, OptArgsPostForm())
79
}
80

81
// ParseArgs reflect param to args
82
func (c *Context) ParseArgs(args interface{}, opts ...ServerOption) error {
83
	if err := parseArgs(c, args, opts...); err != nil {
84
		return NewError(http.StatusBadRequest, "Argument", err)
85
	}
86
	return nil
87
}
88

89
// RequestLength read request body length
90
func (c *Context) RequestLength() (int, error) {
91
	cl := c.Request.ContentLength
92
	if cl < 0 {
93
		return 0, fmt.Errorf("Unknown content length in request")
94
	}
95
	return int(cl), nil
96
}
97

98
// Next should be used only inside interceptor.
99
// It executes the pending handlers inside the calling handler.
100
func (c *Context) Next() {
101
	c.index++
102
	for c.index < int8(len(c.handlers)) {
103
		c.handlers[c.index](c)
104
		c.index++
105
	}
106
}
107

108
// IsAborted return aborted or not
109
func (c *Context) IsAborted() bool {
110
	return c.index >= abortIndex
111
}
112

113
// Abort the next handlers
114
func (c *Context) Abort() {
115
	c.index = abortIndex
116
}
117

118
// AbortWithStatus abort with status
119
func (c *Context) AbortWithStatus(statusCode int) {
120
	c.RespondStatus(statusCode)
121
	c.Abort()
122
}
123

124
// AbortWithStatusJSON abort with status and response data
125
func (c *Context) AbortWithStatusJSON(statusCode int, obj interface{}) {
126
	c.RespondStatusData(statusCode, obj)
127
	c.Abort()
128
}
129

130
// AbortWithError abort with error
131
func (c *Context) AbortWithError(err error) {
132
	c.RespondError(err)
133
	c.Abort()
134
}
135

136
// Respond response 200, and Content-Length: 0
137
func (c *Context) Respond() {
138
	c.Writer.Header().Set(HeaderContentLength, "0")
139
	c.RespondStatus(http.StatusOK)
140
}
141

142
// RespondStatus response status code
143
func (c *Context) RespondStatus(statusCode int) {
144
	c.Writer.WriteHeader(statusCode)
145
	c.wroteHeader = true
146
}
147

148
// RespondError response error
149
func (c *Context) RespondError(err error) {
150
	httpErr := Error2HTTPError(err)
151
	if httpErr == nil {
152
		c.Respond()
153
		return
154
	}
155
	c.RespondStatusData(httpErr.StatusCode(), errorResponse{
156
		Error: httpErr.Error(),
157
		Code:  httpErr.ErrorCode(),
158
	})
159
}
160

161
// RespondJSON response json
162
func (c *Context) RespondJSON(obj interface{}) {
163
	c.RespondStatusData(http.StatusOK, obj)
164
}
165

166
// RespondStatusData response data with code
167
func (c *Context) RespondStatusData(statusCode int, obj interface{}) {
168
	body, err := marshalObj(obj)
169
	if err != nil {
170
		c.RespondError(err)
171
		return
172
	}
173
	c.RespondWithReader(statusCode, body.ContentLength, body.ContentType, body.Body, nil)
174
}
175

176
// RespondWith response with code, content-type, bytes
177
func (c *Context) RespondWith(statusCode int, contentType string, body []byte) {
178
	c.Writer.Header().Set(HeaderContentType, contentType)
179
	c.Writer.Header().Set(HeaderContentLength, strconv.Itoa(len(body)))
180

181
	c.Writer.WriteHeader(statusCode)
182
	c.wroteHeader = true
183
	c.Writer.Write(body)
184
}
185

186
// RespondWithReader response with code, content-length, content-type, an io.Reader and extra headers
187
func (c *Context) RespondWithReader(statusCode int, contentLength int, contentType string,
188
	body io.Reader, extraHeaders map[string]string) {
189
	c.Writer.Header().Set(HeaderContentType, contentType)
190
	c.Writer.Header().Set(HeaderContentLength, strconv.Itoa(contentLength))
191
	for key, val := range extraHeaders {
192
		c.Writer.Header().Set(key, val)
193
	}
194

195
	c.Writer.WriteHeader(statusCode)
196
	c.wroteHeader = true
197
	io.CopyN(c.Writer, body, int64(contentLength))
198
}
199

200
// Stream sends a streaming response and returns a boolean
201
// indicates "Is client disconnected in middle of stream"
202
func (c *Context) Stream(step func(w io.Writer) bool) bool {
203
	w := c.Writer
204
	clientGone := c.Request.Context().Done()
205
	for {
206
		select {
207
		case <-clientGone:
208
			return true
209
		default:
210
			keepOpen := step(w)
211
			c.Flush()
212
			if !keepOpen {
213
				return false
214
			}
215
		}
216
	}
217
}
218

219
// Set is used to store a new key/value pair exclusively for this context.
220
func (c *Context) Set(key string, val interface{}) {
221
	c.mu.Lock()
222
	if c.Meta == nil {
223
		c.Meta = make(map[string]interface{})
224
	}
225
	c.Meta[key] = val
226
	c.mu.Unlock()
227
}
228

229
// Get returns the value for the given key,
230
// If the value does not exists it returns (nil, false).
231
func (c *Context) Get(key string) (val interface{}, exists bool) {
232
	c.mu.RLock()
233
	val, exists = c.Meta[key]
234
	c.mu.RUnlock()
235
	return
236
}
237

238
// RemoteIP parses the IP from Request.RemoteAddr, returns the net.IP (without the port).
239
func (c *Context) RemoteIP() (net.IP, bool) {
240
	ip, _, err := net.SplitHostPort(strings.TrimSpace(c.Request.RemoteAddr))
241
	if err != nil {
242
		return nil, false
243
	}
244
	remoteIP := net.ParseIP(ip)
245
	if remoteIP == nil {
246
		return nil, false
247
	}
248
	return remoteIP, true
249
}
250

251
// Hijack implements the http.Hijacker interface.
252
func (c *Context) Hijack() (net.Conn, *bufio.ReadWriter, error) {
253
	c.wroteHeader = true
254
	return c.Writer.(http.Hijacker).Hijack()
255
}
256

257
// Flush implements the http.Flush interface.
258
func (c *Context) Flush() {
259
	c.Writer.(http.Flusher).Flush()
260
}
261

262
// Pusher implements the http.Pusher interface.
263
func (c *Context) Pusher() http.Pusher {
264
	if pusher, ok := c.Writer.(http.Pusher); ok {
265
		return pusher
266
	}
267
	return nil
268
}
269