upstream-ory-kratos/oryx/httprouterx/router.go

// Copyright © 2023 Ory Corp
// SPDX-License-Identifier: Apache-2.0

package httprouterx

import (
	"context"
	"net/http"
	"path"
	"strings"
)

const (
	AdminPrefix = "/admin"

	afterMatchHooks contextKey = 1
)

type (
	contextKey         int
	AfterMatchHookFunc func(*http.Request)

	router struct {
		mux    *http.ServeMux
		prefix string
	}
	RouterAdmin  struct{ router }
	RouterPublic struct{ router }

	Router interface {
		http.Handler

		Handle(pattern string, handler http.Handler)
		HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request))

		GET(route string, handle http.HandlerFunc)
		OPTIONS(route string, handle http.HandlerFunc)
		HEAD(route string, handle http.HandlerFunc)
		POST(route string, handle http.HandlerFunc)
		PUT(route string, handle http.HandlerFunc)
		PATCH(route string, handle http.HandlerFunc)
		DELETE(route string, handle http.HandlerFunc)
	}
)

func newRouter() *router { return &router{mux: http.NewServeMux()} }

// NewRouter creates a new general purpose router. It should only be used when neither the admin nor the public router is applicable.
func NewRouter() Router { return newRouter() }

// NewRouterAdmin creates a new admin router.
func NewRouterAdmin() *RouterAdmin {
	return &RouterAdmin{router: *newRouter()}
}

func (r *RouterAdmin) ToPublic() *RouterPublic {
	return &RouterPublic{router: router{
		mux:    r.mux,
		prefix: "", // do not copy the admin prefix
	}}
}

// NewRouterPublic returns a public router.
func NewRouterPublic() *RouterPublic {
	return &RouterPublic{router: *newRouter()}
}

// NewRouterAdminWithPrefix creates a new router with the admin prefix.
func NewRouterAdminWithPrefix() *RouterAdmin {
	r := NewRouterAdmin()
	r.prefix = AdminPrefix
	return r
}

func (r *router) GET(route string, handle http.HandlerFunc) {
	r.handle(http.MethodGet, route, handle)
}

func (r *router) OPTIONS(route string, handle http.HandlerFunc) {
	r.handle(http.MethodOptions, route, handle)
}

func (r *router) HEAD(route string, handle http.HandlerFunc) {
	r.handle(http.MethodHead, route, handle)
}

func (r *router) POST(route string, handle http.HandlerFunc) {
	r.handle(http.MethodPost, route, handle)
}

func (r *router) PUT(route string, handle http.HandlerFunc) {
	r.handle(http.MethodPut, route, handle)
}

func (r *router) PATCH(route string, handle http.HandlerFunc) {
	r.handle(http.MethodPatch, route, handle)
}

func (r *router) DELETE(route string, handle http.HandlerFunc) {
	r.handle(http.MethodDelete, route, handle)
}

func (r *router) handle(method string, route string, handler http.Handler) {
	route = path.Join(r.prefix, route)
	r.mux.Handle(method+" "+route, afterMatchHookExecutor(handler))
}

func afterMatchHookExecutor(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if hooks, ok := r.Context().Value(afterMatchHooks).([]AfterMatchHookFunc); ok {
			for _, hook := range hooks {
				hook(r)
			}
		}
		h.ServeHTTP(w, r)
	})
}

// Handle registers the handler for the given pattern. It does not prepend the admin prefix, so the caller is responsible for ensuring that the pattern is correct.
func (r *router) Handle(pattern string, handler http.Handler) {
	r.mux.Handle(pattern, afterMatchHookExecutor(handler))
}

// HandleFunc registers the handler for the given pattern. It does not prepend the admin prefix, so the caller is responsible for ensuring that the pattern is correct.
func (r *router) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
	r.Handle(pattern, http.HandlerFunc(handler))
}

func (r *router) ServeHTTP(w http.ResponseWriter, req *http.Request) { r.mux.ServeHTTP(w, req) }

func TrimTrailingSlashNegroni(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	r.URL.Path = strings.TrimSuffix(r.URL.Path, "/")

	next(rw, r)
}

func NoCacheNegroni(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	if r.Method == "GET" {
		rw.Header().Set("Cache-Control", "private, no-cache, no-store, must-revalidate")
	}

	next(rw, r)
}

func AddAdminPrefixIfNotPresentNegroni(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	if !strings.HasPrefix(r.URL.Path, AdminPrefix) {
		r.URL.Path = path.Join(AdminPrefix, r.URL.Path)
		if r.URL.RawPath != "" {
			r.URL.RawPath = path.Join(AdminPrefix, r.URL.RawPath)
		}
	}

	next(rw, r)
}

// WithAfterMatchHook adds the passed AfterMatchHookFunc to the request's context. The hook is called after the
// router matched the request, but before it is passed to the handler. This allows middlewares to access the matched
// pattern and other information that is only available after the router has matched the request.
// Note that this hook is called multiple times per request if routers are nested, the hook should anticipate that.
func WithAfterMatchHook(req *http.Request, newHooks ...AfterMatchHookFunc) *http.Request {
	hooks, _ := req.Context().Value(afterMatchHooks).([]AfterMatchHookFunc)
	ctx := context.WithValue(req.Context(), afterMatchHooks, append(hooks, newHooks...))
	return req.WithContext(ctx)
}