diff --git a/entry/http/server.go b/entry/http/server.go index 656b8f8..37b596c 100644 --- a/entry/http/server.go +++ b/entry/http/server.go @@ -24,6 +24,7 @@ type Server struct { uptime time.Time enableDocumentRoot bool fileSystem http.FileSystem + beforeRequests []HandleFunc 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 { return func(writer http.ResponseWriter, request *http.Request, params router.Params) { ctx := svr.applyContext() + ps := make(map[string]string, 4) defer func() { svr.releaseContext(ctx) + ps = make(map[string]string, 0) }() - ps := make(map[string]string) for _, v := range params { ps[v.Key] = v.Value } 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-- { 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) { svr.middleware = append(svr.middleware, middleware...) } @@ -214,11 +228,12 @@ func (svr *Server) Shutdown() (err error) { func New(ctx context.Context) *Server { svr := &Server{ - ctx: ctx, - uptime: time.Now(), - router: router.New(), - anyRequests: make(map[string]http.Handler), - middleware: make([]Middleware, 0, 10), + ctx: ctx, + uptime: time.Now(), + router: router.New(), + beforeRequests: make([]HandleFunc, 0, 10), + anyRequests: make(map[string]http.Handler), + middleware: make([]Middleware, 0, 10), } return svr }