Skip to content

Commit 158723a

Browse files
ethanndicksonethanndickson
committed
fix: include custom agent headers in tailnet
1 parent d2c1562 commit 158723a

File tree

2 files changed

+74
-17
lines changed

2 files changed

+74
-17
lines changed

‎agent/agent.go

+8
Original file line numberDiff line numberDiff line change
@@ -1134,11 +1134,19 @@ func (a *agent) trackGoroutine(fn func()) error {
11341134
}
11351135

11361136
func (a *agent) createTailnet(ctx context.Context, agentID uuid.UUID, derpMap *tailcfg.DERPMap, derpForceWebSockets, disableDirectConnections bool) (_ *tailnet.Conn, err error) {
1137+
// Inject `CODER_AGENT_HEADER` into the DERP header.
1138+
var header http.Header
1139+
if client, ok := a.client.(*agentsdk.Client); ok {
1140+
if headerTransport, ok := client.SDK.HTTPClient.Transport.(*codersdk.HeaderTransport); ok {
1141+
header = headerTransport.Header
1142+
}
1143+
}
11371144
network, err := tailnet.NewConn(&tailnet.Options{
11381145
ID: agentID,
11391146
Addresses: a.wireguardAddresses(agentID),
11401147
DERPMap: derpMap,
11411148
DERPForceWebSockets: derpForceWebSockets,
1149+
DERPHeader: &header,
11421150
Logger: a.logger.Named("net.tailnet"),
11431151
ListenPort: a.tailnetListenPort,
11441152
BlockEndpoints: disableDirectConnections,

‎cli/agent_test.go

+66-17
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"fmt"
66
"net/http"
7-
"net/http/httptest"
87
"os"
98
"path/filepath"
109
"runtime"
@@ -18,6 +17,7 @@ import (
1817

1918
"github.com/coder/coder/v2/agent"
2019
"github.com/coder/coder/v2/cli/clitest"
20+
"github.com/coder/coder/v2/coderd"
2121
"github.com/coder/coder/v2/coderd/coderdtest"
2222
"github.com/coder/coder/v2/coderd/database"
2323
"github.com/coder/coder/v2/coderd/database/dbfake"
@@ -232,39 +232,88 @@ func TestWorkspaceAgent(t *testing.T) {
232232
require.Equal(t, codersdk.AgentSubsystemEnvbox, resources[0].Agents[0].Subsystems[0])
233233
require.Equal(t, codersdk.AgentSubsystemExectrace, resources[0].Agents[0].Subsystems[1])
234234
})
235-
t.Run("Header", func(t *testing.T) {
235+
t.Run("Headers&DERPHeaders", func(t *testing.T) {
236236
t.Parallel()
237237

238-
var url string
238+
// Create a coderd API instance the hard way since we need to change the
239+
// handler to inject our custom /derp handler.
240+
dv := coderdtest.DeploymentValues(t)
241+
dv.DERP.Config.BlockDirect = true
242+
setHandler, cancelFunc, serverURL, newOptions := coderdtest.NewOptions(t, &coderdtest.Options{
243+
DeploymentValues: dv,
244+
})
245+
246+
// We set the handler after server creation for the access URL.
247+
coderAPI := coderd.New(newOptions)
248+
setHandler(coderAPI.RootHandler)
249+
provisionerCloser := coderdtest.NewProvisionerDaemon(t, coderAPI)
250+
t.Cleanup(func() {
251+
_ = provisionerCloser.Close()
252+
})
253+
client := codersdk.New(serverURL)
254+
t.Cleanup(func() {
255+
cancelFunc()
256+
_ = provisionerCloser.Close()
257+
_ = coderAPI.Close()
258+
client.HTTPClient.CloseIdleConnections()
259+
})
260+
261+
var (
262+
admin = coderdtest.CreateFirstUser(t, client)
263+
member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID)
264+
)
265+
239266
var called int64
240-
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
241-
assert.Equal(t, "wow", r.Header.Get("X-Testing"))
242-
assert.Equal(t, "Ethan was Here!", r.Header.Get("Cool-Header"))
243-
assert.Equal(t, "very-wow-"+url, r.Header.Get("X-Process-Testing"))
244-
assert.Equal(t, "more-wow", r.Header.Get("X-Process-Testing2"))
245-
atomic.AddInt64(&called, 1)
246-
w.WriteHeader(http.StatusGone)
267+
var derpCalled int64
268+
setHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
269+
// Ignore client requests
270+
if r.Header.Get("X-Testing") == "agent" {
271+
assert.Equal(t, "Ethan was Here!", r.Header.Get("Cool-Header"))
272+
assert.Equal(t, "very-wow-"+client.URL.String(), r.Header.Get("X-Process-Testing"))
273+
assert.Equal(t, "more-wow", r.Header.Get("X-Process-Testing2"))
274+
if strings.HasPrefix(r.URL.Path, "/derp") {
275+
atomic.AddInt64(&derpCalled, 1)
276+
} else {
277+
atomic.AddInt64(&called, 1)
278+
}
279+
}
280+
coderAPI.RootHandler.ServeHTTP(w, r)
247281
}))
248-
defer srv.Close()
249-
url = srv.URL
282+
r := dbfake.WorkspaceBuild(t, coderAPI.Database, database.Workspace{
283+
OrganizationID: memberUser.OrganizationIDs[0],
284+
OwnerID: memberUser.ID,
285+
}).WithAgent().Do()
286+
250287
coderURLEnv := "$CODER_URL"
251288
if runtime.GOOS == "windows" {
252289
coderURLEnv = "%CODER_URL%"
253290
}
254291

255292
logDir := t.TempDir()
256-
inv, _ := clitest.New(t,
293+
agentInv, _ := clitest.New(t,
257294
"agent",
258295
"--auth", "token",
259-
"--agent-token", "fake-token",
260-
"--agent-url", srv.URL,
296+
"--agent-token", r.AgentToken,
297+
"--agent-url", client.URL.String(),
261298
"--log-dir", logDir,
262-
"--agent-header", "X-Testing=wow",
299+
"--agent-header", "X-Testing=agent",
263300
"--agent-header", "Cool-Header=Ethan was Here!",
264301
"--agent-header-command", "printf X-Process-Testing=very-wow-"+coderURLEnv+"'\\r\\n'X-Process-Testing2=more-wow",
265302
)
303+
clitest.Start(t, agentInv)
304+
coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
305+
MatchResources(matchAgentWithVersion).Wait()
306+
307+
clientInv, root := clitest.New(t,
308+
"-v",
309+
"--no-feature-warning",
310+
"--no-version-warning",
311+
"ping", r.Workspace.Name,
312+
"-n", "1",
313+
)
314+
clitest.SetupConfig(t, member, root)
315+
clitest.Start(t, clientInv)
266316

267-
clitest.Start(t, inv)
268317
require.Eventually(t, func() bool {
269318
return atomic.LoadInt64(&called) > 0
270319
}, testutil.WaitShort, testutil.IntervalFast)

0 commit comments

Comments
 (0)