go-omni/transport/nethttp/router.go

// Package nethttp provides a small standard-library router for GoOmni.
// It is intentionally independent of Gin and serves as the reference for a
// future project-owned network layer.
package nethttp

import (
	"net/http"
	"sort"
	"strings"
	"sync"

	"git.yuntuo.hk/k2safe/go-omni/httpx"
)

// Router implements httpx.Router with net/http only. Register routes before
// serving requests.
type Router struct {
	mu     sync.RWMutex
	routes []route
}

type route struct {
	method   string
	pattern  string
	segments []segment
	handler  http.Handler
	score    int
}

type segment struct {
	literal  string
	name     string
	wildcard bool
}

// New creates an empty standard-library router.
func New() *Router { return &Router{} }

// ServeHTTP dispatches a request and injects route params and the matched
// pattern into its context for the transport-neutral handler chain.
func (r *Router) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	r.mu.RLock()
	routes := append([]route(nil), r.routes...)
	r.mu.RUnlock()

	var selected *route
	var selectedParams map[string]string
	pathMatched := false
	allow := make(map[string]struct{})
	for index := range routes {
		candidate := &routes[index]
		params, matched := match(candidate.segments, request.URL.Path)
		if !matched {
			continue
		}
		pathMatched = true
		if candidate.method != "*" && candidate.method != request.Method {
			allow[candidate.method] = struct{}{}
			continue
		}
		if selected == nil || candidate.score > selected.score {
			selected = candidate
			selectedParams = params
		}
	}

	if selected == nil {
		if pathMatched {
			methods := make([]string, 0, len(allow))
			for method := range allow {
				methods = append(methods, method)
			}
			sort.Strings(methods)
			writer.Header().Set("Allow", strings.Join(methods, ", "))
			http.Error(writer, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
			return
		}
		http.NotFound(writer, request)
		return
	}

	requestContext := httpx.WithRouteParams(request.Context(), selectedParams)
	requestContext = httpx.WithRoutePattern(requestContext, selected.pattern)
	selected.handler.ServeHTTP(writer, request.WithContext(requestContext))
}

func (r *Router) Handle(method, path string, handler http.Handler) {
	if handler == nil {
		panic("go-omni/nethttp: nil route handler")
	}
	segments := parsePattern(path)
	r.mu.Lock()
	r.routes = append(r.routes, route{
		method:   method,
		pattern:  cleanPattern(path),
		segments: segments,
		handler:  handler,
		score:    routeScore(method, segments),
	})
	r.mu.Unlock()
}

func (r *Router) Any(path string, handler http.Handler) { r.Handle("*", path, handler) }

func (r *Router) Group(path string) httpx.RouterGroup {
	return &group{router: r, prefix: cleanPattern(path)}
}

type group struct {
	router *Router
	prefix string
}

func (g *group) Handle(method, path string, handler http.Handler) {
	g.router.Handle(method, joinPattern(g.prefix, path), handler)
}

func (g *group) Any(path string, handler http.Handler) { g.Handle("*", path, handler) }

func (g *group) Group(path string) httpx.RouterGroup {
	return &group{router: g.router, prefix: joinPattern(g.prefix, path)}
}

func cleanPattern(pattern string) string {
	if pattern == "" {
		return "/"
	}
	if !strings.HasPrefix(pattern, "/") {
		pattern = "/" + pattern
	}
	pattern = strings.TrimRight(pattern, "/")
	if pattern == "" {
		return "/"
	}
	return pattern
}

func joinPattern(prefix, path string) string {
	if prefix == "" || prefix == "/" {
		return cleanPattern(path)
	}
	if path == "" || path == "/" {
		return cleanPattern(prefix)
	}
	return cleanPattern(strings.TrimRight(prefix, "/") + "/" + strings.Trim(path, "/"))
}

func parsePattern(pattern string) []segment {
	pattern = cleanPattern(pattern)
	if pattern == "/" {
		return nil
	}
	parts := strings.Split(strings.Trim(pattern, "/"), "/")
	segments := make([]segment, 0, len(parts))
	for _, part := range parts {
		switch {
		case strings.HasPrefix(part, ":") && len(part) > 1:
			segments = append(segments, segment{name: part[1:]})
		case strings.HasPrefix(part, "*") && len(part) > 1:
			segments = append(segments, segment{name: part[1:], wildcard: true})
		default:
			segments = append(segments, segment{literal: part})
		}
	}
	return segments
}

func match(segments []segment, path string) (map[string]string, bool) {
	parts := splitPath(path)
	params := make(map[string]string)
	partIndex := 0
	for segmentIndex, current := range segments {
		if current.wildcard {
			params[current.name] = strings.Join(parts[partIndex:], "/")
			return params, true
		}
		if partIndex >= len(parts) {
			return nil, false
		}
		if current.name != "" {
			params[current.name] = parts[partIndex]
		} else if current.literal != parts[partIndex] {
			return nil, false
		}
		partIndex++
		if segmentIndex == len(segments)-1 && partIndex != len(parts) {
			return nil, false
		}
	}
	if partIndex != len(parts) {
		return nil, false
	}
	return params, true
}

func splitPath(path string) []string {
	path = strings.Trim(path, "/")
	if path == "" {
		return nil
	}
	return strings.Split(path, "/")
}

func routeScore(method string, segments []segment) int {
	score := len(segments)
	if method != "*" {
		score += 1000
	}
	for _, current := range segments {
		if current.literal != "" {
			score += 10
		}
		if current.wildcard {
			score -= 5
		}
	}
	return score
}