violet/router/router.go
2023-07-22 01:11:47 +01:00

120 lines
2.8 KiB
Go

package router
import (
"fmt"
"github.com/1f349/violet/proxy"
"github.com/1f349/violet/target"
"github.com/1f349/violet/utils"
"github.com/MrMelon54/trie"
"net/http"
"strings"
)
type Router struct {
route map[string]*trie.Trie[target.Route]
redirect map[string]*trie.Trie[target.Redirect]
notFound http.Handler
proxy *proxy.HybridTransport
}
func New(proxy *proxy.HybridTransport) *Router {
return &Router{
route: make(map[string]*trie.Trie[target.Route]),
redirect: make(map[string]*trie.Trie[target.Redirect]),
notFound: http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
_, _ = fmt.Fprintf(rw, "%d %s\n", http.StatusNotFound, http.StatusText(http.StatusNotFound))
}),
proxy: proxy,
}
}
func (r *Router) hostRoute(host string) *trie.Trie[target.Route] {
h := r.route[host]
if h == nil {
h = &trie.Trie[target.Route]{}
r.route[host] = h
}
return h
}
func (r *Router) hostRedirect(host string) *trie.Trie[target.Redirect] {
h := r.redirect[host]
if h == nil {
h = &trie.Trie[target.Redirect]{}
r.redirect[host] = h
}
return h
}
func (r *Router) AddRoute(t target.Route) {
t.Proxy = r.proxy
host, path := utils.SplitHostPath(t.Src)
r.hostRoute(host).PutString(path, t)
}
func (r *Router) AddRedirect(t target.Redirect) {
host, path := utils.SplitHostPath(t.Src)
r.hostRedirect(host).PutString(path, t)
}
func (r *Router) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
if req.URL.Path == "" {
req.URL.Path = "/"
}
host, _, _ := utils.SplitDomainPort(req.Host, 0)
if r.serveRedirectHTTP(rw, req, host) {
return
}
if r.serveRouteHTTP(rw, req, host) {
return
}
parentHostDot := strings.IndexByte(host, '.')
if parentHostDot == -1 {
r.notFound.ServeHTTP(rw, req)
return
}
wildcardHost := "*" + host[parentHostDot:]
if r.serveRedirectHTTP(rw, req, wildcardHost) {
return
}
if r.serveRouteHTTP(rw, req, wildcardHost) {
return
}
utils.RespondVioletError(rw, http.StatusTeapot, "No route")
}
func (r *Router) serveRouteHTTP(rw http.ResponseWriter, req *http.Request, host string) bool {
h := r.route[host]
if h != nil {
pairs := h.GetAllKeyValues([]byte(req.URL.Path))
for i := len(pairs) - 1; i >= 0; i-- {
if pairs[i].Value.HasFlag(target.FlagPre) || pairs[i].Key == req.URL.Path {
req.URL.Path = strings.TrimPrefix(req.URL.Path, pairs[i].Key)
pairs[i].Value.ServeHTTP(rw, req)
return true
}
}
}
return false
}
func (r *Router) serveRedirectHTTP(rw http.ResponseWriter, req *http.Request, host string) bool {
h := r.redirect[host]
if h != nil {
pairs := h.GetAllKeyValues([]byte(req.URL.Path))
for i := len(pairs) - 1; i >= 0; i-- {
if pairs[i].Value.Flags.HasFlag(target.FlagPre) || pairs[i].Key == req.URL.Path {
req.URL.Path = strings.TrimPrefix(req.URL.Path, pairs[i].Key)
pairs[i].Value.ServeHTTP(rw, req)
return true
}
}
}
return false
}