From a8eeddfb42e71d7a98c90f30d8814bf732c25442 Mon Sep 17 00:00:00 2001
From: Tao Yang <swulling@gmail.com>
Date: Tue, 15 Nov 2022 14:38:55 +0800
Subject: [PATCH] feat: Let port-forwarding support custom http(s) port

---
 coderd/workspaceagents.go    |  3 +++
 coderd/workspaceapps.go      |  7 ++++++-
 coderd/workspaceapps_test.go | 11 ++++++++---
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
index bf3732ee52849..a92a9950d4d62 100644
--- a/coderd/workspaceagents.go
+++ b/coderd/workspaceagents.go
@@ -123,6 +123,9 @@ func (api *API) workspaceAgentMetadata(rw http.ResponseWriter, r *http.Request)
 			workspace.Name,
 			owner.Username,
 		))
+	if api.AccessURL.Port() != "" {
+		vscodeProxyURI += fmt.Sprintf(":%s", api.AccessURL.Port())
+	}
 
 	httpapi.Write(ctx, rw, http.StatusOK, codersdk.WorkspaceAgentMetadata{
 		Apps:                 convertApps(dbApps),
diff --git a/coderd/workspaceapps.go b/coderd/workspaceapps.go
index 95734387abeb5..910ddb1637a7f 100644
--- a/coderd/workspaceapps.go
+++ b/coderd/workspaceapps.go
@@ -40,8 +40,13 @@ const (
 )
 
 func (api *API) appHost(rw http.ResponseWriter, r *http.Request) {
+	host := api.AppHostname
+	if api.AccessURL.Port() != "" {
+		host += fmt.Sprintf(":%s", api.AccessURL.Port())
+	}
+
 	httpapi.Write(r.Context(), rw, http.StatusOK, codersdk.GetAppHostResponse{
-		Host: api.AppHostname,
+		Host: host,
 	})
 }
 
diff --git a/coderd/workspaceapps_test.go b/coderd/workspaceapps_test.go
index 6929f97025d89..15a1ebe145969 100644
--- a/coderd/workspaceapps_test.go
+++ b/coderd/workspaceapps_test.go
@@ -69,7 +69,8 @@ func TestGetAppHost(t *testing.T) {
 			_ = coderdtest.CreateFirstUser(t, client)
 			host, err = client.GetAppHost(ctx)
 			require.NoError(t, err)
-			require.Equal(t, c, host.Host)
+			domain := strings.Split(host.Host, ":")[0]
+			require.Equal(t, c, domain)
 		})
 	}
 }
@@ -204,13 +205,17 @@ func createWorkspaceWithApps(t *testing.T, client *codersdk.Client, orgID uuid.U
 	if appHost != "" {
 		metadata, err := agentClient.WorkspaceAgentMetadata(context.Background())
 		require.NoError(t, err)
-		require.Equal(t, fmt.Sprintf(
+		proxyURL := fmt.Sprintf(
 			"http://{{port}}--%s--%s--%s%s",
 			proxyTestAgentName,
 			workspace.Name,
 			"testuser",
 			strings.ReplaceAll(appHost, "*", ""),
-		), metadata.VSCodePortProxyURI)
+		)
+		if client.URL.Port() != "" {
+			proxyURL += fmt.Sprintf(":%s", client.URL.Port())
+		}
+		require.Equal(t, proxyURL, metadata.VSCodePortProxyURI)
 	}
 	agentCloser := agent.New(agent.Options{
 		Client: agentClient,