frp/pkg/util/vhost/router.go

137 lines
2.4 KiB
Go
Raw Normal View History

package vhost
import (
2024-02-20 12:01:41 +08:00
"cmp"
2019-07-31 00:41:58 +08:00
"errors"
2024-02-20 12:01:41 +08:00
"slices"
"strings"
"sync"
)
2022-08-29 01:02:53 +08:00
var ErrRouterConfigConflict = errors.New("router config conflict")
2019-07-31 00:41:58 +08:00
type routerByHTTPUser map[string][]*Router
2020-05-24 17:48:37 +08:00
type Routers struct {
indexByDomain map[string]routerByHTTPUser
mutex sync.RWMutex
}
2020-05-24 17:48:37 +08:00
type Router struct {
domain string
location string
httpUser string
2017-12-13 03:27:43 +08:00
// store any object here
2017-12-13 03:27:43 +08:00
payload interface{}
}
2020-05-24 17:48:37 +08:00
func NewRouters() *Routers {
return &Routers{
indexByDomain: make(map[string]routerByHTTPUser),
}
}
func (r *Routers) Add(domain, location, httpUser string, payload interface{}) error {
domain = strings.ToLower(domain)
r.mutex.Lock()
defer r.mutex.Unlock()
if _, exist := r.exist(domain, location, httpUser); exist {
2019-07-31 00:41:58 +08:00
return ErrRouterConfigConflict
}
routersByHTTPUser, found := r.indexByDomain[domain]
if !found {
routersByHTTPUser = make(map[string][]*Router)
}
vrs, found := routersByHTTPUser[httpUser]
if !found {
2020-05-24 17:48:37 +08:00
vrs = make([]*Router, 0, 1)
}
2020-05-24 17:48:37 +08:00
vr := &Router{
2016-12-25 01:53:23 +08:00
domain: domain,
location: location,
httpUser: httpUser,
2017-12-13 03:27:43 +08:00
payload: payload,
}
2016-12-25 01:53:23 +08:00
vrs = append(vrs, vr)
2024-02-20 12:01:41 +08:00
slices.SortFunc(vrs, func(a, b *Router) int {
return -cmp.Compare(a.location, b.location)
})
routersByHTTPUser[httpUser] = vrs
r.indexByDomain[domain] = routersByHTTPUser
2019-07-31 00:41:58 +08:00
return nil
}
func (r *Routers) Del(domain, location, httpUser string) {
domain = strings.ToLower(domain)
r.mutex.Lock()
defer r.mutex.Unlock()
routersByHTTPUser, found := r.indexByDomain[domain]
if !found {
return
}
vrs, found := routersByHTTPUser[httpUser]
if !found {
return
}
2020-05-24 17:48:37 +08:00
newVrs := make([]*Router, 0)
for _, vr := range vrs {
if vr.location != location {
newVrs = append(newVrs, vr)
}
}
routersByHTTPUser[httpUser] = newVrs
}
func (r *Routers) Get(host, path, httpUser string) (vr *Router, exist bool) {
host = strings.ToLower(host)
r.mutex.RLock()
defer r.mutex.RUnlock()
routersByHTTPUser, found := r.indexByDomain[host]
if !found {
return
}
vrs, found := routersByHTTPUser[httpUser]
if !found {
return
}
for _, vr = range vrs {
2016-12-25 01:53:23 +08:00
if strings.HasPrefix(path, vr.location) {
return vr, true
}
}
return
}
2016-12-25 01:53:23 +08:00
func (r *Routers) exist(host, path, httpUser string) (route *Router, exist bool) {
routersByHTTPUser, found := r.indexByDomain[host]
if !found {
return
}
routers, found := routersByHTTPUser[httpUser]
if !found {
return
}
for _, route = range routers {
if path == route.location {
return route, true
}
}
return
}