FastRouter/router.go

package router

import (
	"fmt"

	"github.com/valyala/fasthttp"
)

type Ctx = *fasthttp.RequestCtx

// Handler is a request handler with parameters.
type Handler func(ctx Ctx, params []string)

func (h Handler) Serve(ctx Ctx, params []string) {
	h(ctx, params)
}

type Middleware func(Handler) Handler

type node struct {
	segment    string
	handler    Handler
	children   []*node
	isDynamic  bool
	isWildcard bool
	maxParams  uint8
}

type Router struct {
	get          *node
	post         *node
	put          *node
	patch        *node
	delete       *node
	middleware   []Middleware
	paramsBuffer []string // Pre-allocated buffer for parameters
}

type Group struct {
	router     *Router
	prefix     string
	middleware []Middleware
}

// New creates a new Router instance.
func New() *Router {
	return &Router{
		get:          &node{},
		post:         &node{},
		put:          &node{},
		patch:        &node{},
		delete:       &node{},
		middleware:   []Middleware{},
		paramsBuffer: make([]string, 64),
	}
}

// ServeHTTP implements the Handler interface for fasthttp
func (r *Router) ServeHTTP(ctx *fasthttp.RequestCtx) {
	path := string(ctx.Path())
	method := string(ctx.Method())

	h, params, found := r.Lookup(method, path)
	if !found {
		ctx.SetStatusCode(fasthttp.StatusNotFound)
		return
	}

	h(ctx, params)
}

// Handler returns a fasthttp request handler
func (r *Router) Handler() fasthttp.RequestHandler {
	return r.ServeHTTP
}

// Use adds middleware to the router.
func (r *Router) Use(mw ...Middleware) *Router {
	r.middleware = append(r.middleware, mw...)
	return r
}

// Group creates a new route group.
func (r *Router) Group(prefix string) *Group {
	return &Group{router: r, prefix: prefix, middleware: []Middleware{}}
}

// Use adds middleware to the group.
func (g *Group) Use(mw ...Middleware) *Group {
	g.middleware = append(g.middleware, mw...)
	return g
}

// Group creates a nested group.
func (g *Group) Group(prefix string) *Group {
	return &Group{
		router:     g.router,
		prefix:     g.prefix + prefix,
		middleware: append([]Middleware{}, g.middleware...),
	}
}

// applyMiddleware applies middleware in reverse order.
func applyMiddleware(h Handler, mw []Middleware) Handler {
	for i := len(mw) - 1; i >= 0; i-- {
		h = mw[i](h)
	}
	return h
}

// Handle registers a handler for the given method and path.
func (r *Router) Handle(method, path string, h Handler) error {
	root := r.methodNode(method)
	if root == nil {
		return fmt.Errorf("unsupported method: %s", method)
	}
	return r.addRoute(root, path, h, r.middleware)
}

func (r *Router) methodNode(method string) *node {
	switch method {
	case "GET":
		return r.get
	case "POST":
		return r.post
	case "PUT":
		return r.put
	case "PATCH":
		return r.patch
	case "DELETE":
		return r.delete
	default:
		return nil
	}
}

// Get registers a GET handler.
func (r *Router) Get(path string, h Handler) error {
	return r.Handle("GET", path, h)
}

// Post registers a POST handler.
func (r *Router) Post(path string, h Handler) error {
	return r.Handle("POST", path, h)
}

// Put registers a PUT handler.
func (r *Router) Put(path string, h Handler) error {
	return r.Handle("PUT", path, h)
}

// Patch registers a PATCH handler.
func (r *Router) Patch(path string, h Handler) error {
	return r.Handle("PATCH", path, h)
}

// Delete registers a DELETE handler.
func (r *Router) Delete(path string, h Handler) error {
	return r.Handle("DELETE", path, h)
}

func (g *Group) buildGroupMiddleware() []Middleware {
	mw := append([]Middleware{}, g.router.middleware...)
	return append(mw, g.middleware...)
}

