mux/mux.go

package mux

import (
	"fmt"
	"io"
	"net/http"
	"slices"
	"strings"
	"sync/atomic"
)

// Mux is a wrapper around the go's standard http.ServeMux.
// It's a lean wrapper with methods to make routing easier
type Mux struct {
	mux            *http.ServeMux
	routes         *RouteList
	middlewares    []func(http.Handler) http.Handler
	IsShuttingDown atomic.Bool
}

func New() *Mux {
	m := &Mux{
		mux:    http.NewServeMux(),
		routes: new(RouteList),
	}
	return m
}

// HttpServeMux DO NOT USE it for routing, exposed only for edge cases.
func (m *Mux) HttpServeMux() *http.ServeMux {
	return m.mux
}

// Use will register middleware(s) with router stack
func (m *Mux) Use(h ...func(http.Handler) http.Handler) {
	if m == nil {
		panic("mux: func Use was called on nil")
	}

	m.middlewares = append(m.middlewares, h...)
}

// GET method route
func (m *Mux) GET(pattern string, h http.HandlerFunc, mw ...func(http.Handler) http.Handler) {
	m.handle(http.MethodGet, pattern, h, mw...)
}

// HEAD method route
func (m *Mux) HEAD(pattern string, h http.HandlerFunc, mw ...func(http.Handler) http.Handler) {
	m.handle(http.MethodHead, pattern, h, mw...)
}

// POST method route
func (m *Mux) POST(pattern string, h http.HandlerFunc, mw ...func(http.Handler) http.Handler) {
	m.handle(http.MethodPost, pattern, h, mw...)
}

// PUT method route
func (m *Mux) PUT(pattern string, h http.HandlerFunc, mw ...func(http.Handler) http.Handler) {
	m.handle(http.MethodPut, pattern, h, mw...)
}

// PATCH method route
func (m *Mux) PATCH(pattern string, h http.HandlerFunc, mw ...func(http.Handler) http.Handler) {
	m.handle(http.MethodPatch, pattern, h, mw...)
}

// DELETE method route
func (m *Mux) DELETE(pattern string, h http.HandlerFunc, mw ...func(http.Handler) http.Handler) {
	m.handle(http.MethodDelete, pattern, h, mw...)
}

// CONNECT method route
func (m *Mux) CONNECT(pattern string, h http.HandlerFunc) {
	m.handle(http.MethodConnect, pattern, h)
}

// OPTIONS method route
func (m *Mux) OPTIONS(pattern string, h http.HandlerFunc) {
	m.handle(http.MethodOptions, pattern, h)
}

// TRACE method route
func (m *Mux) TRACE(pattern string, h http.HandlerFunc) {
	m.handle(http.MethodTrace, pattern, h)
}

// handle registers the handler for the given pattern.
// If the given pattern conflicts, with one that is already registered, HandleFunc
// panics.
func (m *Mux) handle(method, pattern string, h http.HandlerFunc, mw ...func(http.Handler) http.Handler) {
	if m == nil {
		panic("mux: func Handle() was called on nil")
	}

	if strings.TrimSpace(pattern) == "" {
		panic("mux: pattern cannot be empty")
	}

	if !strings.HasPrefix(pattern, "/") {
		pattern = "/" + pattern
	}

	path := fmt.Sprintf("%s %s", method, pattern)
	if len(mw) > 0 {
		m.mux.Handle(path, stack(h, copyMW(m.middlewares, mw)))
	} else {
		m.mux.Handle(path, stack(h, m.middlewares))
	}

	m.routes.Add(path)
}

// With adds inline middlewares for an endpoint handler.
func (m *Mux) With(mw ...func(http.Handler) http.Handler) *Mux {
	im := &Mux{
		mux:         m.mux,
		middlewares: copyMW(m.middlewares, mw),
		routes:      m.routes,
	}

	return im
}

// Group adds a new inline-Router along the current routing
// path, with a fresh middleware stack for the inline-Router.
func (m *Mux) Group(fn func(grp *Mux)) {
	if m == nil {
		panic("mux: Group() called on nil")
	}

	if fn == nil {
		panic("mux: Group() requires callback")
	}

	im := m.With()
	fn(im)
}

func (m *Mux) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if m == nil {
		panic("mux: method ServeHTTP called on nil")
	}

	m.mux.ServeHTTP(w, req)
}

func (m *Mux) PrintRoutes(w io.Writer) {
	for _, route := range m.routes.All() {
		w.Write([]byte(route))
		w.Write([]byte("\n"))
	}
}

func (m *Mux) RouteList() []string {
	return m.routes.All()
}

func copyMW(a []func(http.Handler) http.Handler, b []func(http.Handler) http.Handler) []func(http.Handler) http.Handler {
	if len(b) > 0 {
		return slices.Concat(a, b)
	}

	mws := make([]func(http.Handler) http.Handler, len(a))
	copy(mws, a)
	return mws
}