mux/router.go

package mux

import (
	"fmt"
	"net/http"
	"strings"
)

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

func NewRouter() *Router {
	return &Router{
		mux: http.NewServeMux(),
	}
}

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

// Get method route
func (r *Router) Get(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodGet, pattern, h)
}

// Head method route
func (r *Router) Head(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodHead, pattern, h)
}

// Post method route
func (r *Router) Post(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodPost, pattern, h)
}

// Put method route
func (r *Router) Put(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodPut, pattern, h)
}

// Patch method route
func (r *Router) Patch(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodPatch, pattern, h)
}

// Delete method route
func (r *Router) Delete(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodDelete, pattern, h)
}

// Connect method route
func (r *Router) Connect(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodConnect, pattern, h)
}

// Options method route
func (r *Router) Options(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodOptions, pattern, h)
}

// Trace method route
func (r *Router) Trace(pattern string, h http.HandlerFunc) {
	r.handlerFunc(http.MethodTrace, pattern, h)
}

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

	path := fmt.Sprintf("%s %s", method, pattern)
	r.mux.Handle(path, stack(r.middlewares, h))
}

// With adds inline middlewares for an endpoint handler.
func (r *Router) With(middleware ...func(http.Handler) http.Handler) *Router {
	mws := make([]func(http.Handler) http.Handler, len(r.middlewares))
	copy(mws, r.middlewares)
	mws = append(mws, middleware...)

	im := &Router{
		mux:         r.mux,
		middlewares: mws,
	}

	return im
}

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

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

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

// Resource resourceful route provides a mapping between HTTP verbs for given the pattern
func (r *Router) Resource(pattern string, fn func(resource *Resource)) {
	if r == nil {
		panic("mux: Resource() called on nil")
	}

	if strings.TrimSpace(pattern) == "" {
		panic("mux: Resource() requires a patter to work")
	}

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

	mws := make([]func(http.Handler) http.Handler, len(r.middlewares))
	copy(mws, r.middlewares)
	fn(&Resource{
		mux:         r.mux,
		pattern:     pattern,
		middlewares: mws,
	})
}

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

	h, pattern := r.mux.Handler(req)
	if pattern == "" {
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	// ensure we run all the middlewares
	h = stack(r.middlewares, h)
	// serve
	h.ServeHTTP(w, req)
}

// stack middlewares(http handler) in order they are passed (FIFO)
func stack(middlewares []func(http.Handler) http.Handler, endpoint http.Handler) http.Handler {
	// Return ahead of time if there aren't any middlewares for the chain
	if len(middlewares) == 0 {
		return endpoint
	}

	// wrap the end handler with the middleware chain
	h := middlewares[len(middlewares)-1](endpoint)
	for i := len(middlewares) - 2; i >= 0; i-- {
		h = middlewares[i](h)
	}

	return h
}