// Handle registers a handler in the group.
func (g *Group) Handle(method, path string, h Handler) error {
	root := g.router.methodNode(method)
	if root == nil {
		return fmt.Errorf("unsupported method: %s", method)
	}
	return g.router.addRoute(root, g.prefix+path, h, g.buildGroupMiddleware())
}

// Get registers a GET handler in the group.
func (g *Group) Get(path string, h Handler) error {
	return g.Handle("GET", path, h)
}

// Post registers a POST handler in the group.
func (g *Group) Post(path string, h Handler) error {
	return g.Handle("POST", path, h)
}

// Put registers a PUT handler in the group.
func (g *Group) Put(path string, h Handler) error {
	return g.Handle("PUT", path, h)
}

// Patch registers a PATCH handler in the group.
func (g *Group) Patch(path string, h Handler) error {
	return g.Handle("PATCH", path, h)
}

// Delete registers a DELETE handler in the group.
func (g *Group) Delete(path string, h Handler) error {
	return g.Handle("DELETE", path, h)
}

// WithMiddleware applies specific middleware for next registration.
func (r *Router) WithMiddleware(mw ...Middleware) *MiddlewareRouter {
	return &MiddlewareRouter{router: r, middleware: mw}
}

// WithMiddleware applies specific middleware for next group route.
func (g *Group) WithMiddleware(mw ...Middleware) *MiddlewareGroup {
	return &MiddlewareGroup{group: g, middleware: mw}
}

type MiddlewareRouter struct {
	router     *Router
	middleware []Middleware
}

type MiddlewareGroup struct {
	group      *Group
	middleware []Middleware
}

func (mr *MiddlewareRouter) buildMiddleware() []Middleware {
	mw := append([]Middleware{}, mr.router.middleware...)
	return append(mw, mr.middleware...)
}

// Handle registers a handler with middleware router.
func (mr *MiddlewareRouter) Handle(method, path string, h Handler) error {
	root := mr.router.methodNode(method)
	if root == nil {
		return fmt.Errorf("unsupported method: %s", method)
	}
	return mr.router.addRoute(root, path, h, mr.buildMiddleware())
}

// Get registers a GET handler with middleware router.
func (mr *MiddlewareRouter) Get(path string, h Handler) error {
	return mr.Handle("GET", path, h)
}

// Post registers a POST handler with middleware router.
func (mr *MiddlewareRouter) Post(path string, h Handler) error {
	return mr.Handle("POST", path, h)
}

// Put registers a PUT handler with middleware router.
func (mr *MiddlewareRouter) Put(path string, h Handler) error {
	return mr.Handle("PUT", path, h)
}

// Patch registers a PATCH handler with middleware router.
func (mr *MiddlewareRouter) Patch(path string, h Handler) error {
	return mr.Handle("PATCH", path, h)
}

// Delete registers a DELETE handler with middleware router.
func (mr *MiddlewareRouter) Delete(path string, h Handler) error {
	return mr.Handle("DELETE", path, h)
}

func (mg *MiddlewareGroup) buildMiddleware() []Middleware {
	mw := append([]Middleware{}, mg.group.router.middleware...)
	mw = append(mw, mg.group.middleware...)
	return append(mw, mg.middleware...)
}

// Handle registers a handler with middleware group.
func (mg *MiddlewareGroup) Handle(method, path string, h Handler) error {
	root := mg.group.router.methodNode(method)
	if root == nil {
		return fmt.Errorf("unsupported method: %s", method)
	}
	return mg.group.router.addRoute(root, mg.group.prefix+path, h, mg.buildMiddleware())
}

// Get registers a GET handler with middleware group.
func (mg *MiddlewareGroup) Get(path string, h Handler) error {
	return mg.Handle("GET", path, h)
}

// Post registers a POST handler with middleware group.
func (mg *MiddlewareGroup) Post(path string, h Handler) error {
	return mg.Handle("POST", path, h)
}

