124 lines
2.4 KiB
Go
124 lines
2.4 KiB
Go
package router
|
|
|
|
import (
|
|
"encoding/json"
|
|
"sync"
|
|
|
|
"github.com/valyala/fasthttp"
|
|
)
|
|
|
|
// BodyCache 缓存 fasthttp 请求体
|
|
type BodyCache struct {
|
|
once sync.Once
|
|
Data []byte
|
|
Err error
|
|
}
|
|
|
|
// Load 读取并缓存请求体
|
|
func (b *BodyCache) Load(ctx *fasthttp.RequestCtx) ([]byte, error) {
|
|
b.once.Do(func() {
|
|
b.Data = ctx.PostBody()
|
|
})
|
|
return b.Data, b.Err
|
|
}
|
|
|
|
// Context fasthttp 上下文封装
|
|
type Context struct {
|
|
RequestCtx *fasthttp.RequestCtx
|
|
params map[string]string
|
|
index int
|
|
handlers []HandlerFunc
|
|
|
|
bodyCache BodyCache
|
|
keys map[string]interface{}
|
|
aborted bool
|
|
}
|
|
|
|
// Next 执行下一个中间件
|
|
func (c *Context) Next() error {
|
|
c.index++
|
|
if c.index < len(c.handlers) {
|
|
return c.handlers[c.index](c)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Abort 中止执行
|
|
func (c *Context) Abort() {
|
|
c.aborted = true
|
|
}
|
|
|
|
// Param 获取路由参数
|
|
func (c *Context) Param(key string) string {
|
|
return c.params[key]
|
|
}
|
|
|
|
// Header 获取请求头
|
|
func (c *Context) Header(key string) string {
|
|
return string(c.RequestCtx.Request.Header.Peek(key))
|
|
}
|
|
|
|
// Body 获取请求体
|
|
func (c *Context) Body() ([]byte, error) {
|
|
return c.bodyCache.Load(c.RequestCtx)
|
|
}
|
|
|
|
// BindJSON 反序列化 JSON 请求体
|
|
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 {
|
|
return string(c.RequestCtx.FormValue(key))
|
|
}
|
|
|
|
// Set 存储键值
|
|
func (c *Context) Set(key string, value interface{}) {
|
|
if c.keys == nil {
|
|
c.keys = make(map[string]interface{})
|
|
}
|
|
c.keys[key] = value
|
|
}
|
|
|
|
// Get 获取键值
|
|
func (c *Context) Get(key string) (interface{}, bool) {
|
|
if c.keys == nil {
|
|
return nil, false
|
|
}
|
|
val, ok := c.keys[key]
|
|
return val, ok
|
|
}
|
|
|
|
// JSON 返回 JSON 响应，支持返回 error
|
|
func (c *Context) JSON(statusCode int, data interface{}) error {
|
|
c.RequestCtx.SetStatusCode(statusCode)
|
|
c.RequestCtx.SetContentType("application/json; charset=utf-8")
|
|
|
|
if err := json.NewEncoder(c.RequestCtx).Encode(data); err != nil {
|
|
c.Abort()
|
|
return err
|
|
}
|
|
|
|
c.Abort()
|
|
return nil
|
|
}
|
|
|
|
// Text 返回纯文本，支持返回 error
|
|
func (c *Context) Text(statusCode int, msg string) error {
|
|
c.RequestCtx.SetStatusCode(statusCode)
|
|
c.RequestCtx.SetContentType("text/plain; charset=utf-8")
|
|
|
|
if _, err := c.RequestCtx.WriteString(msg); err != nil {
|
|
c.Abort()
|
|
return err
|
|
}
|
|
|
|
c.Abort()
|
|
return nil
|
|
}
|