This commit is contained in:
fatedier
2016-12-25 01:53:23 +08:00
parent 59a34b81e0
commit eb4f779384
11 changed files with 180 additions and 156 deletions

View File

@@ -12,7 +12,6 @@ type VhostRouters struct {
}
type VhostRouter struct {
name string
domain string
location string
listener *Listener
@@ -24,6 +23,65 @@ func NewVhostRouters() *VhostRouters {
}
}
func (r *VhostRouters) Add(domain, location string, l *Listener) {
r.mutex.Lock()
defer r.mutex.Unlock()
vrs, found := r.RouterByDomain[domain]
if !found {
vrs = make([]*VhostRouter, 0)
}
vr := &VhostRouter{
domain: domain,
location: location,
listener: l,
}
vrs = append(vrs, vr)
sort.Reverse(ByLocation(vrs))
r.RouterByDomain[domain] = vrs
}
func (r *VhostRouters) Del(l *Listener) {
r.mutex.Lock()
defer r.mutex.Unlock()
vrs, found := r.RouterByDomain[l.name]
if !found {
return
}
for i, vr := range vrs {
if vr.listener == l {
if len(vrs) > i+1 {
r.RouterByDomain[l.name] = append(vrs[:i], vrs[i+1:]...)
} else {
r.RouterByDomain[l.name] = vrs[:i]
}
}
}
}
func (r *VhostRouters) Get(host, path string) (vr *VhostRouter, exist bool) {
r.mutex.RLock()
defer r.mutex.RUnlock()
vrs, found := r.RouterByDomain[host]
if !found {
return
}
//can't support load balance,will to do
for _, vr = range vrs {
if strings.HasPrefix(path, vr.location) {
return vr, true
}
}
return
}
//sort by location
type ByLocation []*VhostRouter
@@ -36,72 +94,3 @@ func (a ByLocation) Swap(i, j int) {
func (a ByLocation) Less(i, j int) bool {
return strings.Compare(a[i].location, a[j].location) < 0
}
func (r *VhostRouters) add(name, domain string, locations []string, l *Listener) {
r.mutex.Lock()
defer r.mutex.Unlock()
vrs, found := r.RouterByDomain[domain]
if !found {
vrs = make([]*VhostRouter, 0)
}
for _, loc := range locations {
vr := &VhostRouter{
name: name,
domain: domain,
location: loc,
listener: l,
}
vrs = append(vrs, vr)
}
sort.Reverse(ByLocation(vrs))
r.RouterByDomain[domain] = vrs
}
func (r *VhostRouters) del(l *Listener) {
r.mutex.Lock()
defer r.mutex.Unlock()
vrs, found := r.RouterByDomain[l.domain]
if !found {
return
}
for i, vr := range vrs {
if vr.listener == l {
if len(vrs) > i+1 {
r.RouterByDomain[l.domain] = append(vrs[:i], vrs[i+1:]...)
} else {
r.RouterByDomain[l.domain] = vrs[:i]
}
}
}
}
func (r *VhostRouters) get(rname string) (vr *VhostRouter, exist bool) {
r.mutex.RLock()
defer r.mutex.RUnlock()
var domain, url string
tmparray := strings.SplitN(rname, ":", 2)
if len(tmparray) == 2 {
domain = tmparray[0]
url = tmparray[1]
}
vrs, found := r.RouterByDomain[domain]
if !found {
return
}
//can't support load balance,will to do
for _, vr = range vrs {
if strings.HasPrefix(url, vr.location) {
return vr, true
}
}
return
}

View File

@@ -21,9 +21,7 @@ import (
"sync"
"time"
"github.com/fatedier/frp/src/models/config"
"github.com/fatedier/frp/src/utils/conn"
"github.com/fatedier/frp/src/utils/log"
)
type muxFunc func(*conn.Conn) (net.Conn, map[string]string, error)
@@ -54,59 +52,42 @@ func NewVhostMuxer(listener *conn.Listener, vhostFunc muxFunc, authFunc httpAuth
}
// listen for a new domain name, if rewriteHost is not empty and rewriteFunc is not nil, then rewrite the host header to rewriteHost
func (v *VhostMuxer) Listen(p *config.ProxyServerConf) (ls []*Listener) {
func (v *VhostMuxer) Listen(name, location, rewriteHost, userName, passWord string) (l *Listener, err error) {
v.mutex.Lock()
defer v.mutex.Unlock()
ls = make([]*Listener, 0)
for _, domain := range p.CustomDomains {
l := &Listener{
name: p.Name,
domain: domain,
locations: p.Locations,
rewriteHost: p.HostHeaderRewrite,
userName: p.HttpUserName,
passWord: p.HttpPassWord,
mux: v,
accept: make(chan *conn.Conn),
}
v.registryRouter.add(p.Name, domain, p.Locations, l)
ls = append(ls, l)
l = &Listener{
name: name,
rewriteHost: rewriteHost,
userName: userName,
passWord: passWord,
mux: v,
accept: make(chan *conn.Conn),
}
return ls
v.registryRouter.Add(name, location, l)
return l, nil
}
func (v *VhostMuxer) getListener(reqInfoMap map[string]string) (l *Listener, exist bool) {
func (v *VhostMuxer) getListener(name, path string) (l *Listener, exist bool) {
v.mutex.RLock()
defer v.mutex.RUnlock()
//host
name := strings.ToLower(reqInfoMap["Host"])
// http
scheme := strings.ToLower(reqInfoMap["Scheme"])
if scheme == "http" || scheme == "" {
name = name + ":" + reqInfoMap["Path"]
}
// // first we check the full hostname
vr, found := v.registryRouter.get(name)
// first we check the full hostname
// if not exist, then check the wildcard_domain such as *.example.com
vr, found := v.registryRouter.Get(name, path)
if found {
return vr.listener, true
}
//if not exist, then check the wildcard_domain such as *.example.com
domainSplit := strings.Split(name, ".")
if len(domainSplit) < 3 {
log.Warn("can't found the router for %s", name)
return l, false
}
domainSplit[0] = "*"
name = strings.Join(domainSplit, ".")
vr, found = v.registryRouter.get(name)
vr, found = v.registryRouter.Get(name, path)
if !found {
log.Warn("can't found the router for %s", name)
return
}
@@ -135,7 +116,9 @@ func (v *VhostMuxer) handle(c *conn.Conn) {
return
}
l, ok := v.getListener(reqInfoMap)
name := strings.ToLower(reqInfoMap["Host"])
path := strings.ToLower(reqInfoMap["Path"])
l, ok := v.getListener(name, path)
if !ok {
c.Close()
return
@@ -165,8 +148,6 @@ func (v *VhostMuxer) handle(c *conn.Conn) {
type Listener struct {
name string
domain string
locations []string
rewriteHost string
userName string
passWord string
@@ -194,7 +175,7 @@ func (l *Listener) Accept() (*conn.Conn, error) {
}
func (l *Listener) Close() error {
l.mux.registryRouter.del(l)
l.mux.registryRouter.Del(l)
close(l.accept)
return nil
}