connect: router fixes

Author: xcolwell

connect/main.go

@@ -2,7 +2,7 @@ package main
 
 import (
 	"context"
-	// "fmt"
+	"fmt"
 	"net"
 	"net/http"
 	"os"
@@ -92,17 +92,21 @@ Options:
 		}
 	})
 
+	service := strings.ToLower(server.RequireService())
+	envService := strings.ToLower(fmt.Sprintf("%s-%s", server.RequireEnv(), server.RequireService()))
+
 	connectRouter := func(w http.ResponseWriter, r *http.Request) {
 		host := r.Header.Get("X-Forwarded-Host")
 		if host == "" {
 			host = r.Header.Get("Host")
 		}
 
-		sub := strings.SplitAfterN(host, ".", 2)[0]
-		if sub == "connect" {
-			// the host is connect.<domain>
+		sub := strings.ToLower(strings.SplitN(host, ".", 2)[0])
+		switch sub {
+		case service, envService:
+			// the host is connect.<domain> or <env>-connect.<domain>
 			connectHandler.Connect(w, r)
-		} else {
+		default:
 			// the host is <auth>.connect.<domain>
 			proxyConnectHandler.Connect(w, r)
 		}

connect/transport_proxy.go

@@ -185,7 +185,7 @@ func (self *ProxyConnectHandler) Connect(w http.ResponseWriter, r *http.Request)
 		host = r.Header.Get("Host")
 	}
 
-	hostProxyId := strings.SplitAfterN(host, ".", 2)[0]
+	hostProxyId := strings.SplitN(host, ".", 2)[0]
 	proxyId, err := model.ParseSignedProxyId(hostProxyId)
 	if err != nil {
 		proxyId, err = model.ParseEncodedProxyId(hostProxyId)
@@ -213,7 +213,7 @@ func (self *ProxyConnectHandler) Connect(w http.ResponseWriter, r *http.Request)
 					http.Error(w, "Malformed auth", http.StatusInternalServerError)
 					return
 				}
-				signedProxyId := strings.SplitAfterN(string(combinedSignedProxyId), ":", 2)[0]
+				signedProxyId := strings.SplitN(string(combinedSignedProxyId), ":", 2)[0]
 				headerProxyId, err := model.ParseSignedProxyId(signedProxyId)
 				if err != nil || proxyId != headerProxyId {
 					http.Error(w, "Not authorized", http.StatusUnauthorized)