router/tree.go

88 lines
1.6 KiB
Go
Executable File

package router
import (
"errors"
"net/http"
"path"
"strings"
)
var Wildcard = "{}"
var WildcardHeader = "__wildcard_value__"
type tree struct {
next map[string]*tree
handler http.HandlerFunc
}
func newTree() *tree {
return &tree{
next: make(map[string]*tree),
handler: nil,
}
}
func (t *tree) Lookup(path string) http.HandlerFunc {
if path == "/" || path == "" {
if t.handler != nil {
return t.handler
}
if n, ok := t.next[Wildcard+Wildcard]; t.handler == nil && ok {
foo := n.handler
return foo
}
return nil
}
key, following := nextPathSegment(path)
n, ok := t.next[key]
if ok {
n2 := n.Lookup(following)
if n2 != nil {
return n2
}
}
if n, ok := t.next[Wildcard]; ok {
foo := n.Lookup(following)
if foo != nil {
return func(w http.ResponseWriter, r *http.Request) {
r.Header.Add(WildcardHeader, key)
foo(w, r)
}
}
} else if n, ok := t.next[Wildcard+Wildcard]; ok {
foo := n.handler
if foo != nil {
return func(w http.ResponseWriter, r *http.Request) {
r.Header.Add(WildcardHeader, key)
foo(w, r)
}
}
}
return nil
}
func (t *tree) Insert(path string, foo http.HandlerFunc) error {
if path == "/" {
if t.handler != nil {
return errors.New("occupied path")
}
t.handler = foo
return nil
}
key, following := nextPathSegment(path)
_, ok := t.next[key]
if !ok {
t.next[key] = newTree()
}
return t.next[key].Insert(following, foo)
}
func nextPathSegment(p string) (string, string) {
p = path.Clean("/" + p)
i := strings.Index(p[1:], "/") + 1
if i <= 0 {
return p[1:], "/"
}
return p[1:i], p[i:]
}