From 281b44f637b80f39910e64d8ac49d1fb1a091a01 Mon Sep 17 00:00:00 2001
From: fancl <fancl@justcall.cn>
Date: Mon, 15 Apr 2024 11:49:42 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0before=E6=93=8D=E4=BD=9C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 entry/http/server.go | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

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
 }