support for parameters, middleware, and route groups.

2025-10-26 12:25:18 +01:00
parent 6c318a988c
commit 6f87960f3a
4 changed files with 1117 additions and 149 deletions
--- a/sux.go
+++ b/sux.go
@@ -1,128 +1,433 @@
 package sux

 import (
+	"context"
 	"net/http"
+	"sync"
 )

 const (
-	nomatch uint8 = iota
+	noMatch nodeType = iota
 	static
+	param
+	wildcard
 )

 type (
+	nodeType uint8
+
+	// Params holds route parameters extracted from the URL
+	Params map[string]string
+
+	// MiddlewareFunc represents a middleware function
+	MiddlewareFunc func(http.Handler) http.Handler
+
 	node struct {
-		term    byte
-		ntype   uint8
-		handler http.HandlerFunc
-		child   []*node
+		path          string
+		nodeType      nodeType
+		handler       http.HandlerFunc
+		children      []*node
+		paramChild    *node  // For :param routes
+		wildcardChild *node  // For *param routes
+		paramName     string // Name of the parameter
+		middleware    []MiddlewareFunc
 	}
+
 	Router struct {
-		route map[string]*node
+		trees                   map[string]*node // One tree per HTTP method
+		middleware              []MiddlewareFunc
+		notFoundHandler         http.HandlerFunc
+		methodNotAllowedHandler http.HandlerFunc
+		paramsPool              sync.Pool
+		prefix                  string // For route groups
 	}
 )

-var root *Router
+// Context key for storing route parameters
+type contextKey string

