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 }