frp/utils/vhost/router.go

113 lines
1.9 KiB
Go
Raw Normal View History

package vhost
import (
"sort"
"strings"
"sync"
)
type VhostRouters struct {
RouterByDomain map[string][]*VhostRouter
mutex sync.RWMutex
}
type VhostRouter struct {
domain string
location string
2017-12-13 03:27:43 +08:00
payload interface{}
}
func NewVhostRouters() *VhostRouters {
return &VhostRouters{
RouterByDomain: make(map[string][]*VhostRouter),
}
}
2017-12-13 03:27:43 +08:00
func (r *VhostRouters) Add(domain, location string, payload interface{}) {
r.mutex.Lock()
defer r.mutex.Unlock()
vrs, found := r.RouterByDomain[domain]
if !found {
vrs = make([]*VhostRouter, 0, 1)
}
2016-12-25 01:53:23 +08:00
vr := &VhostRouter{
domain: domain,
location: location,
2017-12-13 03:27:43 +08:00
payload: payload,
}
2016-12-25 01:53:23 +08:00
vrs = append(vrs, vr)
sort.Sort(sort.Reverse(ByLocation(vrs)))
r.RouterByDomain[domain] = vrs
}
func (r *VhostRouters) Del(domain, location string) {
r.mutex.Lock()
defer r.mutex.Unlock()
vrs, found := r.RouterByDomain[domain]
if !found {
return
}
newVrs := make([]*VhostRouter, 0)
for _, vr := range vrs {
if vr.location != location {
newVrs = append(newVrs, vr)
}
}
r.RouterByDomain[domain] = newVrs
}
2016-12-25 01:53:23 +08:00
func (r *VhostRouters) Get(host, path string) (vr *VhostRouter, exist bool) {
r.mutex.RLock()
defer r.mutex.RUnlock()
2016-12-25 01:53:23 +08:00
vrs, found := r.RouterByDomain[host]
if !found {
return
}
// can't support load balance, will to do
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 *VhostRouters) Exist(host, path string) (vr *VhostRouter, exist bool) {
r.mutex.RLock()
defer r.mutex.RUnlock()
vrs, found := r.RouterByDomain[host]
if !found {
return
}
for _, vr = range vrs {
if path == vr.location {
return vr, true
}
}
return
}
// sort by location
2016-12-25 01:53:23 +08:00
type ByLocation []*VhostRouter
func (a ByLocation) Len() int {
return len(a)
}
func (a ByLocation) Swap(i, j int) {
a[i], a[j] = a[j], a[i]
}
func (a ByLocation) Less(i, j int) bool {
return strings.Compare(a[i].location, a[j].location) < 0
}