kit/rest/router/patrouter.go

package router

import (
	"errors"
	"git.diulo.com/mogfee/kit/rest/httpx"
	"git.diulo.com/mogfee/kit/rest/pathvar"
	"git.diulo.com/mogfee/kit/rest/search"
	"net/http"
	"path"
	"strings"
)

const (
	allowHeader          = "Allow"
	allowMethodSeparator = ", "
)

var (
	ErrInvalidMethod = errors.New("not a valid http method")
	ErrInvalidPath   = errors.New("path must begin with '/'")
)

type patRouter struct {
	trees      map[string]*search.TreeNode
	notFound   http.Handler
	notAllowed http.Handler
}

func NewRouter() httpx.Router {
	return &patRouter{
		trees: make(map[string]*search.TreeNode),
	}
}

func (p *patRouter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	reqPath := path.Clean(r.URL.Path)
	if tree, ok := p.trees[r.Method]; ok {
		if result, ok := tree.Search(reqPath); ok {
			if len(result.Params) > 0 {
				r = pathvar.WithVars(r, result.Params)
			}
			result.Item.(http.Handler).ServeHTTP(w, r)
			return
		}
	}
	allows, ok := p.methodsAllowed(r.Method, reqPath)
	if !ok {
		p.handleNotFound(w, r)
		return
	}
	if p.notAllowed != nil {
		p.notAllowed.ServeHTTP(w, r)
	} else {
		w.Header().Set(allowHeader, allows)
		w.WriteHeader(http.StatusMethodNotAllowed)
	}
}

func (p *patRouter) Handle(method string, reqPath string, handler http.Handler) error {
	if !validMethod(method) {
		return ErrInvalidMethod
	}
	if len(reqPath) == 0 || reqPath[0] != '/' {
		return ErrInvalidPath
	}
	cleanPath := path.Clean(reqPath)
	tree, ok := p.trees[method]
	if ok {
		return tree.Add(cleanPath, handler)
	}

	tree = search.NewTree()
	p.trees[method] = tree
	return tree.Add(cleanPath, handler)
}

func (p *patRouter) SetNotFoundHandler(handler http.Handler) {
	p.notFound = handler
}

func (p *patRouter) SetNotAllowedHandler(handler http.Handler) {
	p.notAllowed = handler
}
func (p *patRouter) handleNotFound(w http.ResponseWriter, r *http.Request) {
	if p.notFound != nil {
		p.notFound.ServeHTTP(w, r)
	} else {
		http.NotFound(w, r)
	}
}
func (p *patRouter) methodsAllowed(method, path string) (string, bool) {
	var allows []string
	for treeMethod, tree := range p.trees {
		if treeMethod == method {
			continue
		}
		_, ok := tree.Search(path)
		if ok {
			allows = append(allows, treeMethod)
		}
	}
	if len(allows) > 0 {
		return strings.Join(allows, allowMethodSeparator), true
	}
	return "", false
}
func validMethod(method string) bool {
	return method == http.MethodDelete || method == http.MethodGet ||
		method == http.MethodHead || method == http.MethodOptions ||
		method == http.MethodPatch || method == http.MethodPost ||
		method == http.MethodPut
}