From 80456a3bacffca3d3a16e2c4a321564436692c53 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Mon, 21 Oct 2024 01:37:05 +0000 Subject: [PATCH 1/2] fix: include custom agent headers in tailnet --- agent/agent.go | 8 +++++ cli/agent_test.go | 83 +++++++++++++++++++++++++++++++++++++---------- 2 files changed, 74 insertions(+), 17 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index 6e4c3fa476bb3..cb0037dd0ed48 100644 --- 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 index f30d12b012d88..f3ef06edd52fa 100644 --- 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) From eb5a0d910cccd46ce564db4708d98827eca0362a Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Mon, 21 Oct 2024 02:50:17 +0000 Subject: [PATCH 2/2] attempt ci fix --- cli/agent_test.go | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/cli/agent_test.go b/cli/agent_test.go index f3ef06edd52fa..5d37f2c54c9d4 100644 --- a/cli/agent_test.go +++ b/cli/agent_test.go @@ -261,10 +261,10 @@ func TestWorkspaceAgent(t *testing.T) { var ( admin = coderdtest.CreateFirstUser(t, client) member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID) + called int64 + derpCalled int64 ) - var called int64 - var derpCalled int64 setHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Ignore client requests if r.Header.Get("X-Testing") == "agent" { @@ -304,6 +304,7 @@ func TestWorkspaceAgent(t *testing.T) { coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID). MatchResources(matchAgentWithVersion).Wait() + ctx := testutil.Context(t, testutil.WaitLong) clientInv, root := clitest.New(t, "-v", "--no-feature-warning", @@ -312,11 +313,11 @@ func TestWorkspaceAgent(t *testing.T) { "-n", "1", ) clitest.SetupConfig(t, member, root) - clitest.Start(t, clientInv) + err := clientInv.WithContext(ctx).Run() + require.NoError(t, err) - require.Eventually(t, func() bool { - return atomic.LoadInt64(&called) > 0 - }, testutil.WaitShort, testutil.IntervalFast) + require.Greater(t, atomic.LoadInt64(&called), int64(0), "expected coderd to be reached with custom headers") + require.Greater(t, atomic.LoadInt64(&derpCalled), int64(0), "expected /derp to be called with custom headers") }) }