From bd20e90e6ba87dbe6f27b81b895e53d22860489e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Efe=20=C3=87etin?= Date: Thu, 10 Mar 2022 10:35:15 +0300 Subject: [PATCH] :sparkles: feature: add initial support for hooks (#1777) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add initial support for hooks. * release ctx, mutex. * Add unit tests. * add comment lines. * update * update * remove unnecessary code. * fix race condition. * fix gosec. * skip error handling for onshutdown and onresponse. * update * separate hooks from app.go * make hooks field private, hook struct public and Hooks() func. * remove onreq and onres because of they can be done by middlewares. * OnGroupName method. * Update hooks.go Co-authored-by: hi019 <65871571+hi019@users.noreply.github.com> * handle errors for name and groupname * fix tests. * Update app.go * use struct fields instead of map * add multi-handler. * add onGroup, make prefix field public on Group struct. * Update hooks.go * add newhooks method. * ✨ feature: add initial support for hooks * remove ctx from hooks. Co-authored-by: hi019 <65871571+hi019@users.noreply.github.com> Co-authored-by: wernerr --- app.go | 63 ++++++++++++------ group.go | 31 +++++---- hooks.go | 141 +++++++++++++++++++++++++++++++++++++++ hooks_test.go | 178 ++++++++++++++++++++++++++++++++++++++++++++++++++ router.go | 10 ++- 5 files changed, 389 insertions(+), 34 deletions(-) create mode 100644 hooks.go create mode 100644 hooks_test.go diff --git a/app.go b/app.go index bcad97893b..1a84773f1c 100644 --- a/app.go +++ b/app.go @@ -112,9 +112,13 @@ type App struct { getBytes func(s string) (b []byte) // Converts byte slice to a string getString func(b []byte) string - // Mounted and main apps appList map[string]*App + // Hooks + hooks *hooks + // Latest route & group + latestRoute *Route + latestGroup *Group } // Config is a struct holding the server settings. @@ -424,14 +428,6 @@ const ( DefaultCompressedFileSuffix = ".fiber.gz" ) -// Variables for Name & GetRoute -var latestRoute struct { - route *Route - mu sync.Mutex -} - -var latestGroup Group - // DefaultErrorHandler that process return errors from handlers var DefaultErrorHandler = func(c *Ctx, err error) error { code := StatusInternalServerError @@ -462,11 +458,17 @@ func New(config ...Config) *App { }, }, // Create config - config: Config{}, - getBytes: utils.UnsafeBytes, - getString: utils.UnsafeString, - appList: make(map[string]*App), + config: Config{}, + getBytes: utils.UnsafeBytes, + getString: utils.UnsafeString, + appList: make(map[string]*App), + latestRoute: &Route{}, + latestGroup: &Group{}, } + + // Define hooks + app.hooks = newHooks(app) + // Override config if provided if len(config) > 0 { app.config = config[0] @@ -570,13 +572,18 @@ func (app *App) Mount(prefix string, fiber *App) Router { // Assign name to specific route. func (app *App) Name(name string) Router { - latestRoute.mu.Lock() - if strings.HasPrefix(latestRoute.route.path, latestGroup.prefix) { - latestRoute.route.Name = latestGroup.name + name + app.mutex.Lock() + if strings.HasPrefix(app.latestRoute.path, app.latestGroup.Prefix) { + app.latestRoute.Name = app.latestGroup.name + name } else { - latestRoute.route.Name = name + app.latestRoute.Name = name } - latestRoute.mu.Unlock() + + if err := app.hooks.executeOnNameHooks(*app.latestRoute); err != nil { + panic(err) + } + app.mutex.Unlock() + return app } @@ -703,7 +710,12 @@ func (app *App) Group(prefix string, handlers ...Handler) Router { if len(handlers) > 0 { app.register(methodUse, prefix, handlers...) } - return &Group{prefix: prefix, app: app} + grp := &Group{Prefix: prefix, app: app} + if err := app.hooks.executeOnGroupHooks(*grp); err != nil { + panic(err) + } + + return grp } // Route is used to define routes with a common prefix inside the common function. @@ -919,6 +931,10 @@ func (app *App) HandlersCount() uint32 { // // Shutdown does not close keepalive connections so its recommended to set ReadTimeout to something else than 0. func (app *App) Shutdown() error { + if app.hooks != nil { + defer app.hooks.executeOnShutdownHooks() + } + app.mutex.Lock() defer app.mutex.Unlock() if app.server == nil { @@ -932,6 +948,11 @@ func (app *App) Server() *fasthttp.Server { return app.server } +// Hooks returns the hook struct to register hooks. +func (app *App) Hooks() *hooks { + return app.hooks +} + // Test is used for internal debugging by passing a *http.Request. // Timeout is optional and defaults to 1s, -1 will disable it completely. func (app *App) Test(req *http.Request, msTimeout ...int) (resp *http.Response, err error) { @@ -1098,6 +1119,10 @@ func (app *App) serverErrorHandler(fctx *fasthttp.RequestCtx, err error) { // startupProcess Is the method which executes all the necessary processes just before the start of the server. func (app *App) startupProcess() *App { + if err := app.hooks.executeOnListenHooks(); err != nil { + panic(err) + } + app.mutex.Lock() app.buildTree() app.mutex.Unlock() diff --git a/group.go b/group.go index 0e4efaf218..20caf1fa68 100644 --- a/group.go +++ b/group.go @@ -13,9 +13,10 @@ import ( // Group struct type Group struct { - app *App - prefix string - name string + app *App + name string + + Prefix string } // Mount attaches another app instance as a sub-router along a routing path. @@ -23,7 +24,7 @@ type Group struct { // compose them as a single service using Mount. func (grp *Group) Mount(prefix string, fiber *App) Router { stack := fiber.Stack() - groupPath := getGroupPath(grp.prefix, prefix) + groupPath := getGroupPath(grp.Prefix, prefix) for m := range stack { for r := range stack[m] { @@ -46,13 +47,19 @@ func (grp *Group) Mount(prefix string, fiber *App) Router { // Assign name to specific route. func (grp *Group) Name(name string) Router { - if strings.HasPrefix(grp.prefix, latestGroup.prefix) { - grp.name = latestGroup.name + name + grp.app.mutex.Lock() + if strings.HasPrefix(grp.Prefix, grp.app.latestGroup.Prefix) { + grp.name = grp.app.latestGroup.name + name } else { grp.name = name } - latestGroup = *grp + grp.app.latestGroup = grp + + if err := grp.app.hooks.executeOnGroupNameHooks(*grp.app.latestGroup); err != nil { + panic(err) + } + grp.app.mutex.Unlock() return grp } @@ -84,14 +91,14 @@ func (grp *Group) Use(args ...interface{}) Router { panic(fmt.Sprintf("use: invalid handler %v\n", reflect.TypeOf(arg))) } } - grp.app.register(methodUse, getGroupPath(grp.prefix, prefix), handlers...) + grp.app.register(methodUse, getGroupPath(grp.Prefix, prefix), handlers...) return grp } // Get registers a route for GET methods that requests a representation // of the specified resource. Requests using GET should only retrieve data. func (grp *Group) Get(path string, handlers ...Handler) Router { - path = getGroupPath(grp.prefix, path) + path = getGroupPath(grp.Prefix, path) return grp.app.Add(MethodHead, path, handlers...).Add(MethodGet, path, handlers...) } @@ -144,12 +151,12 @@ func (grp *Group) Patch(path string, handlers ...Handler) Router { // Add allows you to specify a HTTP method to register a route func (grp *Group) Add(method, path string, handlers ...Handler) Router { - return grp.app.register(method, getGroupPath(grp.prefix, path), handlers...) + return grp.app.register(method, getGroupPath(grp.Prefix, path), handlers...) } // Static will create a file server serving static files func (grp *Group) Static(prefix, root string, config ...Static) Router { - return grp.app.registerStatic(getGroupPath(grp.prefix, prefix), root, config...) + return grp.app.registerStatic(getGroupPath(grp.Prefix, prefix), root, config...) } // All will register the handler on all HTTP methods @@ -164,7 +171,7 @@ func (grp *Group) All(path string, handlers ...Handler) Router { // api := app.Group("/api") // api.Get("/users", handler) func (grp *Group) Group(prefix string, handlers ...Handler) Router { - prefix = getGroupPath(grp.prefix, prefix) + prefix = getGroupPath(grp.Prefix, prefix) if len(handlers) > 0 { _ = grp.app.register(methodUse, prefix, handlers...) } diff --git a/hooks.go b/hooks.go new file mode 100644 index 0000000000..b13d491b98 --- /dev/null +++ b/hooks.go @@ -0,0 +1,141 @@ +package fiber + +// Handlers define a function to create hooks for Fiber. +type OnRouteHandler = func(Route) error +type OnNameHandler = OnRouteHandler +type OnGroupHandler = func(Group) error +type OnGroupNameHandler = OnGroupHandler +type OnListenHandler = func() error +type OnShutdownHandler = OnListenHandler + +type hooks struct { + // Embed app + app *App + + // Hooks + onRoute []OnRouteHandler + onName []OnNameHandler + onGroup []OnGroupHandler + onGroupName []OnGroupNameHandler + onListen []OnListenHandler + onShutdown []OnShutdownHandler +} + +func newHooks(app *App) *hooks { + return &hooks{ + app: app, + onRoute: make([]OnRouteHandler, 0), + onGroup: make([]OnGroupHandler, 0), + onGroupName: make([]OnGroupNameHandler, 0), + onName: make([]OnNameHandler, 0), + onListen: make([]OnListenHandler, 0), + onShutdown: make([]OnShutdownHandler, 0), + } +} + +// OnRoute is a hook to execute user functions on each route registeration. +// Also you can get route properties by route parameter. +func (h *hooks) OnRoute(handler ...OnRouteHandler) { + h.app.mutex.Lock() + h.onRoute = append(h.onRoute, handler...) + h.app.mutex.Unlock() +} + +// OnName is a hook to execute user functions on each route naming. +// Also you can get route properties by route parameter. +// +// WARN: OnName only works with naming routes, not groups. +func (h *hooks) OnName(handler ...OnNameHandler) { + h.app.mutex.Lock() + h.onName = append(h.onName, handler...) + h.app.mutex.Unlock() +} + +// OnGroup is a hook to execute user functions on each group registeration. +// Also you can get group properties by group parameter. +func (h *hooks) OnGroup(handler ...OnGroupHandler) { + h.app.mutex.Lock() + h.onGroup = append(h.onGroup, handler...) + h.app.mutex.Unlock() +} + +// OnGroupName is a hook to execute user functions on each group naming. +// Also you can get group properties by group parameter. +// +// WARN: OnGroupName only works with naming groups, not routes. +func (h *hooks) OnGroupName(handler ...OnGroupNameHandler) { + h.app.mutex.Lock() + h.onGroupName = append(h.onGroupName, handler...) + h.app.mutex.Unlock() +} + +// OnListen is a hook to execute user functions on Listen, ListenTLS, Listener. +func (h *hooks) OnListen(handler ...OnListenHandler) { + h.app.mutex.Lock() + h.onListen = append(h.onListen, handler...) + h.app.mutex.Unlock() +} + +// OnShutdown is a hook to execute user functions after Shutdown. +func (h *hooks) OnShutdown(handler ...OnShutdownHandler) { + h.app.mutex.Lock() + h.onShutdown = append(h.onShutdown, handler...) + h.app.mutex.Unlock() +} + +func (h *hooks) executeOnRouteHooks(route Route) error { + for _, v := range h.onRoute { + if err := v(route); err != nil { + return err + } + } + + return nil +} + +func (h *hooks) executeOnNameHooks(route Route) error { + + for _, v := range h.onName { + if err := v(route); err != nil { + return err + } + } + + return nil +} + +func (h *hooks) executeOnGroupHooks(group Group) error { + for _, v := range h.onGroup { + if err := v(group); err != nil { + return err + } + } + + return nil +} + +func (h *hooks) executeOnGroupNameHooks(group Group) error { + for _, v := range h.onGroupName { + if err := v(group); err != nil { + return err + } + } + + return nil +} + +func (h *hooks) executeOnListenHooks() error { + for _, v := range h.onListen { + if err := v(); err != nil { + return err + } + } + + return nil +} + +func (h *hooks) executeOnShutdownHooks() { + for _, v := range h.onShutdown { + _ = v() + } +} diff --git a/hooks_test.go b/hooks_test.go new file mode 100644 index 0000000000..2339e5306c --- /dev/null +++ b/hooks_test.go @@ -0,0 +1,178 @@ +package fiber + +import ( + "errors" + "fmt" + "testing" + "time" + + "github.com/gofiber/fiber/v2/internal/bytebufferpool" + "github.com/gofiber/fiber/v2/utils" +) + +var testSimpleHandler = func(c *Ctx) error { + return c.SendString("simple") +} + +func Test_Hook_OnRoute(t *testing.T) { + t.Parallel() + + app := New() + + app.Hooks().OnRoute(func(r Route) error { + utils.AssertEqual(t, "", r.Name) + + return nil + }) + + app.Get("/", testSimpleHandler).Name("x") + + subApp := New() + subApp.Get("/test", testSimpleHandler) + + app.Mount("/sub", subApp) +} + +func Test_Hook_OnName(t *testing.T) { + t.Parallel() + + app := New() + + buf := bytebufferpool.Get() + defer bytebufferpool.Put(buf) + + app.Hooks().OnName(func(r Route) error { + buf.WriteString(r.Name) + + return nil + }) + + app.Get("/", testSimpleHandler).Name("index") + + subApp := New() + subApp.Get("/test", testSimpleHandler) + subApp.Get("/test2", testSimpleHandler) + + app.Mount("/sub", subApp) + + utils.AssertEqual(t, "index", buf.String()) +} + +func Test_Hook_OnName_Error(t *testing.T) { + t.Parallel() + + app := New() + defer func() { + if err := recover(); err != nil { + utils.AssertEqual(t, "unknown error", fmt.Sprintf("%v", err)) + } + }() + + app.Hooks().OnName(func(r Route) error { + return errors.New("unknown error") + }) + + app.Get("/", testSimpleHandler).Name("index") +} + +func Test_Hook_OnGroup(t *testing.T) { + t.Parallel() + + app := New() + + buf := bytebufferpool.Get() + defer bytebufferpool.Put(buf) + + app.Hooks().OnGroup(func(g Group) error { + buf.WriteString(g.Prefix) + + return nil + }) + + grp := app.Group("/x").Name("x.") + grp.Group("/a") + + utils.AssertEqual(t, "/x/x/a", buf.String()) +} + +func Test_Hook_OnGroupName(t *testing.T) { + t.Parallel() + + app := New() + + buf := bytebufferpool.Get() + defer bytebufferpool.Put(buf) + + app.Hooks().OnGroupName(func(g Group) error { + buf.WriteString(g.name) + + return nil + }) + + grp := app.Group("/x").Name("x.") + grp.Get("/test", testSimpleHandler) + grp.Get("/test2", testSimpleHandler) + + utils.AssertEqual(t, "x.", buf.String()) +} + +func Test_Hook_OnGroupName_Error(t *testing.T) { + t.Parallel() + + app := New() + defer func() { + if err := recover(); err != nil { + utils.AssertEqual(t, "unknown error", fmt.Sprintf("%v", err)) + } + }() + + app.Hooks().OnGroupName(func(g Group) error { + return errors.New("unknown error") + }) + + grp := app.Group("/x").Name("x.") + grp.Get("/test", testSimpleHandler) +} + +func Test_Hook_OnShutdown(t *testing.T) { + t.Parallel() + + app := New() + + buf := bytebufferpool.Get() + defer bytebufferpool.Put(buf) + + app.Hooks().OnShutdown(func() error { + buf.WriteString("shutdowning") + + return nil + }) + + utils.AssertEqual(t, nil, app.Shutdown()) + utils.AssertEqual(t, "shutdowning", buf.String()) +} + +func Test_Hook_OnListen(t *testing.T) { + t.Parallel() + + app := New(Config{ + DisableStartupMessage: true, + }) + + buf := bytebufferpool.Get() + defer bytebufferpool.Put(buf) + + app.Hooks().OnListen(func() error { + buf.WriteString("ready") + + return nil + }) + + go func() { + time.Sleep(1000 * time.Millisecond) + utils.AssertEqual(t, nil, app.Shutdown()) + }() + utils.AssertEqual(t, nil, app.Listen(":9000")) + + utils.AssertEqual(t, "ready", buf.String()) +} diff --git a/router.go b/router.go index e7764546f7..53150e21c7 100644 --- a/router.go +++ b/router.go @@ -167,6 +167,7 @@ func (app *App) handler(rctx *fasthttp.RequestCtx) { if match && app.config.ETag { setETag(c, false) } + // Release Ctx app.ReleaseCtx(c) } @@ -435,9 +436,12 @@ func (app *App) addRoute(method string, route *Route) { app.routesRefreshed = true } - latestRoute.mu.Lock() - latestRoute.route = route - latestRoute.mu.Unlock() + app.mutex.Lock() + app.latestRoute = route + if err := app.hooks.executeOnRouteHooks(*route); err != nil { + panic(err) + } + app.mutex.Unlock() } // buildTree build the prefix tree from the previously registered routes