go-omni/transport/nethttp/router.go

// Package nethttp provides GoOmni's standard-library router. Register routes
// before serving requests.
package nethttp

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

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

type Router struct {
	routes                     []route
	notFound, methodNotAllowed http.Handler
}

type route struct {
	method   string
	pattern  string
	segments []routing.Segment
	handler  http.Handler
}

func New() *Router { return &Router{} }

// SetFallbacks must be called before serving requests. Nil preserves the plain
// net/http fallback, useful when using this router without an Omni application.
func (r *Router) SetFallbacks(notFound, methodNotAllowed http.Handler) {
	r.notFound, r.methodNotAllowed = notFound, methodNotAllowed
}

func (r *Router) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	var selected *route
	var selectedParams map[string]string
	allow := make(map[string]bool)
	for index := range r.routes {
		candidate := &r.routes[index]
		params, matched := routing.Match(candidate.segments, request.URL.Path)
		if !matched {
			continue
		}
		if candidate.method != request.Method {
			allow[candidate.method] = true
			continue
		}
		if selected == nil || moreSpecific(candidate.segments, selected.segments) {
			selected, selectedParams = candidate, params
		}
	}
	if selected == nil {
		if len(allow) > 0 {
			methods := make([]string, 0, len(allow))
			for method := range allow {
				methods = append(methods, method)
			}
			sort.Strings(methods)
			writer.Header().Set("Allow", strings.Join(methods, ", "))
			if r.methodNotAllowed != nil {
				r.methodNotAllowed.ServeHTTP(writer, request)
			} else {
				http.Error(writer, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
			}
		} else {
			if r.notFound != nil {
				r.notFound.ServeHTTP(writer, request)
			} else {
				http.NotFound(writer, request)
			}
		}
		return
	}
	ctx := httpx.WithRouteParams(request.Context(), selectedParams)
	ctx = httpx.WithRoutePattern(ctx, selected.pattern)
	selected.handler.ServeHTTP(writer, request.WithContext(ctx))
}

func (r *Router) Handle(method, path string, handler http.Handler) {
	segments := routing.Parse(method, path, handler)
	for _, existing := range r.routes {
		if existing.method == method {
			routing.CheckConflict(existing.segments, segments)
		}
	}
	r.routes = append(r.routes, route{method, path, segments, handler})
}

func (r *Router) Any(path string, handler http.Handler) {
	for _, method := range routing.Methods {
		r.Handle(method, path, handler)
	}
}

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

type group struct {
	router *Router
	prefix string
}

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

func (g *group) Any(path string, handler http.Handler) {
	g.router.Any(routing.Join(g.prefix, path), handler)
}

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

// Compare from the left: an earlier static segment wins even if another
// matching route has more static segments later in its path.
func moreSpecific(a, b []routing.Segment) bool {
	for i := 0; i < len(a) && i < len(b); i++ {
		if (a[i].Name == "") != (b[i].Name == "") {
			return a[i].Name == ""
		}
	}
	return false
}