添加before操作

This commit is contained in:
fancl 2024-04-15 11:49:42 +08:00
parent 0cd890fbb0
commit 281b44f637
1 changed files with 21 additions and 6 deletions

View File

@ -24,6 +24,7 @@ type Server struct {
uptime time.Time uptime time.Time
enableDocumentRoot bool enableDocumentRoot bool
fileSystem http.FileSystem fileSystem http.FileSystem
beforeRequests []HandleFunc
anyRequests map[string]http.Handler anyRequests map[string]http.Handler
} }
@ -43,14 +44,23 @@ func (svr *Server) releaseContext(ctx *Context) {
func (svr *Server) wrapHandle(cb HandleFunc, middleware ...Middleware) router.Handle { func (svr *Server) wrapHandle(cb HandleFunc, middleware ...Middleware) router.Handle {
return func(writer http.ResponseWriter, request *http.Request, params router.Params) { return func(writer http.ResponseWriter, request *http.Request, params router.Params) {
ctx := svr.applyContext() ctx := svr.applyContext()
ps := make(map[string]string, 4)
defer func() { defer func() {
svr.releaseContext(ctx) svr.releaseContext(ctx)
ps = make(map[string]string, 0)
}() }()
ps := make(map[string]string)
for _, v := range params { for _, v := range params {
ps[v.Key] = v.Value ps[v.Key] = v.Value
} }
ctx.reset(request, writer, ps) ctx.reset(request, writer, ps)
if len(svr.beforeRequests) > 0 {
for i := len(svr.beforeRequests) - 1; i >= 0; i-- {
if err := svr.beforeRequests[i](ctx); err != nil {
ctx.Status(http.StatusServiceUnavailable)
return
}
}
}
for i := len(svr.middleware) - 1; i >= 0; i-- { for i := len(svr.middleware) - 1; i >= 0; i-- {
cb = svr.middleware[i](cb) cb = svr.middleware[i](cb)
} }
@ -63,6 +73,10 @@ func (svr *Server) wrapHandle(cb HandleFunc, middleware ...Middleware) router.Ha
} }
} }
func (svr *Server) Before(cb ...HandleFunc) {
svr.beforeRequests = append(svr.beforeRequests, cb...)
}
func (svr *Server) Use(middleware ...Middleware) { func (svr *Server) Use(middleware ...Middleware) {
svr.middleware = append(svr.middleware, middleware...) svr.middleware = append(svr.middleware, middleware...)
} }
@ -217,6 +231,7 @@ func New(ctx context.Context) *Server {
ctx: ctx, ctx: ctx,
uptime: time.Now(), uptime: time.Now(),
router: router.New(), router: router.New(),
beforeRequests: make([]HandleFunc, 0, 10),
anyRequests: make(map[string]http.Handler), anyRequests: make(map[string]http.Handler),
middleware: make([]Middleware, 0, 10), middleware: make([]Middleware, 0, 10),
} }