feat: allow multiple duplicate proxies registered with tcpmux for load balancing

server/service.go

@@ -108,7 +108,7 @@ func NewService(cfg config.ServerCommonConf) (svr *Service, err error) {
 			TcpPortManager: ports.NewPortManager("tcp", cfg.ProxyBindAddr, cfg.AllowPorts),
 			UdpPortManager: ports.NewPortManager("udp", cfg.ProxyBindAddr, cfg.AllowPorts),
 		},
-		httpVhostRouter: vhost.NewVhostRouters(),
+		httpVhostRouter: vhost.NewVhostRouters(false),
 		authVerifier:    auth.NewAuthVerifier(cfg.AuthServerConfig),
 		tlsConfig:       generateTLSConfig(),
 		cfg:             cfg,

utils/tcpmux/httpconnect.go

@@ -31,7 +31,7 @@ type HttpConnectTcpMuxer struct {
 }
 
 func NewHttpConnectTcpMuxer(listener net.Listener, timeout time.Duration) (*HttpConnectTcpMuxer, error) {
-	mux, err := vhost.NewVhostMuxer(listener, getHostFromHttpConnect, nil, sendHttpOk, nil, timeout)
+	mux, err := vhost.NewVhostMuxer(listener, getHostFromHttpConnect, nil, sendHttpOk, nil, timeout, true)
 	return &HttpConnectTcpMuxer{mux}, err
 }
 

utils/vhost/http.go

