refactor: clean up code (#5308)

pkg/util/vhost/http.go

@@ -24,7 +24,6 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
-	"strings"
 	"time"
 
 	libio "github.com/fatedier/golib/io"
@@ -160,7 +159,7 @@ func (rp *HTTPReverseProxy) UnRegister(routeCfg RouteConfig) {
 }
 
 func (rp *HTTPReverseProxy) GetRouteConfig(domain, location, routeByHTTPUser string) *RouteConfig {
-	vr, ok := rp.getVhost(domain, location, routeByHTTPUser)
+	vr, ok := rp.vhostRouter.getByRoute(domain, location, routeByHTTPUser)
 	if ok {
 		log.Debugf("get new http request host [%s] path [%s] httpuser [%s]", domain, location, routeByHTTPUser)
 		return vr.payload.(*RouteConfig)
@@ -171,7 +170,7 @@ func (rp *HTTPReverseProxy) GetRouteConfig(domain, location, routeByHTTPUser str
 // CreateConnection create a new connection by route config
 func (rp *HTTPReverseProxy) CreateConnection(reqRouteInfo *RequestRouteInfo, byEndpoint bool) (net.Conn, error) {
 	host, _ := httppkg.CanonicalHost(reqRouteInfo.Host)
-	vr, ok := rp.getVhost(host, reqRouteInfo.URL, reqRouteInfo.HTTPUser)
+	vr, ok := rp.vhostRouter.getByRoute(host, reqRouteInfo.URL, reqRouteInfo.HTTPUser)
 	if ok {
 		if byEndpoint {
 			fn := vr.payload.(*RouteConfig).CreateConnByEndpointFn
@@ -208,50 +207,6 @@ func checkRouteAuthByRequest(req *http.Request, rc *RouteConfig) bool {
 	return ok && user == rc.Username && passwd == rc.Password
 }
 
-// getVhost tries to get vhost router by route policy.
-func (rp *HTTPReverseProxy) getVhost(domain, location, routeByHTTPUser string) (*Router, bool) {
-	findRouter := func(inDomain, inLocation, inRouteByHTTPUser string) (*Router, bool) {
-		vr, ok := rp.vhostRouter.Get(inDomain, inLocation, inRouteByHTTPUser)
-		if ok {
-			return vr, ok
-		}
-		// Try to check if there is one proxy that doesn't specify routerByHTTPUser, it means match all.
-		vr, ok = rp.vhostRouter.Get(inDomain, inLocation, "")
-		if ok {
-			return vr, ok
-		}
-		return nil, false
-	}
-
-	// First we check the full hostname
-	// if not exist, then check the wildcard_domain such as *.example.com
-	vr, ok := findRouter(domain, location, routeByHTTPUser)
-	if ok {
-		return vr, ok
-	}
-
-	// e.g. domain = test.example.com, try to match wildcard domains.
-	// *.example.com
-	// *.com
-	domainSplit := strings.Split(domain, ".")
-	for len(domainSplit) >= 3 {
-		domainSplit[0] = "*"
-		domain = strings.Join(domainSplit, ".")
-		vr, ok = findRouter(domain, location, routeByHTTPUser)
-		if ok {
-			return vr, true
-		}
-		domainSplit = domainSplit[1:]
-	}
-
-	// Finally, try to check if there is one proxy that domain is "*" means match all domains.
-	vr, ok = findRouter("*", location, routeByHTTPUser)
-	if ok {
-		return vr, true
-	}
-	return nil, false
-}
-
 func (rp *HTTPReverseProxy) connectHandler(rw http.ResponseWriter, req *http.Request) {
 	hj, ok := rw.(http.Hijacker)
 	if !ok {

pkg/util/vhost/https.go

@@ -29,11 +29,11 @@ type HTTPSMuxer struct {
 
 func NewHTTPSMuxer(listener net.Listener, timeout time.Duration) (*HTTPSMuxer, error) {
 	mux, err := NewMuxer(listener, GetHTTPSHostname, timeout)
-	mux.SetFailHookFunc(vhostFailed)
 	if err != nil {
 		return nil, err
 	}
-	return &HTTPSMuxer{mux}, err
+	mux.SetFailHookFunc(vhostFailed)
+	return &HTTPSMuxer{mux}, nil
 }
 
 func GetHTTPSHostname(c net.Conn) (_ net.Conn, _ map[string]string, err error) {

pkg/util/vhost/https_test.go

@@ -16,23 +16,53 @@ func TestGetHTTPSHostname(t *testing.T) {
 	require.NoError(err)
 	defer l.Close()
 
-	var conn net.Conn
+	connCh := make(chan net.Conn, 1)
+	acceptErrCh := make(chan error, 1)
 	go func() {
-		conn, _ = l.Accept()
-		require.NotNil(conn)
+		conn, err := l.Accept()
+		if err != nil {
+			acceptErrCh <- err
+			return
+		}
+		connCh <- conn
 	}()
 
+	clientErrCh := make(chan error, 1)
 	go func() {
 		time.Sleep(100 * time.Millisecond)
-		tls.Dial("tcp", l.Addr().String(), &tls.Config{
+		conn, err := tls.Dial("tcp", l.Addr().String(), &tls.Config{
 			InsecureSkipVerify: true,
 			ServerName:         "example.com",
 		})
+		if conn != nil {
+			_ = conn.Close()
+		}
+		clientErrCh <- err
 	}()
 
-	time.Sleep(200 * time.Millisecond)
-	_, infos, err := GetHTTPSHostname(conn)
+	var conn net.Conn
+	select {
+	case conn = <-connCh:
+	case err := <-acceptErrCh:
+		require.NoError(err)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for accepted connection")
+	}
+	require.NotNil(conn)
+
+	serverConn, infos, err := GetHTTPSHostname(conn)
+	if serverConn != nil {
+		_ = serverConn.Close()
+	} else {
+		_ = conn.Close()
+	}
 	require.NoError(err)
 	require.Equal("example.com", infos["Host"])
 	require.Equal("https", infos["Scheme"])
+
+	select {
+	case <-clientErrCh:
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for TLS client")
+	}
 }

pkg/util/vhost/router.go

@@ -93,12 +93,19 @@ func (r *Routers) Del(domain, location, httpUser string) {
 	routersByHTTPUser[httpUser] = newVrs
 }
 
+// Get returns the best location match for an exact host and exact HTTP user.
+// It does not apply all-users, wildcard-domain, or catch-all-domain fallback.
 func (r *Routers) Get(host, path, httpUser string) (vr *Router, exist bool) {
 	host = strings.ToLower(host)
 
 	r.mutex.RLock()
 	defer r.mutex.RUnlock()
 
+	return r.getLocked(host, path, httpUser)
+}
+
+// getLocked performs an exact-host lookup; host must already be lower-cased.
+func (r *Routers) getLocked(host, path, httpUser string) (vr *Router, exist bool) {
 	routersByHTTPUser, found := r.indexByDomain[host]
 	if !found {
 		return
@@ -117,6 +124,49 @@ func (r *Routers) Get(host, path, httpUser string) (vr *Router, exist bool) {
 	return
 }
 
+func (r *Routers) getByRoute(host, path, httpUser string) (*Router, bool) {
+	host = strings.ToLower(host)
+
+	r.mutex.RLock()
+	defer r.mutex.RUnlock()
+
+	// First we check the full hostname; if it doesn't exist, then check wildcard domains.
+	// For example, test.example.com checks *.example.com before falling back to "*".
+	vr, ok := r.getExactOrAllUsersLocked(host, path, httpUser)
+	if ok {
+		return vr, true
+	}
+
+	hostSplit := strings.Split(host, ".")
+	// Keep two-label hosts out of the wildcard walk, so example.com does not match *.com.
+	for len(hostSplit) >= 3 {
+		// Replace the leftmost remaining label with the wildcard marker.
+		hostSplit[0] = "*"
+		host = strings.Join(hostSplit, ".")
+		vr, ok = r.getExactOrAllUsersLocked(host, path, httpUser)
+		if ok {
+			return vr, true
+		}
+		hostSplit = hostSplit[1:]
+	}
+
+	// Finally, try to check if there is one proxy whose domain is "*", which means match all domains.
+	return r.getExactOrAllUsersLocked("*", path, httpUser)
+}
+
+func (r *Routers) getExactOrAllUsersLocked(host, path, httpUser string) (*Router, bool) {
+	vr, ok := r.getLocked(host, path, httpUser)
+	if ok {
+		return vr, true
+	}
+	// Try to check if there is one proxy that doesn't specify routeByHTTPUser, it means match all.
+	vr, ok = r.getLocked(host, path, "")
+	if ok {
+		return vr, true
+	}
+	return nil, false
+}
+
 func (r *Routers) exist(host, path, httpUser string) (route *Router, exist bool) {
 	routersByHTTPUser, found := r.indexByDomain[host]
 	if !found {

pkg/util/vhost/router_test.go

@@ -0,0 +1,257 @@
+// Copyright 2017 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package vhost
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRoutersGet(t *testing.T) {
+	routers := NewRouters()
+	require.NoError(t, routers.Add("example.com", "/api", "alice", "exact-user"))
+	require.NoError(t, routers.Add("example.com", "/public", "", "exact-all-users"))
+	require.NoError(t, routers.Add("*.example.com", "/api", "", "wildcard-all-users"))
+	require.NoError(t, routers.Add("*", "/api", "", "all-domains"))
+
+	t.Run("exact host and user match with normalized domain", func(t *testing.T) {
+		router, ok := routers.Get("EXAMPLE.COM", "/api/users", "alice")
+		require.True(t, ok)
+		require.Equal(t, "exact-user", router.payload)
+	})
+
+	t.Run("exact host and empty user match", func(t *testing.T) {
+		router, ok := routers.Get("EXAMPLE.COM", "/public/docs", "")
+		require.True(t, ok)
+		require.Equal(t, "exact-all-users", router.payload)
+	})
+
+	t.Run("does not fall back from named user to empty user", func(t *testing.T) {
+		// Get intentionally requires an exact HTTP user; route-level fallbacks live in getByRoute.
+		_, ok := routers.Get("EXAMPLE.COM", "/public/docs", "alice")
+		require.False(t, ok)
+	})
+
+	t.Run("does not fall back to wildcard domains", func(t *testing.T) {
+		_, ok := routers.Get("foo.example.com", "/api/users", "")
+		require.False(t, ok)
+	})
+
+	t.Run("does not fall back to catch-all domain", func(t *testing.T) {
+		_, ok := routers.Get("missing.test", "/api/users", "")
+		require.False(t, ok)
+	})
+}
+
+func TestRoutersGetByRoute(t *testing.T) {
+	routers := NewRouters()
+	require.NoError(t, routers.Add("example.com", "/api", "alice", "exact-user"))
+	require.NoError(t, routers.Add("example.com", "/api", "", "exact-all-users"))
+	require.NoError(t, routers.Add("exact.example.com", "/api", "", "exact-subdomain"))
+	require.NoError(t, routers.Add("*.example.com", "/api", "", "wildcard-all-users"))
+	require.NoError(t, routers.Add("*.foo.example.com", "/api", "", "specific-wildcard"))
+	require.NoError(t, routers.Add("*.bar.com", "/api", "", "wildcard-parent-domain"))
+	require.NoError(t, routers.Add("*", "/admin", "root", "all-domains-user"))
+	require.NoError(t, routers.Add("*", "/", "", "all-domains"))
+
+	tests := []struct {
+		name     string
+		domain   string
+		location string
+		httpUser string
+		want     string
+	}{
+		{
+			name:     "exact domain and http user",
+			domain:   "example.com",
+			location: "/api/users",
+			httpUser: "alice",
+			want:     "exact-user",
+		},
+		{
+			name:     "exact domain falls back to all users",
+			domain:   "example.com",
+			location: "/api/users",
+			httpUser: "bob",
+			want:     "exact-all-users",
+		},
+		{
+			name:     "wildcard domain uses all users fallback",
+			domain:   "foo.example.com",
+			location: "/api/users",
+			httpUser: "bob",
+			want:     "wildcard-all-users",
+		},
+		{
+			name:     "mixed-case domain is normalized",
+			domain:   "Foo.Example.Com",
+			location: "/api/users",
+			httpUser: "bob",
+			want:     "wildcard-all-users",
+		},
+		{
+			name:     "exact domain wins over wildcard domain",
+			domain:   "exact.example.com",
+			location: "/api/users",
+			httpUser: "bob",
+			want:     "exact-subdomain",
+		},
+		{
+			name:     "more specific wildcard wins over broader wildcard",
+			domain:   "bar.foo.example.com",
+			location: "/api/users",
+			httpUser: "bob",
+			want:     "specific-wildcard",
+		},
+		{
+			name:     "wildcard walk checks parent domains",
+			domain:   "a.b.bar.com",
+			location: "/api/users",
+			httpUser: "bob",
+			want:     "wildcard-parent-domain",
+		},
+		{
+			name:     "catch-all domain fallback",
+			domain:   "foo.test.com",
+			location: "/other",
+			httpUser: "bob",
+			want:     "all-domains",
+		},
+		{
+			name:     "catch-all domain honors http user",
+			domain:   "foo.test.com",
+			location: "/admin/panel",
+			httpUser: "root",
+			want:     "all-domains-user",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			router, ok := routers.getByRoute(tt.domain, tt.location, tt.httpUser)
+			require.True(t, ok)
+			require.Equal(t, tt.want, router.payload)
+		})
+	}
+}
+
+func TestRoutersGetByRouteNoMatch(t *testing.T) {
+	routers := NewRouters()
+	require.NoError(t, routers.Add("*.example.com", "/api", "", "wildcard-all-users"))
+	require.NoError(t, routers.Add("*.com", "/api", "", "top-level-wildcard"))
+
+	tests := []struct {
+		name     string
+		domain   string
+		location string
+	}{
+		{
+			name:     "two-label domain does not enter wildcard walk",
+			domain:   "example.com",
+			location: "/api/users",
+		},
+		{
+			name:     "missing catch-all remains no match",
+			domain:   "foo.test.com",
+			location: "/api/users",
+		},
+		{
+			name:     "wrong path remains no match",
+			domain:   "foo.example.com",
+			location: "/other",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			router, ok := routers.getByRoute(tt.domain, tt.location, "bob")
+			require.False(t, ok)
+			require.Nil(t, router)
+		})
+	}
+}
+
+func TestRoutersConcurrentGetByRouteAndAdd(t *testing.T) {
+	routers := NewRouters()
+	require.NoError(t, routers.Add("*.example.com", "/api", "", "wildcard"))
+
+	const readers = 8
+	const iterations = 200
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	start := make(chan struct{})
+	errCh := make(chan error, readers+1)
+	var wg sync.WaitGroup
+
+	for id := range readers {
+		wg.Go(func() {
+			<-start
+			for j := range iterations {
+				select {
+				case <-ctx.Done():
+					return
+				default:
+				}
+				router, ok := routers.getByRoute("foo.example.com", "/api/users", "")
+				if !ok || router == nil || router.payload != "wildcard" {
+					errCh <- fmt.Errorf("reader %d iteration %d got router=%v ok=%v", id, j, router, ok)
+					return
+				}
+			}
+		})
+	}
+
+	wg.Go(func() {
+		<-start
+		for i := range iterations {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
+			err := routers.Add(fmt.Sprintf("host-%d.example.com", i), "/api", "", i)
+			if err != nil {
+				errCh <- err
+				return
+			}
+		}
+	})
+
+	close(start)
+
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		cancel()
+		t.Fatal("concurrent route lookup and add timed out")
+	}
+
+	close(errCh)
+	for err := range errCh {
+		require.NoError(t, err)
+	}
+}

pkg/util/vhost/vhost.go

@@ -148,41 +148,9 @@ func (v *Muxer) Listen(ctx context.Context, cfg *RouteConfig) (l *Listener, err
 }
 
 func (v *Muxer) getListener(name, path, httpUser string) (*Listener, bool) {
-	findRouter := func(inName, inPath, inHTTPUser string) (*Listener, bool) {
-		vr, ok := v.registryRouter.Get(inName, inPath, inHTTPUser)
-		if ok {
-			return vr.payload.(*Listener), true
-		}
-		// Try to check if there is one proxy that doesn't specify routerByHTTPUser, it means match all.
-		vr, ok = v.registryRouter.Get(inName, inPath, "")
-		if ok {
-			return vr.payload.(*Listener), true
-		}
-		return nil, false
-	}
-
-	// first we check the full hostname
-	// if not exist, then check the wildcard_domain such as *.example.com
-	l, ok := findRouter(name, path, httpUser)
+	vr, ok := v.registryRouter.getByRoute(name, path, httpUser)
 	if ok {
-		return l, true
-	}
-
-	domainSplit := strings.Split(name, ".")
-	for len(domainSplit) >= 3 {
-		domainSplit[0] = "*"
-		name = strings.Join(domainSplit, ".")
-
-		l, ok = findRouter(name, path, httpUser)
-		if ok {
-			return l, true
-		}
-		domainSplit = domainSplit[1:]
-	}
-	// Finally, try to check if there is one proxy that domain is "*" means match all domains.
-	l, ok = findRouter("*", path, httpUser)
-	if ok {
-		return l, true
+		return vr.payload.(*Listener), true
 	}
 	return nil, false
 }