support http load balancing

utils/vhost/http.go

@@ -23,7 +23,6 @@ import (
 	"net"
 	"net/http"
 	"strings"
-	"sync"
 	"time"
 
 	frpLog "github.com/fatedier/frp/utils/log"
@@ -32,8 +31,7 @@ import (
 )
 
 var (
-	ErrRouterConfigConflict = errors.New("router config conflict")
-	ErrNoDomain             = errors.New("no such domain")
+	ErrNoDomain = errors.New("no such domain")
 )
 
 func getHostFromAddr(addr string) (host string) {
@@ -51,21 +49,19 @@ type HttpReverseProxyOptions struct {
 }
 
 type HttpReverseProxy struct {
-	proxy *ReverseProxy
-
+	proxy       *ReverseProxy
 	vhostRouter *VhostRouters
 
 	responseHeaderTimeout time.Duration
-	cfgMu                 sync.RWMutex
 }
 
-func NewHttpReverseProxy(option HttpReverseProxyOptions) *HttpReverseProxy {
+func NewHttpReverseProxy(option HttpReverseProxyOptions, vhostRouter *VhostRouters) *HttpReverseProxy {
 	if option.ResponseHeaderTimeoutS <= 0 {
 		option.ResponseHeaderTimeoutS = 60
 	}
 	rp := &HttpReverseProxy{
 		responseHeaderTimeout: time.Duration(option.ResponseHeaderTimeoutS) * time.Second,
-		vhostRouter:           NewVhostRouters(),
+		vhostRouter:           vhostRouter,
 	}
 	proxy := &ReverseProxy{
 		Director: func(req *http.Request) {
@@ -106,21 +102,18 @@ func NewHttpReverseProxy(option HttpReverseProxyOptions) *HttpReverseProxy {
 	return rp
 }
 
+// Register register the route config to reverse proxy
+// reverse proxy will use CreateConnFn from routeCfg to create a connection to the remote service
 func (rp *HttpReverseProxy) Register(routeCfg VhostRouteConfig) error {
-	rp.cfgMu.Lock()
-	defer rp.cfgMu.Unlock()
-	_, ok := rp.vhostRouter.Exist(routeCfg.Domain, routeCfg.Location)
-	if ok {
-		return ErrRouterConfigConflict
-	} else {
-		rp.vhostRouter.Add(routeCfg.Domain, routeCfg.Location, &routeCfg)
+	err := rp.vhostRouter.Add(routeCfg.Domain, routeCfg.Location, &routeCfg)
+	if err != nil {
+		return err
 	}
 	return nil
 }
 
+// UnRegister unregister route config by domain and location
 func (rp *HttpReverseProxy) UnRegister(domain string, location string) {
-	rp.cfgMu.Lock()
-	defer rp.cfgMu.Unlock()
 	rp.vhostRouter.Del(domain, location)
 }
 
@@ -140,6 +133,7 @@ func (rp *HttpReverseProxy) GetHeaders(domain string, location string) (headers
 	return
 }
 
+// CreateConnection create a new connection by route config
 func (rp *HttpReverseProxy) CreateConnection(domain string, location string, remoteAddr string) (net.Conn, error) {
 	vr, ok := rp.getVhost(domain, location)
 	if ok {
@@ -163,10 +157,8 @@ func (rp *HttpReverseProxy) CheckAuth(domain, location, user, passwd string) boo
 	return true
 }
 
+// getVhost get vhost router by domain and location
 func (rp *HttpReverseProxy) getVhost(domain string, location string) (vr *VhostRouter, ok bool) {
-	rp.cfgMu.RLock()
-	defer rp.cfgMu.RUnlock()
-
 	// first we check the full hostname
 	// if not exist, then check the wildcard_domain such as *.example.com
 	vr, ok = rp.vhostRouter.Get(domain, location)

utils/vhost/router.go

@@ -1,11 +1,16 @@
 package vhost
 
 import (
+	"errors"
 	"sort"
 	"strings"
 	"sync"
 )
 
+var (
+	ErrRouterConfigConflict = errors.New("router config conflict")
+)
+
 type VhostRouters struct {
 	RouterByDomain map[string][]*VhostRouter
 	mutex          sync.RWMutex
@@ -24,10 +29,14 @@ func NewVhostRouters() *VhostRouters {
 	}
 }
 
-func (r *VhostRouters) Add(domain, location string, payload interface{}) {
+func (r *VhostRouters) Add(domain, location string, payload interface{}) error {
 	r.mutex.Lock()
 	defer r.mutex.Unlock()
 
+	if _, exist := r.exist(domain, location); exist {
+		return ErrRouterConfigConflict
+	}
+
 	vrs, found := r.RouterByDomain[domain]
 	if !found {
 		vrs = make([]*VhostRouter, 0, 1)
@@ -42,6 +51,7 @@ func (r *VhostRouters) Add(domain, location string, payload interface{}) {
 
 	sort.Sort(sort.Reverse(ByLocation(vrs)))
 	r.RouterByDomain[domain] = vrs
+	return nil
 }
 
 func (r *VhostRouters) Del(domain, location string) {
@@ -80,10 +90,7 @@ func (r *VhostRouters) Get(host, path string) (vr *VhostRouter, exist bool) {
 	return
 }
 
-func (r *VhostRouters) Exist(host, path string) (vr *VhostRouter, exist bool) {
-	r.mutex.RLock()
-	defer r.mutex.RUnlock()
-
+func (r *VhostRouters) exist(host, path string) (vr *VhostRouter, exist bool) {
 	vrs, found := r.RouterByDomain[host]
 	if !found {
 		return

utils/vhost/vhost.go

@@ -15,7 +15,6 @@ package vhost
 import (
 	"fmt"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/fatedier/frp/utils/log"
@@ -35,7 +34,6 @@ type VhostMuxer struct {
 	authFunc       httpAuthFunc
 	rewriteFunc    hostRewriteFunc
 	registryRouter *VhostRouters
-	mutex          sync.RWMutex
 }
 
 func NewVhostMuxer(listener frpNet.Listener, vhostFunc muxFunc, authFunc httpAuthFunc, rewriteFunc hostRewriteFunc, timeout time.Duration) (mux *VhostMuxer, err error) {
@@ -53,6 +51,7 @@ func NewVhostMuxer(listener frpNet.Listener, vhostFunc muxFunc, authFunc httpAut
 
 type CreateConnFunc func(remoteAddr string) (frpNet.Conn, error)
 
+// VhostRouteConfig is the params used to match HTTP requests
 type VhostRouteConfig struct {
 	Domain      string
 	Location    string
@@ -67,14 +66,6 @@ type VhostRouteConfig struct {
 // 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(cfg *VhostRouteConfig) (l *Listener, err error) {
-	v.mutex.Lock()
-	defer v.mutex.Unlock()
-
-	_, ok := v.registryRouter.Exist(cfg.Domain, cfg.Location)
-	if ok {
-		return nil, fmt.Errorf("hostname [%s] location [%s] is already registered", cfg.Domain, cfg.Location)
-	}
-
 	l = &Listener{
 		name:        cfg.Domain,
 		location:    cfg.Location,
@@ -85,14 +76,14 @@ func (v *VhostMuxer) Listen(cfg *VhostRouteConfig) (l *Listener, err error) {
 		accept:      make(chan frpNet.Conn),
 		Logger:      log.NewPrefixLogger(""),
 	}
-	v.registryRouter.Add(cfg.Domain, cfg.Location, l)
+	err = v.registryRouter.Add(cfg.Domain, cfg.Location, l)
+	if err != nil {
+		return
+	}
 	return l, nil
 }
 
 func (v *VhostMuxer) getListener(name, path string) (l *Listener, exist bool) {
-	v.mutex.RLock()
-	defer v.mutex.RUnlock()
-
 	// 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)