Function TestDERPHeaders

cli/root_test.go:202–293

TestDERPHeaders ensures that the client sends the global `--header`s and `--header-command` to the DERP server when connecting.

(t *testing.T)

// TestDERPHeaders ensures that the client sends the global `--header`s and
// `--header-command` to the DERP server when connecting.
func TestDERPHeaders(t *testing.T) {
	t.Parallel()

	// 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, codersdk.WithHTTPClient(coderdtest.NewIsolatedHTTPClient(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)
		workspace          = runAgent(t, client, memberUser.ID, newOptions.Database)
	)

	// Inject custom /derp handler so we can inspect the headers.
	var (
		expectedHeaders = map[string]string{
			"X-Test-Header":     "test-value",
			"Cool-Header":       "Dean was Here!",
			"X-Process-Testing": "very-wow",
		}
		derpCalled atomic.Int64
	)
	setHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/derp") {
			ok := true
			for k, v := range expectedHeaders {
				if r.Header.Get(k) != v {
					ok = false
					break
				}
			}
			if ok {
				// Only increment if all the headers are set, because the agent
				// calls derp also.
				derpCalled.Add(1)
			}
		}

		coderAPI.RootHandler.ServeHTTP(w, r)