fix: include custom agent headers in tailnet

ethanndickson · ethanndickson · commit 80456a3bacff · 2024-10-21T04:19:00.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -1134,11 +1134,19 @@ func (a *agent) trackGoroutine(fn func()) error {
 }
 
 func (a *agent) createTailnet(ctx context.Context, agentID uuid.UUID, derpMap *tailcfg.DERPMap, derpForceWebSockets, disableDirectConnections bool) (_ *tailnet.Conn, err error) {
+	// Inject `CODER_AGENT_HEADER` into the DERP header.
+	var header http.Header
+	if client, ok := a.client.(*agentsdk.Client); ok {
+		if headerTransport, ok := client.SDK.HTTPClient.Transport.(*codersdk.HeaderTransport); ok {
+			header = headerTransport.Header
+		}
+	}
 	network, err := tailnet.NewConn(&tailnet.Options{
 		ID:                  agentID,
 		Addresses:           a.wireguardAddresses(agentID),
 		DERPMap:             derpMap,
 		DERPForceWebSockets: derpForceWebSockets,
+		DERPHeader:          &header,
 		Logger:              a.logger.Named("net.tailnet"),
 		ListenPort:          a.tailnetListenPort,
 		BlockEndpoints:      disableDirectConnections,
diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"net/http"
-	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -18,6 +17,7 @@ import (
 
 	"github.com/coder/coder/v2/agent"
 	"github.com/coder/coder/v2/cli/clitest"
+	"github.com/coder/coder/v2/coderd"
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
@@ -232,39 +232,88 @@ func TestWorkspaceAgent(t *testing.T) {
 		require.Equal(t, codersdk.AgentSubsystemEnvbox, resources[0].Agents[0].Subsystems[0])
 		require.Equal(t, codersdk.AgentSubsystemExectrace, resources[0].Agents[0].Subsystems[1])
 	})
-	t.Run("Header", func(t *testing.T) {
+	t.Run("Headers&DERPHeaders", func(t *testing.T) {
 		t.Parallel()
 
-		var url string
+		// Create a coderd API instance the hard way since we need to change the
+		// handler to inject our custom /derp handler.
+		dv := coderdtest.DeploymentValues(t)
+		dv.DERP.Config.BlockDirect = true
+		setHandler, cancelFunc, serverURL, newOptions := coderdtest.NewOptions(t, &coderdtest.Options{
+			DeploymentValues: dv,
+		})
+
+		// We set the handler after server creation for the access URL.
+		coderAPI := coderd.New(newOptions)
+		setHandler(coderAPI.RootHandler)
+		provisionerCloser := coderdtest.NewProvisionerDaemon(t, coderAPI)
+		t.Cleanup(func() {
+			_ = provisionerCloser.Close()
+		})
+		client := codersdk.New(serverURL)
+		t.Cleanup(func() {
+			cancelFunc()
+			_ = provisionerCloser.Close()
+			_ = coderAPI.Close()
+			client.HTTPClient.CloseIdleConnections()
+		})
+
+		var (
+			admin              = coderdtest.CreateFirstUser(t, client)
+			member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID)
+		)
+
 		var called int64
-		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			assert.Equal(t, "wow", r.Header.Get("X-Testing"))
-			assert.Equal(t, "Ethan was Here!", r.Header.Get("Cool-Header"))
-			assert.Equal(t, "very-wow-"+url, r.Header.Get("X-Process-Testing"))
-			assert.Equal(t, "more-wow", r.Header.Get("X-Process-Testing2"))
-			atomic.AddInt64(&called, 1)
-			w.WriteHeader(http.StatusGone)
+		var derpCalled int64
+		setHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// Ignore client requests
+			if r.Header.Get("X-Testing") == "agent" {
+				assert.Equal(t, "Ethan was Here!", r.Header.Get("Cool-Header"))
+				assert.Equal(t, "very-wow-"+client.URL.String(), r.Header.Get("X-Process-Testing"))
+				assert.Equal(t, "more-wow", r.Header.Get("X-Process-Testing2"))
+				if strings.HasPrefix(r.URL.Path, "/derp") {
+					atomic.AddInt64(&derpCalled, 1)
+				} else {
+					atomic.AddInt64(&called, 1)
+				}
+			}
+			coderAPI.RootHandler.ServeHTTP(w, r)
 		}))
-		defer srv.Close()
-		url = srv.URL
+		r := dbfake.WorkspaceBuild(t, coderAPI.Database, database.Workspace{
+			OrganizationID: memberUser.OrganizationIDs[0],
+			OwnerID:        memberUser.ID,
+		}).WithAgent().Do()
+
 		coderURLEnv := "$CODER_URL"
 		if runtime.GOOS == "windows" {
 			coderURLEnv = "%CODER_URL%"
 		}
 
 		logDir := t.TempDir()
-		inv, _ := clitest.New(t,
+		agentInv, _ := clitest.New(t,
 			"agent",
 			"--auth", "token",
-			"--agent-token", "fake-token",
-			"--agent-url", srv.URL,
+			"--agent-token", r.AgentToken,
+			"--agent-url", client.URL.String(),
 			"--log-dir", logDir,
-			"--agent-header", "X-Testing=wow",
+			"--agent-header", "X-Testing=agent",
 			"--agent-header", "Cool-Header=Ethan was Here!",
 			"--agent-header-command", "printf X-Process-Testing=very-wow-"+coderURLEnv+"'\\r\\n'X-Process-Testing2=more-wow",
 		)
+		clitest.Start(t, agentInv)
+		coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
+			MatchResources(matchAgentWithVersion).Wait()
+
+		clientInv, root := clitest.New(t,
+			"-v",
+			"--no-feature-warning",
+			"--no-version-warning",
+			"ping", r.Workspace.Name,
+			"-n", "1",
+		)
+		clitest.SetupConfig(t, member, root)
+		clitest.Start(t, clientInv)
 
-		clitest.Start(t, inv)
 		require.Eventually(t, func() bool {
 			return atomic.LoadInt64(&called) > 0
 		}, testutil.WaitShort, testutil.IntervalFast)
