fancl пре 11 месеци
родитељ
комит
281b44f637
1 измењених фајлова са 21 додато и 6 уклоњено
  1. 21 6
      entry/http/server.go

+ 21 - 6
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
 }