golib/features/mux/mux.go

package mux

import (
	"net/http"
	"path/filepath"
	"strings"

	"github.com/gorilla/mux"
)

type (
	MiddlewareFunc = mux.MiddlewareFunc
	Router         = mux.Router
)

var (
	defaultMethods = []string{http.MethodGet, http.MethodPost}
)

func New() *mux.Router {
	return mux.NewRouter()
}

func RegisterWith(router *mux.Router, path string, handler http.HandlerFunc, methods ...string) {
	r := router.HandleFunc(path, handler)
	if len(methods) == 0 {
		methods = defaultMethods
	}
	r.Methods(methods...)
}

type GroupMux struct {
	router *mux.Router
}

func (g *GroupMux) Register(path string, handler http.Handler, methods ...string) {
	r := g.router.Handle(path, handler)
	if len(methods) == 0 {
		methods = defaultMethods
	}
	r.Methods(methods...)
}

func (g *GroupMux) RegisterFunc(path string, handler http.HandlerFunc, methods ...string) {
	g.Register(path, handler, methods...)
}

func (g *GroupMux) Use(mwf ...mux.MiddlewareFunc) {
	g.router.Use(mwf...)
}

func GroupWith(router *Router, prefix string) *GroupMux {
	return &GroupMux{router: router.PathPrefix(prefix).Subrouter()}
}

func Params(r *http.Request) map[string]string {
	return mux.Vars(r)
}

func RegisterStaticDir(router *mux.Router, path, dir string) {
	if strings.LastIndex(path, "/") == -1 {
		path = path + "/"
	}
	handler := http.StripPrefix(path, http.FileServer(http.Dir(filepath.Join(dir))))
	router.PathPrefix(path).Handler(handler).Methods(http.MethodGet)
}

func RegisterStaticFile(router *mux.Router, path, file string) {
	router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, file)
	}).Methods(http.MethodGet)
}

func CORS() mux.MiddlewareFunc {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			w.Header().Add("Access-Control-Allow-Origin", "*")
			w.Header().Add("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE")
			w.Header().Add("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept, Authorization")
			w.Header().Add("Access-Control-Expose-Headers", "Content-Length, Access-Control-Allow-Origin, Access-Control-Allow-Headers, Cache-Control, Content-Language, Content-Type")
			w.Header().Add("Access-Control-Allow-Credentials", "true")
			next.ServeHTTP(w, req)
		})
	}
}

func NoCache() mux.MiddlewareFunc {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			w.Header().Set("Cache-Control", "no-cache")
			next.ServeHTTP(w, req)
		})
	}
}

func SetAuthenticate(w http.ResponseWriter) {
	w.Header().Set("WWW-Authenticate", `Basic realm="Restricted Area"`)
}

func Authorization(handler func(username, password string) bool) mux.MiddlewareFunc {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			username, password, ok := r.BasicAuth()
			if !ok {
				SetAuthenticate(w)
				http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
				return
			}
			if !handler(username, password) {
				SetAuthenticate(w)
				http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}