fix: rewrite url to agent ip in single tailnet

coadler · coadler · commit b7ba1c964360 · 2024-01-24T19:50:11.000Z
This restores previous behavior of being able to cache connections
across agents in single tailnet.
diff --git a/coderd/tailnet.go b/coderd/tailnet.go
@@ -103,14 +103,7 @@ func NewServerTailnet(
 		transport:            tailnetTransport.Clone(),
 	}
 	tn.transport.DialContext = tn.dialContext
-
-	// Bugfix: for some reason all calls to tn.dialContext come from
-	// "localhost", causing connections to be cached and requests to go to the
-	// wrong workspaces. This disables keepalives for now until the root cause
-	// can be found.
-	tn.transport.MaxIdleConnsPerHost = -1
-	tn.transport.DisableKeepAlives = true
-
+	tn.transport.MaxIdleConnsPerHost = 10
 	tn.transport.MaxIdleConns = 0
 	// We intentionally don't verify the certificate chain here.
 	// The connection to the workspace is already established and most
@@ -312,7 +305,15 @@ type ServerTailnet struct {
 }
 
 func (s *ServerTailnet) ReverseProxy(targetURL, dashboardURL *url.URL, agentID uuid.UUID) (_ *httputil.ReverseProxy, release func(), _ error) {
-	proxy := httputil.NewSingleHostReverseProxy(targetURL)
+	// Rewrite the targetURL's Host to point to the agent's IP. This is
+	// necessary because due to TCP connection caching, each agent needs to be
+	// addressed invidivually. Otherwise, all connections get dialed as
+	// "localhost:port", causing connections to be shared across agents.
+	tgt := *targetURL
+	_, port, _ := net.SplitHostPort(tgt.Host)
+	tgt.Host = net.JoinHostPort(tailnet.IPFromUUID(agentID).String(), port)
+
+	proxy := httputil.NewSingleHostReverseProxy(&tgt)
 	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
 		site.RenderStaticErrorPage(w, r, site.ErrorPageData{
 			Status:       http.StatusBadGateway,
@@ -460,7 +461,7 @@ func (s *ServerTailnet) Close() error {
 	s.cancel()
 	_ = s.cache.Close()
 	_ = s.conn.Close()
-	s.transport.CloseIdleConnections()
+	// s.transport.CloseIdleConnections()
 	<-s.derpMapUpdaterClosed
 	return nil
 }
diff --git a/coderd/tailnet_test.go b/coderd/tailnet_test.go
@@ -95,6 +95,32 @@ func TestServerTailnet_ReverseProxy(t *testing.T) {
 		assert.Equal(t, http.StatusOK, res.StatusCode)
 	})
 
+	t.Run("HostRewrite", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentID, _, serverTailnet := setupAgent(t, nil)
+
+		u, err := url.Parse(fmt.Sprintf("http://127.0.0.1:%d", codersdk.WorkspaceAgentHTTPAPIServerPort))
+		require.NoError(t, err)
+
+		rp, release, err := serverTailnet.ReverseProxy(u, u, agentID)
+		require.NoError(t, err)
+		defer release()
+
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
+		require.NoError(t, err)
+
+		// Ensure the reverse proxy director rewrites the url host to the agent's IP.
+		rp.Director(req)
+		assert.Equal(t,
+			fmt.Sprintf("[%s]:%d", tailnet.IPFromUUID(agentID).String(), codersdk.WorkspaceAgentHTTPAPIServerPort),
+			req.URL.Host,
+		)
+	})
+
 	t.Run("HTTPSProxy", func(t *testing.T) {
 		t.Parallel()
 
