coder · kylecarbs · Nov 4, 2022 · Oct 28, 2022 · Nov 1, 2022 · Nov 1, 2022
diff --git a/agent/agent.go b/agent/agent.go
@@ -221,12 +221,18 @@ func (a *agent) run(ctx context.Context) error {
 }
 
 func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (*tailnet.Conn, error) {
+	a.closeMutex.Lock()
+	if a.isClosed() {
+		a.closeMutex.Unlock()
+		return nil, xerrors.New("closed")
+	}
 	network, err := tailnet.NewConn(&tailnet.Options{
 		Addresses: []netip.Prefix{netip.PrefixFrom(codersdk.TailnetIP, 128)},
 		DERPMap:   derpMap,
 		Logger:    a.logger.Named("tailnet"),
 	})
 	if err != nil {
+		a.closeMutex.Unlock()
 		return nil, xerrors.Errorf("create tailnet: %w", err)
 	}
 	a.network = network
@@ -237,14 +243,13 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (*t
 		}
 		return a.stats.wrapConn(conn)
 	})
+	a.connCloseWait.Add(4)
+	a.closeMutex.Unlock()
 
 	sshListener, err := network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetSSHPort))
 	if err != nil {
 		return nil, xerrors.Errorf("listen on the ssh port: %w", err)
 	}
