diff options
Diffstat (limited to 'pkg/ext/router.go')
-rw-r--r-- | pkg/ext/router.go | 51 |
1 files changed, 51 insertions, 0 deletions
diff --git a/pkg/ext/router.go b/pkg/ext/router.go new file mode 100644 index 0000000..74f0a95 --- /dev/null +++ b/pkg/ext/router.go @@ -0,0 +1,51 @@ +package ext + +import ( + "github.com/fasthttp/router" + "github.com/valyala/fasthttp" +) + +type ( + Router struct { + middlewares []Middleware + fastRouter *router.Router + } + Middleware func(next fasthttp.RequestHandler) fasthttp.RequestHandler + ErrorRequestHandler func(ctx *fasthttp.RequestCtx) error +) + +func NewRouter(nestedRouter *router.Router) *Router { + return &Router{ + fastRouter: nestedRouter, + } +} + +func (self *Router) AddMiddleware(middleware Middleware) { + self.middlewares = append(self.middlewares, middleware) +} + +func wrapError(next ErrorRequestHandler) fasthttp.RequestHandler { + return func(ctx *fasthttp.RequestCtx) { + if err := next(ctx); err != nil { + ctx.Response.SetStatusCode(500) + InternalServerError(ctx, err) + } + } +} + +func (self *Router) run(next ErrorRequestHandler) fasthttp.RequestHandler { + return func(ctx *fasthttp.RequestCtx) { + req := wrapError(next) + for _, r := range self.middlewares { + req = r(req) + } + req(ctx) + } +} + +func (self *Router) GET(path string, handler ErrorRequestHandler) { + self.fastRouter.GET(path, self.run(handler)) +} +func (self *Router) POST(path string, handler ErrorRequestHandler) { + self.fastRouter.POST(path, self.run(handler)) +} |