-func makenode(term byte) *node {
+const paramsKey contextKey = "params"
+
+// ParamsFromContext extracts route parameters from the request context
+func ParamsFromContext(r *http.Request) Params {
+	if params, ok := r.Context().Value(paramsKey).(Params); ok {
+		return params
+	}
+	return nil
+}
+
+// makeNode creates a new node with the given path
+func makeNode(path string) *node {
 	return &node{
-		term:  term,
-		child: make([]*node, 0),
+		path:     path,
+		children: make([]*node, 0),
 	}
 }

-func (n *node) addchild(term byte) *node {
-	if fn := n.findchild(term); fn == nil {
-		nn := makenode(term)
-		n.child = append(n.child, nn)
-		return nn
+// addChild adds a child node for the given path segment
+func (n *node) addChild(path string) *node {
+	if child := n.findChild(path); child == nil {
+		newNode := makeNode(path)
+		n.children = append(n.children, newNode)
+		return newNode
 	} else {
-		return fn
+		return child
 	}
 }

-func (n *node) maketree(word []byte, handler http.HandlerFunc) {
-	m := n
-	for i, l := 1, len(word); i < l; i++ {
-		m = m.addchild(word[i])
+// addParamChild adds a parameter child node
+func (n *node) addParamChild(paramName string) *node {
+	if n.paramChild == nil {
+		n.paramChild = makeNode(":" + paramName)
+		n.paramChild.nodeType = param
+		n.paramChild.paramName = paramName
 	}
-	m.ntype = static
-	m.handler = handler
+	return n.paramChild
 }

-func (n *node) findchild(term byte) *node {
-	for _, v := range n.child {
-		if v.term == term {
-			return v
+// addWildcardChild adds a wildcard child node
+func (n *node) addWildcardChild(paramName string) *node {
+	if n.wildcardChild == nil {
+		n.wildcardChild = makeNode("*" + paramName)
+		n.wildcardChild.nodeType = wildcard
+		n.wildcardChild.paramName = paramName
+	}
+	return n.wildcardChild
+}
+
+// findChild finds a static child node by path
+func (n *node) findChild(path string) *node {
+	for _, child := range n.children {
+		if child.path == path {
+			return child
 		}
 	}
 	return nil
 }

-func (n *node) find(word string) *node {
-	ss := []byte(word)
-	m := n
-	for i, l := 1, len(ss); i < l; i++ {
-		m = m.findchild(ss[i])
-		if m == nil {
-			return nil
+// parsePath splits a path into segments and identifies parameters
+func parsePath(path string) []string {
+	if path == "" || path == "/" {
+		return []string{}
+	}
+
+	segments := make([]string, 0)
+	start := 1 // Skip the leading slash
+
+	for i := 1; i < len(path); i++ {
+		if path[i] == '/' {
+			if start < i {
+				segments = append(segments, path[start:i])
+			}
+			start = i + 1
 		}
 	}
-	return m
+
+	if start < len(path) {
+		segments = append(segments, path[start:])
+	}
+
+	return segments
 }

-func (n *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	if m, exists := n.route[r.Method]; exists {
-		m = m.find(r.URL.Path)
-		if m == nil {
-			http.NotFound(w, r)
-			return
-		}
-		if m.ntype == nomatch {
-			http.NotFound(w, r)
-			return
-		}
-		m.handler(w, r)
+// isParam checks if a path segment is a parameter (:param)
+func isParam(segment string) bool {
+	return len(segment) > 0 && segment[0] == ':'
+}
+
+// isWildcard checks if a path segment is a wildcard (*param)
+func isWildcard(segment string) bool {
+	return len(segment) > 0 && segment[0] == '*'
+}
+
+// getParamName extracts the parameter name from a segment
+func getParamName(segment string) string {
+	if len(segment) > 1 {
+		return segment[1:]
+	}
+	return ""
+}
+
+// addRoute adds a route to the tree
+func (n *node) addRoute(segments []string, handler http.HandlerFunc) {
+	n.addRouteWithMiddleware(segments, handler, nil)
+}
+
+// addRouteWithMiddleware adds a route to the tree with middleware
+func (n *node) addRouteWithMiddleware(segments []string, handler http.HandlerFunc, middleware []MiddlewareFunc) {
+	current := n
+
+	// Handle root path ("/") case
+	if len(segments) == 0 {
+		n.nodeType = static
+		n.handler = handler
+		n.middleware = middleware
 		return
 	}
+
+	for i, segment := range segments {
+		isLast := i == len(segments)-1
+
+		if isParam(segment) {
+			paramName := getParamName(segment)
+			current = current.addParamChild(paramName)
+		} else if isWildcard(segment) {
+			paramName := getParamName(segment)
+			current = current.addWildcardChild(paramName)
+		} else {
+			current = current.addChild(segment)
+		}
+
+		if isLast {
+			current.nodeType = static
+			current.handler = handler
+			current.middleware = middleware
+		}
+	}
 }

-func (n *Router) GET(path string, handler http.HandlerFunc) *Router {
-	n.route["GET"].maketree([]byte(path), handler)
-	return n
-}
-func (n *Router) POST(path string, handler http.HandlerFunc) *Router {
-	n.route["POST"].maketree([]byte(path), handler)
-	return n
-}
-func (n *Router) PUT(path string, handler http.HandlerFunc) *Router {
-	n.route["PUT"].maketree([]byte(path), handler)
-	return n
-}
-func (n *Router) PATCH(path string, handler http.HandlerFunc) *Router {
-	n.route["PATCH"].maketree([]byte(path), handler)
-	return n
-}
-func (n *Router) DELETE(path string, handler http.HandlerFunc) *Router {
-	n.route["DELETE"].maketree([]byte(path), handler)
-	return n
-}
-func (n *Router) OPTIONS(path string, handler http.HandlerFunc) *Router {
-	n.route["OPTIONS"].maketree([]byte(path), handler)
-	return n
-}
-func (n *Router) HEAD(path string, handler http.HandlerFunc) *Router {
-	n.route["HEAD"].maketree([]byte(path), handler)
-	return n
+// find searches for a route matching the given path
+func (n *node) find(segments []string, params Params) *node {
+	current := n
+
+	for i, segment := range segments {
+		isLast := i == len(segments)-1
+
+		// Try static match first (fastest path)
+		if child := current.findChild(segment); child != nil {
+			if isLast && child.nodeType == static {
+				return child
+			}
+			current = child
+			continue
+		}
+
+		// Try parameter match
+		if current.paramChild != nil {
+			if params != nil {
+				params[current.paramChild.paramName] = segment
+			}
+			if isLast && current.paramChild.nodeType == param {
+				return current.paramChild
+			}
+			current = current.paramChild
+			continue
+		}
+
+		// Try wildcard match (catches everything)
+		if current.wildcardChild != nil {
+			if params != nil {
+				// Wildcard captures the rest of the path
+				wildcardValue := segment
+				if i < len(segments)-1 {
+					for j := i + 1; j < len(segments); j++ {
+						wildcardValue += "/" + segments[j]
+					}
+				}
+				params[current.wildcardChild.paramName] = wildcardValue
+			}
+			return current.wildcardChild
+		}
+
+		return nil // No match found
+	}
+
+	if current.nodeType == static {
+		return current
+	}
+
+	return nil
 }

+// ServeHTTP implements the http.Handler interface
+func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	if tree, exists := r.trees[req.Method]; exists {
+		segments := parsePath(req.URL.Path)
+
+		// Get params from pool for better performance
+		params := r.paramsPool.Get().(Params)
+		for k := range params {
+			delete(params, k)
+		}
+
+		node := tree.find(segments, params)
+
+		if node == nil || node.handler == nil {
+			// Check if the path exists for other methods to return 405 instead of 404
+			hasOtherMethod := false
+			for method, otherTree := range r.trees {
+				if method != req.Method {
+					if otherNode := otherTree.find(segments, nil); otherNode != nil {
+						hasOtherMethod = true
+						break
+					}
+				}
+			}
+
+			// Return params to pool
+			r.paramsPool.Put(params)
+
+			if hasOtherMethod {
+				r.handleMethodNotAllowed(w, req)
+			} else {
+				r.handleNotFound(w, req)
+			}
+			return
+		}
+
+		// Add params to request context
+		ctx := context.WithValue(req.Context(), paramsKey, params)
+		req = req.WithContext(ctx)
+
+		// Create handler chain with middleware
+		handler := http.Handler(http.HandlerFunc(node.handler))
+
+		// Apply node-specific middleware first, then router middleware
+		for i := len(node.middleware) - 1; i >= 0; i-- {
+			handler = node.middleware[i](handler)
+		}
+		for i := len(r.middleware) - 1; i >= 0; i-- {
+			handler = r.middleware[i](handler)
+		}
+
+		// Defer returning params to pool after handler completes
+		defer func() {
+			r.paramsPool.Put(params)
+		}()
+
+		handler.ServeHTTP(w, req)
+		return
+	}
+
+	r.handleMethodNotAllowed(w, req)
+}
+
+// handleNotFound handles 404 responses
+func (r *Router) handleNotFound(w http.ResponseWriter, req *http.Request) {
+	if r.notFoundHandler != nil {
+		r.notFoundHandler(w, req)
+	} else {
+		http.NotFound(w, req)
+	}
+}
+
+// handleMethodNotAllowed handles 405 responses
+func (r *Router) handleMethodNotAllowed(w http.ResponseWriter, req *http.Request) {
+	if r.methodNotAllowedHandler != nil {
+		r.methodNotAllowedHandler(w, req)
+	} else {
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		w.Write([]byte("Method Not Allowed"))
+	}
+}
+
+// Use adds global middleware to the router
+func (r *Router) Use(middleware ...MiddlewareFunc) {
+	r.middleware = append(r.middleware, middleware...)
+}
+
+// Group creates a new route group with a prefix and optional middleware
+func (r *Router) Group(prefix string, middleware ...MiddlewareFunc) *Router {
+	// Ensure prefix starts with /
+	if len(prefix) > 0 && prefix[0] != '/' {
+		prefix = "/" + prefix
+	}
+
+	return &Router{
+		trees:                   r.trees, // Share the same trees
+		middleware:              append(r.middleware, middleware...),
+		notFoundHandler:         r.notFoundHandler,
+		methodNotAllowedHandler: r.methodNotAllowedHandler,
+		paramsPool:              r.paramsPool,
+		prefix:                  r.prefix + prefix,
+	}
+}
+
+// NotFound sets a custom 404 handler
+func (r *Router) NotFound(handler http.HandlerFunc) {
+	r.notFoundHandler = handler
+}
+
+// MethodNotAllowed sets a custom 405 handler
+func (r *Router) MethodNotAllowed(handler http.HandlerFunc) {
+	r.methodNotAllowedHandler = handler
+}
+
+// addRoute is a helper method to add routes for different HTTP methods
+func (r *Router) addRoute(method, path string, handler http.HandlerFunc) *Router {
+	if r.trees == nil {
+		r.trees = make(map[string]*node)
+	}
+
+	// Find the root tree for this method (from the main router or create new)
+	var rootTree *node
+	if tree, exists := r.trees[method]; exists {
+		rootTree = tree
+	} else {
+		rootTree = makeNode("/")
+		r.trees[method] = rootTree
+	}
+
+	// Apply group prefix if present
+	fullPath := r.prefix + path
+	segments := parsePath(fullPath)
+	rootTree.addRouteWithMiddleware(segments, handler, r.middleware)
+	return r
+}
+
+// GET adds a GET route
+func (r *Router) GET(path string, handler http.HandlerFunc) *Router {
+	return r.addRoute("GET", path, handler)
+}
+
+// POST adds a POST route
+func (r *Router) POST(path string, handler http.HandlerFunc) *Router {
+	return r.addRoute("POST", path, handler)
+}
+
+// PUT adds a PUT route
+func (r *Router) PUT(path string, handler http.HandlerFunc) *Router {
+	return r.addRoute("PUT", path, handler)
+}
+
+// PATCH adds a PATCH route
+func (r *Router) PATCH(path string, handler http.HandlerFunc) *Router {
+	return r.addRoute("PATCH", path, handler)
+}
+
+// DELETE adds a DELETE route
+func (r *Router) DELETE(path string, handler http.HandlerFunc) *Router {
+	return r.addRoute("DELETE", path, handler)
+}
+
+// OPTIONS adds an OPTIONS route
+func (r *Router) OPTIONS(path string, handler http.HandlerFunc) *Router {
+	return r.addRoute("OPTIONS", path, handler)
+}
+
+// HEAD adds a HEAD route
+func (r *Router) HEAD(path string, handler http.HandlerFunc) *Router {
+	return r.addRoute("HEAD", path, handler)
+}
+
+// New creates a new Router instance
 func New() *Router {
-	root = &Router{route: make(map[string]*node)}
-	root.route["GET"] = makenode([]byte("/")[0])
-	root.route["POST"] = makenode([]byte("/")[0])
-	root.route["PUT"] = makenode([]byte("/")[0])
-	root.route["PATCH"] = makenode([]byte("/")[0])
-	root.route["DELETE"] = makenode([]byte("/")[0])
-	root.route["OPTIONS"] = makenode([]byte("/")[0])
-	root.route["HEAD"] = makenode([]byte("/")[0])
-	return root
+	router := &Router{
+		trees:      make(map[string]*node),
+		middleware: make([]MiddlewareFunc, 0),
+		paramsPool: sync.Pool{
+			New: func() interface{} {
+				return make(Params)
+			},
+		},
+		prefix: "",
+	}
+
+	// Initialize trees for common HTTP methods
+	methods := []string{"GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"}
+	for _, method := range methods {
+		router.trees[method] = makeNode("/")
+	}
+
+	return router
 }