-	a.closeMutex.Lock()
-	a.connCloseWait.Add(1)
-	a.closeMutex.Unlock()
 	go func() {
 		defer a.connCloseWait.Done()
 		for {
@@ -260,9 +265,6 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (*t
 	if err != nil {
 		return nil, xerrors.Errorf("listen for reconnecting pty: %w", err)
 	}
-	a.closeMutex.Lock()
-	a.connCloseWait.Add(1)
-	a.closeMutex.Unlock()
 	go func() {
 		defer a.connCloseWait.Done()
 		for {
@@ -298,9 +300,6 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (*t
 	if err != nil {
 		return nil, xerrors.Errorf("listen for speedtest: %w", err)
 	}
-	a.closeMutex.Lock()
-	a.connCloseWait.Add(1)
-	a.closeMutex.Unlock()
 	go func() {
 		defer a.connCloseWait.Done()
 		for {
@@ -323,9 +322,6 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (*t
 	if err != nil {
 		return nil, xerrors.Errorf("listen for statistics: %w", err)
 	}
-	a.closeMutex.Lock()
-	a.connCloseWait.Add(1)
-	a.closeMutex.Unlock()
 	go func() {
 		defer a.connCloseWait.Done()
 		defer statisticsListener.Close()
@@ -569,7 +565,6 @@ func (a *agent) createCommand(ctx context.Context, rawCommand string, env []stri
 	// Set environment variables reliable detection of being inside a
 	// Coder workspace.
 	cmd.Env = append(cmd.Env, "CODER=true")
-
 	cmd.Env = append(cmd.Env, fmt.Sprintf("USER=%s", username))
 	// Git on Windows resolves with UNIX-style paths.
 	// If using backslashes, it's unable to find the executable.
@@ -585,6 +580,10 @@ func (a *agent) createCommand(ctx context.Context, rawCommand string, env []stri
 	cmd.Env = append(cmd.Env, fmt.Sprintf("SSH_CLIENT=%s %s %s", srcAddr, srcPort, dstPort))
 	cmd.Env = append(cmd.Env, fmt.Sprintf("SSH_CONNECTION=%s %s %s %s", srcAddr, srcPort, dstAddr, dstPort))
 
+	// This adds the ports dialog to code-server that enables
+	// proxying a port dynamically.
+	cmd.Env = append(cmd.Env, fmt.Sprintf("VSCODE_PROXY_URI=%s", metadata.VSCodePortProxyURI))
+
 	// Hide Coder message on code-server's "Getting Started" page
 	cmd.Env = append(cmd.Env, "CS_DISABLE_GETTING_STARTED_OVERRIDE=true")
 

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -23,6 +23,7 @@ import (
 
 	"golang.org/x/xerrors"
 	"tailscale.com/net/speedtest"
+	"tailscale.com/tailcfg"
 
 	scp "github.com/bramvdbogaerde/go-scp"
 	"github.com/google/uuid"
@@ -559,6 +560,7 @@ func TestAgent(t *testing.T) {
 			agentID: uuid.New(),
 			metadata: codersdk.WorkspaceAgentMetadata{
 				GitAuthConfigs: 1,
+				DERPMap:        &tailcfg.DERPMap{},
 			},
 			statsChan:   make(chan *codersdk.AgentStats),
 			coordinator: tailnet.NewCoordinator(),

diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go
@@ -31,7 +31,7 @@ func TestWorkspaceActivityBump(t *testing.T) {
 		})
 		user := coderdtest.CreateFirstUser(t, client)
 
-		workspace = createWorkspaceWithApps(t, client, user.OrganizationID, 1234, func(cwr *codersdk.CreateWorkspaceRequest) {
+		workspace = createWorkspaceWithApps(t, client, user.OrganizationID, "", 1234, func(cwr *codersdk.CreateWorkspaceRequest) {
 			cwr.TTLMillis = &ttlMillis
 		})
 

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -83,6 +83,46 @@ func (api *API) workspaceAgentMetadata(rw http.ResponseWriter, r *http.Request)
 		})
 		return
 	}
+	resource, err := api.Database.GetWorkspaceResourceByID(r.Context(), workspaceAgent.ResourceID)
+	if err != nil {
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching workspace resource.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	build, err := api.Database.GetWorkspaceBuildByJobID(r.Context(), resource.JobID)
+	if err != nil {
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching workspace build.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	workspace, err := api.Database.GetWorkspaceByID(r.Context(), build.WorkspaceID)
+	if err != nil {
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching workspace.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	owner, err := api.Database.GetUserByID(r.Context(), workspace.OwnerID)
+	if err != nil {
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching workspace owner.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	vscodeProxyURI := strings.ReplaceAll(api.AppHostname, "*",
+		fmt.Sprintf("%s://{{port}}--%s--%s--%s",
+			api.AccessURL.Scheme,
+			workspaceAgent.Name,
+			workspace.Name,
+			owner.Username,
+		))
 
 	httpapi.Write(ctx, rw, http.StatusOK, codersdk.WorkspaceAgentMetadata{
 		Apps:                 convertApps(dbApps),
@@ -91,6 +131,7 @@ func (api *API) workspaceAgentMetadata(rw http.ResponseWriter, r *http.Request)
 		EnvironmentVariables: apiAgent.EnvironmentVariables,
 		StartupScript:        apiAgent.StartupScript,
 		Directory:            apiAgent.Directory,
+		VSCodePortProxyURI:   vscodeProxyURI,
 	})
 }
 

diff --git a/coderd/workspaceapps_test.go b/coderd/workspaceapps_test.go
@@ -121,7 +121,7 @@ func setupProxyTest(t *testing.T, customAppHost ...string) (*codersdk.Client, co
 	})
 	user := coderdtest.CreateFirstUser(t, client)
 
-	workspace := createWorkspaceWithApps(t, client, user.OrganizationID, uint16(tcpAddr.Port))
+	workspace := createWorkspaceWithApps(t, client, user.OrganizationID, appHost, uint16(tcpAddr.Port))
 
 	// Configure the HTTP client to not follow redirects and to route all
 	// requests regardless of hostname to the coderd test server.
@@ -139,7 +139,7 @@ func setupProxyTest(t *testing.T, customAppHost ...string) (*codersdk.Client, co
 	return client, user, workspace, uint16(tcpAddr.Port)
 }
 
-func createWorkspaceWithApps(t *testing.T, client *codersdk.Client, orgID uuid.UUID, port uint16, workspaceMutators ...func(*codersdk.CreateWorkspaceRequest)) codersdk.Workspace {
+func createWorkspaceWithApps(t *testing.T, client *codersdk.Client, orgID uuid.UUID, appHost string, port uint16, workspaceMutators ...func(*codersdk.CreateWorkspaceRequest)) codersdk.Workspace {
 	authToken := uuid.NewString()
 
 	appURL := fmt.Sprintf("http://127.0.0.1:%d?%s", port, proxyTestAppQuery)
@@ -198,6 +198,17 @@ func createWorkspaceWithApps(t *testing.T, client *codersdk.Client, orgID uuid.U
 
 	agentClient := codersdk.New(client.URL)
 	agentClient.SessionToken = authToken
+	if appHost != "" {
+		metadata, err := agentClient.WorkspaceAgentMetadata(context.Background())
+		require.NoError(t, err)
+		require.Equal(t, fmt.Sprintf(
+			"http://{{port}}--%s--%s--%s%s",
+			proxyTestAgentName,
+			workspace.Name,
+			"testuser",
+			strings.ReplaceAll(appHost, "*", ""),
+		), metadata.VSCodePortProxyURI)
+	}
 	agentCloser := agent.New(agent.Options{
 		Client: agentClient,
 		Logger: slogtest.Make(t, nil).Named("agent"),

diff --git a/coderd/wsconncache/wsconncache.go b/coderd/wsconncache/wsconncache.go
@@ -86,8 +86,18 @@ func (c *Cache) Acquire(r *http.Request, id uuid.UUID) (*Conn, func(), error) {
 		// A singleflight group is used to allow for concurrent requests to the
 		// same identifier to resolve.
 		rawConn, err, _ = c.connGroup.Do(id.String(), func() (interface{}, error) {
+			c.closeMutex.Lock()
+			select {
+			case <-c.closed:
+				c.closeMutex.Unlock()
+				return nil, xerrors.New("closed")
+			default:
+			}
+			c.closeGroup.Add(1)
+			c.closeMutex.Unlock()
 			agentConn, err := c.dialer(r, id)
 			if err != nil {
+				c.closeGroup.Done()
 				return nil, xerrors.Errorf("dial: %w", err)
 			}
 			timeoutCtx, timeoutCancelFunc := context.WithCancel(context.Background())
@@ -102,9 +112,6 @@ func (c *Cache) Acquire(r *http.Request, id uuid.UUID) (*Conn, func(), error) {
 				timeoutCancel: timeoutCancelFunc,
 				transport:     transport,
 			}
-			c.closeMutex.Lock()
-			c.closeGroup.Add(1)
-			c.closeMutex.Unlock()
 			go func() {
 				defer c.closeGroup.Done()
 				var err error

diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -123,6 +123,7 @@ type WorkspaceAgentMetadata struct {
 	// the Coder deployment has. If this number is >0, we
 	// set up special configuration in the workspace.
 	GitAuthConfigs       int               `json:"git_auth_configs"`
+	VSCodePortProxyURI   string            `json:"vscode_port_proxy_uri"`
 	Apps                 []WorkspaceApp    `json:"apps"`
 	DERPMap              *tailcfg.DERPMap  `json:"derpmap"`
 	EnvironmentVariables map[string]string `json:"environment_variables"`