@@ -110,7 +110,7 @@ func (rp *HttpReverseProxy) UnRegister(domain string, location string) {
 func (rp *HttpReverseProxy) GetRealHost(domain string, location string) (host string) {
 	vr, ok := rp.getVhost(domain, location)
 	if ok {
-		host = vr.payload.(*VhostRouteConfig).RewriteHost
+		host = vr.getPayload().(*VhostRouteConfig).RewriteHost
 	}
 	return
 }
@@ -118,7 +118,7 @@ func (rp *HttpReverseProxy) GetRealHost(domain string, location string) (host st
 func (rp *HttpReverseProxy) GetHeaders(domain string, location string) (headers map[string]string) {
 	vr, ok := rp.getVhost(domain, location)
 	if ok {
-		headers = vr.payload.(*VhostRouteConfig).Headers
+		headers = vr.getPayload().(*VhostRouteConfig).Headers
 	}
 	return
 }
@@ -127,7 +127,7 @@ func (rp *HttpReverseProxy) GetHeaders(domain string, location string) (headers
 func (rp *HttpReverseProxy) CreateConnection(domain string, location string, remoteAddr string) (net.Conn, error) {
 	vr, ok := rp.getVhost(domain, location)
 	if ok {
-		fn := vr.payload.(*VhostRouteConfig).CreateConnFn
+		fn := vr.getPayload().(*VhostRouteConfig).CreateConnFn
 		if fn != nil {
 			return fn(remoteAddr)
 		}
@@ -138,8 +138,9 @@ func (rp *HttpReverseProxy) CreateConnection(domain string, location string, rem
 func (rp *HttpReverseProxy) CheckAuth(domain, location, user, passwd string) bool {
 	vr, ok := rp.getVhost(domain, location)
 	if ok {
-		checkUser := vr.payload.(*VhostRouteConfig).Username
+		routeCfg := vr.getPayload().(*VhostRouteConfig)
-		checkPasswd := vr.payload.(*VhostRouteConfig).Password
+		checkUser := routeCfg.Username
+		checkPasswd := routeCfg.Password
 		if (checkUser != "" || checkPasswd != "") && (checkUser != user || checkPasswd != passwd) {
 			return false
 		}

utils/vhost/https.go

@@ -48,7 +48,7 @@ type HttpsMuxer struct {
 }
 
 func NewHttpsMuxer(listener net.Listener, timeout time.Duration) (*HttpsMuxer, error) {
-	mux, err := NewVhostMuxer(listener, GetHttpsHostname, nil, nil, nil, timeout)
+	mux, err := NewVhostMuxer(listener, GetHttpsHostname, nil, nil, nil, timeout, false)
 	return &HttpsMuxer{mux}, err
 }
 

utils/vhost/router.go

@@ -2,6 +2,7 @@ package vhost
 
 import (
 	"errors"
+	"math/rand"
 	"sort"
 	"strings"
 	"sync"
@@ -12,20 +13,30 @@ var (
 )
 
 type VhostRouters struct {
-	RouterByDomain map[string][]*VhostRouter
+	RouterByDomain  map[string][]*VhostRouter
-	mutex          sync.RWMutex
+	allowDuplicates bool
+	mutex           sync.RWMutex
 }
 
 type VhostRouter struct {
 	domain   string
 	location string
 
-	payload interface{}
+	allowDuplicates bool
+	payloads        []interface{}
 }
 
-func NewVhostRouters() *VhostRouters {
+func (vr *VhostRouter) getPayload() interface{} {
+	if !vr.allowDuplicates {
+		return vr.payloads[0]
+	}
+	return vr.payloads[rand.Intn(len(vr.payloads))]
+}
+
+func NewVhostRouters(allowDuplicates bool) *VhostRouters {
 	return &VhostRouters{
-		RouterByDomain: make(map[string][]*VhostRouter),
+		allowDuplicates: allowDuplicates,
+		RouterByDomain:  make(map[string][]*VhostRouter),
 	}
 }
 
@@ -33,8 +44,12 @@ func (r *VhostRouters) Add(domain, location string, payload interface{}) error {
 	r.mutex.Lock()
 	defer r.mutex.Unlock()
 
-	if _, exist := r.exist(domain, location); exist {
+	if vr, exist := r.exist(domain, location); exist {
-		return ErrRouterConfigConflict
+		if !r.allowDuplicates {
+			return ErrRouterConfigConflict
+		}
+		vr.payloads = append(vr.payloads, payload)
+		return nil
 	}
 
 	vrs, found := r.RouterByDomain[domain]
@@ -43,10 +58,12 @@ func (r *VhostRouters) Add(domain, location string, payload interface{}) error {
 	}
 
 	vr := &VhostRouter{
-		domain:   domain,
+		domain:          domain,
-		location: location,
+		location:        location,
-		payload:  payload,
+		allowDuplicates: r.allowDuplicates,
+		payloads:        make([]interface{}, 1),
 	}
+	vr.payloads[0] = payload
 	vrs = append(vrs, vr)
 
 	sort.Sort(sort.Reverse(ByLocation(vrs)))
@@ -68,7 +85,41 @@ func (r *VhostRouters) Del(domain, location string) {
 			newVrs = append(newVrs, vr)
 		}
 	}
-	r.RouterByDomain[domain] = newVrs
+	if len(newVrs) == 0 {
+		delete(r.RouterByDomain, domain)
+	} else {
+		r.RouterByDomain[domain] = newVrs
+	}
+}
+
+func (r *VhostRouters) DelPayloadFromLocation(domain, location string, payload interface{}) {
+	r.mutex.Lock()
+
+	vrs, found := r.RouterByDomain[domain]
+	if !found {
+		r.mutex.Unlock()
+		return
+	}
+	newPayloadsLen := -1
+	for _, vr := range vrs {
+		if vr.location == location {
+			newPayloads := make([]interface{}, 0)
+			for _, payloadIter := range vr.payloads {
+				if payloadIter != payload {
+					newPayloads = append(newPayloads, payloadIter)
+				}
+			}
+			vr.payloads = newPayloads
+			newPayloadsLen = len(newPayloads)
+			break
+		}
+	}
+
+	r.mutex.Unlock()
+
+	if newPayloadsLen == 0 {
+		r.Del(domain, location)
+	}
 }
 
 func (r *VhostRouters) Get(host, path string) (vr *VhostRouter, exist bool) {
@@ -80,7 +131,6 @@ func (r *VhostRouters) Get(host, path string) (vr *VhostRouter, exist bool) {
 		return
 	}
 
-	// can't support load balance, will to do
 	for _, vr = range vrs {
 		if strings.HasPrefix(path, vr.location) {
 			return vr, true

utils/vhost/vhost.go

@@ -41,7 +41,7 @@ type VhostMuxer struct {
 	registryRouter *VhostRouters
 }
 
-func NewVhostMuxer(listener net.Listener, vhostFunc muxFunc, authFunc httpAuthFunc, successFunc successFunc, rewriteFunc hostRewriteFunc, timeout time.Duration) (mux *VhostMuxer, err error) {
+func NewVhostMuxer(listener net.Listener, vhostFunc muxFunc, authFunc httpAuthFunc, successFunc successFunc, rewriteFunc hostRewriteFunc, timeout time.Duration, allowDuplicates bool) (mux *VhostMuxer, err error) {
 	mux = &VhostMuxer{
 		listener:       listener,
 		timeout:        timeout,
@@ -49,7 +49,7 @@ func NewVhostMuxer(listener net.Listener, vhostFunc muxFunc, authFunc httpAuthFu
 		authFunc:       authFunc,
 		successFunc:    successFunc,
 		rewriteFunc:    rewriteFunc,
-		registryRouter: NewVhostRouters(),
+		registryRouter: NewVhostRouters(allowDuplicates),
 	}
 	go mux.run()
 	return mux, nil
@@ -94,7 +94,7 @@ func (v *VhostMuxer) getListener(name, path string) (l *Listener, exist bool) {
 	// if not exist, then check the wildcard_domain such as *.example.com
 	vr, found := v.registryRouter.Get(name, path)
 	if found {
-		return vr.payload.(*Listener), true
+		return vr.getPayload().(*Listener), true
 	}
 
 	domainSplit := strings.Split(name, ".")
@@ -112,7 +112,7 @@ func (v *VhostMuxer) getListener(name, path string) (l *Listener, exist bool) {
 
 		vr, found = v.registryRouter.Get(name, path)
 		if found {
-			return vr.payload.(*Listener), true
+			return vr.getPayload().(*Listener), true
 		}
 		domainSplit = domainSplit[1:]
 	}
@@ -223,7 +223,7 @@ func (l *Listener) Accept() (net.Conn, error) {
 }
 
 func (l *Listener) Close() error {
-	l.mux.registryRouter.Del(l.name, l.location)
+	l.mux.registryRouter.DelPayloadFromLocation(l.name, l.location, l)
 	close(l.accept)
 	return nil
 }