// Put registers a PUT handler with middleware group.
func (mg *MiddlewareGroup) Put(path string, h Handler) error {
	return mg.Handle("PUT", path, h)
}

// Patch registers a PATCH handler with middleware group.
func (mg *MiddlewareGroup) Patch(path string, h Handler) error {
	return mg.Handle("PATCH", path, h)
}

// Delete registers a DELETE handler with middleware group.
func (mg *MiddlewareGroup) Delete(path string, h Handler) error {
	return mg.Handle("DELETE", path, h)
}

// StandardHandler adapts a standard fasthttp.RequestHandler to the router's Handler
func StandardHandler(handler fasthttp.RequestHandler) Handler {
	return func(ctx Ctx, _ []string) {
		handler(ctx)
	}
}

// readSegment extracts the next path segment.
func readSegment(path string, start int) (segment string, end int, hasMore bool) {
	if start >= len(path) {
		return "", start, false
	}
	if path[start] == '/' {
		start++
	}
	if start >= len(path) {
		return "", start, false
	}
	end = start
	for end < len(path) && path[end] != '/' {
		end++
	}
	return path[start:end], end, end < len(path)
}

// addRoute adds a new route to the trie.
func (r *Router) addRoute(root *node, path string, h Handler, mw []Middleware) error {
	h = applyMiddleware(h, mw)
	if path == "/" {
		root.handler = h
		return nil
	}
	current := root
	pos := 0
	lastWC := false
	count := uint8(0)
	for {
		seg, newPos, more := readSegment(path, pos)
		if seg == "" {
			break
		}
		isDyn := len(seg) > 2 && seg[0] == '[' && seg[len(seg)-1] == ']'
		isWC := len(seg) > 0 && seg[0] == '*'
		if isWC {
			if lastWC || more {
				return fmt.Errorf("wildcard must be the last segment in the path")
			}
			lastWC = true
		}
		if isDyn || isWC {
			count++
		}
		var child *node
		for _, c := range current.children {
			if c.segment == seg {
				child = c
				break
			}
		}
		if child == nil {
			child = &node{segment: seg, isDynamic: isDyn, isWildcard: isWC}
			current.children = append(current.children, child)
		}
		if child.maxParams < count {
			child.maxParams = count
		}
		current = child
		pos = newPos
	}
	current.handler = h
	return nil
}

// Lookup finds a handler matching method and path.
func (r *Router) Lookup(method, path string) (Handler, []string, bool) {
	root := r.methodNode(method)
	if root == nil {
		return nil, nil, false
	}
	if path == "/" {
		return root.handler, nil, root.handler != nil
	}

	buffer := r.paramsBuffer
	if cap(buffer) < int(root.maxParams) {
		buffer = make([]string, root.maxParams)
		r.paramsBuffer = buffer
	}
	buffer = buffer[:0]

	h, paramCount, found := match(root, path, 0, &buffer)
	if !found {
		return nil, nil, false
	}

	return h, buffer[:paramCount], true
}

// match traverses the trie to find a handler.
func match(current *node, path string, start int, params *[]string) (Handler, int, bool) {
	paramCount := 0

	for _, c := range current.children {
		if c.isWildcard {
			rem := path[start:]
			if len(rem) > 0 && rem[0] == '/' {
				rem = rem[1:]
			}
			*params = append(*params, rem)
			return c.handler, 1, c.handler != nil
		}
	}

	seg, pos, more := readSegment(path, start)
	if seg == "" {
		return current.handler, 0, current.handler != nil
	}

	for _, c := range current.children {
		if c.segment == seg || c.isDynamic {
			if c.isDynamic {
				*params = append(*params, seg)
				paramCount++
			}
			if !more {
				return c.handler, paramCount, c.handler != nil
			}
			h, nestedCount, ok := match(c, path, pos, params)
			if ok {
				return h, paramCount + nestedCount, true
			}
		}
	}

	return nil, 0, false
}