99 lines
2.1 KiB
Go
99 lines
2.1 KiB
Go
package router
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"sync"
|
|
)
|
|
|
|
// BodyCache 专门负责缓存请求体，保证只读一次
|
|
type BodyCache struct {
|
|
once sync.Once
|
|
Data []byte
|
|
Err error
|
|
}
|
|
|
|
// Load 读取请求体并缓存，只执行一次
|
|
func (b *BodyCache) Load(r *http.Request) ([]byte, error) {
|
|
b.once.Do(func() {
|
|
if r.Body == nil {
|
|
b.Err = http.ErrBodyNotAllowed
|
|
return
|
|
}
|
|
var buf bytes.Buffer
|
|
b.Err = func() error {
|
|
_, err := io.Copy(&buf, r.Body)
|
|
return err
|
|
}()
|
|
if b.Err != nil {
|
|
return
|
|
}
|
|
b.Data = buf.Bytes()
|
|
// 重新设置请求体方便后续读取
|
|
r.Body = io.NopCloser(bytes.NewReader(b.Data))
|
|
})
|
|
return b.Data, b.Err
|
|
}
|
|
|
|
// Context 自定义请求上下文，组合 BodyCache
|
|
type Context struct {
|
|
Writer http.ResponseWriter
|
|
Request *http.Request
|
|
Params map[string]string
|
|
Index int
|
|
Handlers []HandlerFunc
|
|
|
|
BodyCache BodyCache // 请求体缓存
|
|
}
|
|
|
|
// Next 执行下一个中间件或处理函数
|
|
func (c *Context) Next() {
|
|
c.Index++
|
|
if c.Index < len(c.Handlers) {
|
|
c.Handlers[c.Index](c)
|
|
}
|
|
}
|
|
|
|
// Param 获取路由参数
|
|
func (c *Context) Param(key string) string {
|
|
return c.Params[key]
|
|
}
|
|
|
|
// Header 获取请求头
|
|
func (c *Context) Header(key string) string {
|
|
return c.Request.Header.Get(key)
|
|
}
|
|
|
|
// JSON 返回 JSON 格式响应
|
|
func (c *Context) JSON(statusCode int, data interface{}) {
|
|
c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8")
|
|
c.Writer.WriteHeader(statusCode)
|
|
if err := json.NewEncoder(c.Writer).Encode(data); err != nil {
|
|
http.Error(c.Writer, err.Error(), http.StatusInternalServerError)
|
|
}
|
|
}
|
|
|
|
// Body 方便读取请求体，实际调用 BodyCache 的 Load 方法
|
|
func (c *Context) Body() ([]byte, error) {
|
|
return c.BodyCache.Load(c.Request)
|
|
}
|
|
|
|
// BindJSON 反序列化 JSON 请求体到 obj
|
|
func (c *Context) BindJSON(obj interface{}) error {
|
|
body, err := c.Body()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return json.Unmarshal(body, obj)
|
|
}
|
|
|
|
// PostForm 获取表单参数
|
|
func (c *Context) PostForm(key string) string {
|
|
if err := c.Request.ParseForm(); err != nil {
|
|
return ""
|
|
}
|
|
return c.Request.FormValue(key)
|
|
}
|