feat: refresh header commands based on JWT expiration (#29760)

Adds opportunistic JWT parsing for values returned from the
`--header-command`. If a JWT is parsed we use the value of its `exp`
field to determine when we should re-run the command to refresh the
values.

agent/agent.go

@@ -1865,10 +1865,19 @@ func (a *agent) createTailnet(
 	keySeed int64,
 ) (_ *tailnet.Conn, err error) {
 	// Inject `CODER_AGENT_HEADER` into the DERP header.
-	var header http.Header
+	var getHeaders func() http.Header
 	if client, ok := a.client.(*agentsdk.Client); ok {
-		if headerTransport, ok := client.SDK.HTTPClient.Transport.(*codersdk.HeaderTransport); ok {
-			header = headerTransport.Header
+		if headerTransport, ok := client.SDK.HTTPClient.Transport.(*codersdk.HeaderTransport); ok && headerTransport.Provider != nil {
+			getHeaders = func() http.Header {
+				refreshCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+				defer cancel()
+				headers, err := headerTransport.Provider.Headers(refreshCtx)
+				if err != nil {
+					a.logger.Error(ctx, "get connection headers", slog.Error(err))
+					return nil
+				}
+				return headers
+			}
 		}
 	}
 	network, err := tailnet.NewConn(&tailnet.Options{
@@ -1876,7 +1885,7 @@ func (a *agent) createTailnet(
 		Addresses:           a.wireguardAddresses(agentID),
 		DERPMap:             derpMap,
 		DERPForceWebSockets: derpForceWebSockets,
-		DERPHeader:          &header,
+		DERPGetHeaders:      getHeaders,
 		DERPTLSConfig:       a.derpTLSConfig,
 		Logger:              a.logger.Named("net.tailnet"),
 		ListenPort:          a.tailnetListenPort,

cli/exp_scaletest.go

@@ -630,9 +630,9 @@ func (r *RootCmd) scaletestCleanup() *serpent.Command {
 			client.HTTPClient = &http.Client{
 				Transport: &codersdk.HeaderTransport{
 					Transport: http.DefaultTransport,
-					Header: map[string][]string{
+					Provider: codersdk.StaticHeaderProvider{Header: http.Header{
 						codersdk.BypassRatelimitHeader: {"true"},
-					},
+					}},
 				},
 			}
 

cli/exp_scaletest_bridge.go

@@ -67,9 +67,9 @@ Examples:
 			client.HTTPClient = &http.Client{
 				Transport: &codersdk.HeaderTransport{
 					Transport: http.DefaultTransport,
-					Header: map[string][]string{
+					Provider: codersdk.StaticHeaderProvider{Header: http.Header{
 						codersdk.BypassRatelimitHeader: {"true"},
-					},
+					}},
 				},
 			}
 			reg := prometheus.NewRegistry()

cli/exp_scaletest_chat.go

@@ -62,10 +62,6 @@ func (r *RootCmd) scaletestChat() *serpent.Command {
 			if err != nil {
 				return err
 			}
-			client.HTTPClient.Transport = &codersdk.HeaderTransport{
-				Transport: client.HTTPClient.Transport,
-				Header:    BypassHeader,
-			}
 
 			workspaces, err := targetFlags.getTargetedWorkspaces(ctx, client, me.OrganizationIDs, inv.Stdout)
 			if err != nil {

cli/header.go

@@ -0,0 +1,186 @@
+package cli
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"net/url"
+	"os"
+	"os/exec"
+	"runtime"
+	"strings"
+	"time"
+
+	"github.com/go-jose/go-jose/v4"
+	"github.com/go-jose/go-jose/v4/jwt"
+	"golang.org/x/sync/singleflight"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/quartz"
+)
+
+// jwtExpirationSkew accounts for clock skew between the client and issuer.
+const jwtExpirationSkew = 10 * time.Second
+
+// headerTransport resolves initial headers so command errors surface at startup.
+func headerTransport(ctx context.Context, serverURL *url.URL, header []string, headerCommand string) (*codersdk.HeaderTransport, error) {
+	headers, err := parseHeaders(header)
+	if err != nil {
+		return nil, err
+	}
+	var provider codersdk.HeaderProvider
+	if headerCommand == "" {
+		provider = codersdk.StaticHeaderProvider{Header: headers}
+	} else {
+		provider = &commandHeaderProvider{
+			serverURL:     serverURL,
+			staticHeaders: headers,
+			command:       headerCommand,
+			clock:         quartz.NewReal(),
+		}
+		if _, err := provider.Headers(ctx); err != nil {
+			return nil, err
+		}
+	}
+	return &codersdk.HeaderTransport{
+		Transport: http.DefaultTransport,
+		Provider:  provider,
+	}, nil
+}
+
+type commandHeaderProvider struct {
+	serverURL     *url.URL
+	staticHeaders http.Header
+	command       string
+	clock         quartz.Clock
+
+	sf singleflight.Group
+	// Cache access is serialized inside sf.Do. Published maps are immutable.
+	cached  http.Header
+	expires time.Time
+}
+
+func (p *commandHeaderProvider) Headers(ctx context.Context) (http.Header, error) {
+	value, err, _ := p.sf.Do("headers", func() (any, error) {
+		if p.cached != nil && (p.expires.IsZero() || p.clock.Now().Before(p.expires)) {
+			return p.cached.Clone(), nil
+		}
+		lines, err := runHeaderCommand(ctx, p.serverURL, p.command)
+		if err != nil {
+			return nil, err
+		}
+		headers, err := parseHeaders(lines)
+		if err != nil {
+			return nil, err
+		}
+		p.cached, p.expires = headers, headerExpiry(headers)
+		return headers.Clone(), nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	headers, ok := value.(http.Header)
+	if !ok {
+		return nil, xerrors.New("unexpected header cache result")
+	}
+	merged := p.staticHeaders.Clone()
+	if merged == nil {
+		merged = make(http.Header)
+	}
+	for name, values := range headers {
+		for _, value := range values {
+			merged.Add(name, value)
+		}
+	}
+	return merged, nil
+}
+
+func runHeaderCommand(ctx context.Context, serverURL *url.URL, command string) ([]string, error) {
+	shell, caller := "sh", "-c"
+	if runtime.GOOS == "windows" {
+		shell, caller = "cmd.exe", "/c"
+	}
+	var stdout bytes.Buffer
+	// #nosec G204 - Executing the configured command is the purpose of this option.
+	cmd := exec.CommandContext(ctx, shell, caller, command)
+	cmd.Env = append(os.Environ(), "CODER_URL="+serverURL.String())
+	cmd.Stdout = &stdout
+	cmd.Stderr = io.Discard
+	if err := cmd.Run(); err != nil {
+		return nil, xerrors.Errorf("run header command: %w", err)
+	}
+	var lines []string
+	scanner := bufio.NewScanner(&stdout)
+	for scanner.Scan() {
+		lines = append(lines, scanner.Text())
+	}
+	if err := scanner.Err(); err != nil {
+		return nil, xerrors.Errorf("scan header command output: %w", err)
+	}
+	return lines, nil
+}
+
+func parseHeaders(lines []string) (http.Header, error) {
+	headers := make(http.Header)
+	for i, line := range lines {
+		key, value, ok := strings.Cut(line, "=")
+		if !ok {
+			return nil, xerrors.Errorf("parse header %d: expected key=value", i+1)
+		}
+		headers.Add(key, value)
+	}
+	return headers, nil
+}
+
+// headerExpiry uses unverified JWT claims only to schedule refreshes. It does
+// not authenticate the headers, which are passed to the server unchanged.
+func headerExpiry(headers http.Header) time.Time {
+	var earliest time.Time
+	for _, values := range headers {
+		for _, value := range values {
+			if expiry, ok := jwtExpiry(value); ok && (earliest.IsZero() || expiry.Before(earliest)) {
+				earliest = expiry
+			}
+		}
+	}
+	if earliest.IsZero() {
+		return time.Time{}
+	}
+	return earliest.Add(-jwtExpirationSkew)
+}
+
+var headerJWTAlgorithms = []jose.SignatureAlgorithm{
+	jose.HS256, jose.HS384, jose.HS512,
+	jose.RS256, jose.RS384, jose.RS512,
+	jose.ES256, jose.ES384, jose.ES512,
+	jose.PS256, jose.PS384, jose.PS512,
+	jose.EdDSA,
+}
+
+func jwtExpiry(value string) (time.Time, bool) {
+	fields := strings.Fields(value)
+	switch {
+	case len(fields) == 1:
+		value = fields[0]
+	case len(fields) == 2 && strings.EqualFold(fields[0], "Bearer"):
+		value = fields[1]
+	default:
+		return time.Time{}, false
+	}
+	token, err := jwt.ParseSigned(value, headerJWTAlgorithms)
+	if err != nil {
+		return time.Time{}, false
+	}
+	var claims struct {
+		Expiry *jwt.NumericDate `json:"exp"`
+	}
+	// We don't have the issuer's key. exp is only a refresh hint, not proof
+	// of authentication.
+	if err := token.UnsafeClaimsWithoutVerification(&claims); err != nil || claims.Expiry == nil {
+		return time.Time{}, false
+	}
+	return claims.Expiry.Time(), true
+}

cli/header_internal_test.go

@@ -0,0 +1,192 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/go-jose/go-jose/v4"
+	"github.com/go-jose/go-jose/v4/jwt"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/quartz"
+)
+
+func TestCommandHeaderProvider(t *testing.T) {
+	t.Parallel()
+
+	t.Run("JWTRefresh", func(t *testing.T) {
+		t.Parallel()
+
+		t.Run("CacheExpiry", func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitShort)
+			clock := quartz.NewMock(t)
+			first := headerTestJWT(t, jwt.NewNumericDate(clock.Now().Add(time.Minute)))
+			later := headerTestJWT(t, jwt.NewNumericDate(clock.Now().Add(time.Hour)))
+			// Avoid nested quotes: cmd.exe does not unquote arguments like sh.
+			command := fmt.Sprintf("echo Authorization=Bearer %s&&echo X-Token=%s&&echo X-Removed=old", first, later)
+			provider := newTestHeaderProvider(clock, http.Header{"X-Static": {"unchanged"}}, command)
+
+			headers, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, "Bearer "+first, headers.Get("Authorization"))
+			require.Equal(t, "unchanged", headers.Get("X-Static"))
+			provider.command = fmt.Sprintf("echo Authorization=Bearer %s&&echo X-Value=updated", later)
+
+			clock.Advance(time.Minute - jwtExpirationSkew - time.Second).MustWait(ctx)
+			cached, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, headers, cached)
+			cached["Authorization"][0] = "mutated"
+			cached["X-Static"][0] = "mutated"
+			isolated, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, "Bearer "+first, isolated.Get("Authorization"))
+			require.Equal(t, "unchanged", isolated.Get("X-Static"))
+			require.Equal(t, headers, isolated)
+
+			clock.Advance(time.Second).MustWait(ctx)
+			refreshed, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, []string{"Bearer " + later}, refreshed.Values("Authorization"))
+			require.Equal(t, "unchanged", refreshed.Get("X-Static"))
+			require.Equal(t, "updated", refreshed.Get("X-Value"))
+			require.Empty(t, refreshed.Get("X-Removed"))
+			require.Empty(t, refreshed.Get("X-Token"))
+		})
+
+		t.Run("WithinClockSkew", func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitShort)
+			clock := quartz.NewMock(t)
+			shortLived := headerTestJWT(t, jwt.NewNumericDate(clock.Now().Add(jwtExpirationSkew/2)))
+			provider := newTestHeaderProvider(clock, nil, "echo Authorization=Bearer "+shortLived)
+			headers, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, "Bearer "+shortLived, headers.Get("Authorization"))
+
+			replacement := headerTestJWT(t, jwt.NewNumericDate(clock.Now().Add(time.Hour)))
+			provider.command = "echo Authorization=Bearer " + replacement
+			// The current token is already within the skew window, so the next call
+			// refreshes without advancing the clock.
+			refreshed, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, "Bearer "+replacement, refreshed.Get("Authorization"))
+
+			provider.command = "echo X-Value=changed"
+			cached, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, refreshed, cached)
+		})
+	})
+
+	t.Run("StaticOutput", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
+		clock := quartz.NewMock(t)
+		command := "echo X-Value=static&&echo X-Token=" + headerTestJWT(t, nil)
+		staticToken := headerTestJWT(t, jwt.NewNumericDate(clock.Now().Add(-time.Minute)))
+		provider := newTestHeaderProvider(clock, http.Header{"X-Static-Token": {staticToken}}, command)
+		headers, err := provider.Headers(ctx)
+		require.NoError(t, err)
+		require.Equal(t, "static", headers.Get("X-Value"))
+		require.Equal(t, staticToken, headers.Get("X-Static-Token"))
+
+		provider.command = "echo X-Value=changed"
+		clock.Advance(24 * time.Hour).MustWait(ctx)
+		cached, err := provider.Headers(ctx)
+		require.NoError(t, err)
+		require.Equal(t, headers, cached)
+	})
+
+	t.Run("CommandFailure", func(t *testing.T) {
+		t.Parallel()
+
+		t.Run("CommandError", func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitShort)
+			clock := quartz.NewMock(t)
+			command := "echo Authorization=Bearer " + headerTestJWT(t, jwt.NewNumericDate(clock.Now().Add(time.Minute)))
+			provider := newTestHeaderProvider(clock, nil, command)
+			_, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			clock.Advance(time.Minute).MustWait(ctx)
+			provider.command = "exit 1"
+
+			sent := false
+			transport := &codersdk.HeaderTransport{
+				Provider: provider,
+				Transport: roundTripper(func(*http.Request) (*http.Response, error) {
+					sent = true
+					return httptest.NewRecorder().Result(), nil
+				}),
+			}
+			req := httptest.NewRequest(http.MethodGet, provider.serverURL.String(), nil).WithContext(ctx)
+			res, err := transport.RoundTrip(req)
+			if res != nil {
+				require.NoError(t, res.Body.Close())
+			}
+			require.ErrorContains(t, err, "run header command")
+			require.NotContains(t, err.Error(), provider.command)
+			require.Nil(t, res)
+			require.False(t, sent, "a failed refresh must not send stale headers")
+
+			provider.command = "echo X-Value=recovered"
+			recovered, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, "recovered", recovered.Get("X-Value"))
+			require.Empty(t, recovered.Get("Authorization"))
+		})
+
+		t.Run("CanceledRefresh", func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitShort)
+			clock := quartz.NewMock(t)
+			token := headerTestJWT(t, jwt.NewNumericDate(clock.Now().Add(time.Minute)))
+			provider := newTestHeaderProvider(clock, nil, "echo Authorization=Bearer "+token)
+			_, err := provider.Headers(ctx)
+			require.NoError(t, err)
+			clock.Advance(time.Minute).MustWait(ctx)
+
+			canceled, cancel := context.WithCancel(ctx)
+			cancel()
+			headers, err := provider.Headers(canceled)
+			require.ErrorIs(t, err, context.Canceled)
+			require.Nil(t, headers)
+
+			provider.command = "echo X-Value=recovered"
+			headers, err = provider.Headers(ctx)
+			require.NoError(t, err)
+			require.Equal(t, "recovered", headers.Get("X-Value"))
+		})
+	})
+}
+
+func newTestHeaderProvider(clock quartz.Clock, static http.Header, command string) *commandHeaderProvider {
+	return &commandHeaderProvider{
+		serverURL:     &url.URL{Scheme: "https", Host: "coder.example.com"},
+		staticHeaders: static,
+		command:       command,
+		clock:         clock,
+	}
+}
+
+func headerTestJWT(t *testing.T, expiry *jwt.NumericDate) string {
+	t.Helper()
+	signer, err := jose.NewSigner(jose.SigningKey{
+		Algorithm: jose.HS512,
+		Key:       []byte(strings.Repeat("k", 64)),
+	}, nil)
+	require.NoError(t, err)
+	value, err := jwt.Signed(signer).Claims(jwt.Claims{Expiry: expiry}).Serialize()
+	require.NoError(t, err)
+	return value
+}

cli/root.go

@@ -1,8 +1,6 @@
 package cli
 
 import (
-	"bufio"
-	"bytes"
 	"context"
 	"crypto/tls"
 	"crypto/x509"
@@ -15,7 +13,6 @@ import (
 	"net/http"
 	"net/url"
 	"os"
-	"os/exec"
 	"os/signal"
 	"path/filepath"
 	"runtime"
@@ -473,7 +470,7 @@ func (r *RootCmd) Command(subcommands []*serpent.Command) (*serpent.Command, err
 		{
 			Flag:        varHeaderCommand,
 			Env:         "CODER_HEADER_COMMAND",
-			Description: "An external command that outputs additional HTTP headers added to all requests. The command must output each header as `key=value` on its own line.",
+			Description: "An external command that outputs additional HTTP headers added to all requests. The command must output each header as `key=value` on its own line. If a header value is a JWT, it will be refreshed based on the value in its `exp` field.",
 			Value:       serpent.StringOf(&r.headerCommand),
 			Group:       globalGroup,
 		},
@@ -1909,49 +1906,6 @@ func (r roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	return r(req)
 }
 
-// HeaderTransport creates a new transport that executes `--header-command`
-// if it is set to add headers for all outbound requests.
-func headerTransport(ctx context.Context, serverURL *url.URL, header []string, headerCommand string) (*codersdk.HeaderTransport, error) {
-	transport := &codersdk.HeaderTransport{
-		Transport: http.DefaultTransport,
-		Header:    http.Header{},
-	}
-	headers := header
-	if headerCommand != "" {
-		shell := "sh"
-		caller := "-c"
-		if runtime.GOOS == "windows" {
-			shell = "cmd.exe"
-			caller = "/c"
-		}
-		var outBuf bytes.Buffer
-		// #nosec
-		cmd := exec.CommandContext(ctx, shell, caller, headerCommand)
-		cmd.Env = append(os.Environ(), "CODER_URL="+serverURL.String())
-		cmd.Stdout = &outBuf
-		cmd.Stderr = io.Discard
-		err := cmd.Run()
-		if err != nil {
-			return nil, xerrors.Errorf("failed to run %v: %w", cmd.Args, err)
-		}
-		scanner := bufio.NewScanner(&outBuf)
-		for scanner.Scan() {
-			headers = append(headers, scanner.Text())
-		}
-		if err := scanner.Err(); err != nil {
-			return nil, xerrors.Errorf("scan %v: %w", cmd.Args, err)
-		}
-	}
-	for _, header := range headers {
-		parts := strings.SplitN(header, "=", 2)
-		if len(parts) < 2 {
-			return nil, xerrors.Errorf("split header %q had less than two parts", header)
-		}
-		transport.Header.Add(parts[0], parts[1])
-	}
-	return transport, nil
-}
-
 // PrintDeprecatedOptions loops through all command options, and
 // prints a warning for usage of deprecated options.
 func PrintDeprecatedOptions() serpent.MiddlewareFunc {

cli/support.go

@@ -134,7 +134,7 @@ func (r *RootCmd) supportBundle() *serpent.Command {
 				cliLog.Debug(inv.Context(), "running as owner")
 				client.HTTPClient.Transport = &codersdk.HeaderTransport{
 					Transport: client.HTTPClient.Transport,
-					Header:    http.Header{codersdk.BypassRatelimitHeader: {"true"}},
+					Provider:  codersdk.StaticHeaderProvider{Header: http.Header{codersdk.BypassRatelimitHeader: {"true"}}},
 				}
 			} else if !ok {
 				cliLog.Warn(inv.Context(), "not running as owner, not all information available")

cli/testdata/coder_--help.golden

@@ -111,7 +111,8 @@ variables or flags.
       --header-command string, $CODER_HEADER_COMMAND
           An external command that outputs additional HTTP headers added to all
           requests. The command must output each header as `key=value` on its
-          own line.
+          own line. If a header value is a JWT, it will be refreshed based on
+          the value in its `exp` field.
 
       --no-feature-warning bool, $CODER_NO_FEATURE_WARNING
           Suppress warnings about unlicensed features.

codersdk/client.go

@@ -876,19 +876,45 @@ func WithQueryParam(key, value string) RequestOption {
 	}
 }
 
+// HeaderProvider returns the headers to add to an outgoing request.
+// @typescript-ignore HeaderProvider
+type HeaderProvider interface {
+	Headers(ctx context.Context) (http.Header, error)
+}
+
+// StaticHeaderProvider always returns a copy of its configured headers.
+// @typescript-ignore StaticHeaderProvider
+type StaticHeaderProvider struct {
+	Header http.Header
+}
+
+// Headers returns a copy of the configured headers.
+func (p StaticHeaderProvider) Headers(context.Context) (http.Header, error) {
+	return p.Header.Clone(), nil
+}
+
 // HeaderTransport is a http.RoundTripper that adds some headers to all requests.
 // @typescript-ignore HeaderTransport
 type HeaderTransport struct {
 	Transport http.RoundTripper
-	Header    http.Header
+	Provider  HeaderProvider
 }
 
 var _ http.RoundTripper = &HeaderTransport{}
 
 func (h *HeaderTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	for k, v := range h.Header {
-		for _, vv := range v {
-			req.Header.Add(k, vv)
+	if h.Provider != nil {
+		headers, err := h.Provider.Headers(req.Context())
+		if err != nil {
+			if req.Body != nil {
+				_ = req.Body.Close()
+			}
+			return nil, xerrors.Errorf("get headers: %w", err)
+		}
+		for k, v := range headers {
+			for _, vv := range v {
+				req.Header.Add(k, vv)
+			}
 		}
 	}
 	if h.Transport == nil {

codersdk/workspacesdk/workspacesdk.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/google/uuid"
 	"go.opentelemetry.io/otel/baggage"
@@ -288,9 +289,18 @@ func (c *Client) DialAgent(dialCtx context.Context, agentID uuid.UUID, options *
 	controller.ResumeTokenCtrl = tailnet.NewBasicResumeTokenController(options.Logger, clk)
 
 	ip := tailnet.TailscaleServicePrefix.RandomAddr()
-	var header http.Header
-	if headerTransport, ok := c.client.HTTPClient.Transport.(*codersdk.HeaderTransport); ok {
-		header = headerTransport.Header
+	var getHeaders func() http.Header
+	if headerTransport, ok := c.client.HTTPClient.Transport.(*codersdk.HeaderTransport); ok && headerTransport.Provider != nil {
+		getHeaders = func() http.Header {
+			refreshCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+			defer cancel()
+			headers, err := headerTransport.Provider.Headers(refreshCtx)
+			if err != nil {
+				options.Logger.Error(ctx, "get connection headers", slog.Error(err))
+				return nil
+			}
+			return headers
+		}
 	}
 	var telemetrySink tailnet.TelemetrySink
 	if options.EnableTelemetry {
@@ -303,7 +313,7 @@ func (c *Client) DialAgent(dialCtx context.Context, agentID uuid.UUID, options *
 	conn, err := tailnet.NewConn(&tailnet.Options{
 		Addresses:           []netip.Prefix{netip.PrefixFrom(ip, 128)},
 		DERPMap:             connInfo.DERPMap,
-		DERPHeader:          &header,
+		DERPGetHeaders:      getHeaders,
 		DERPTLSConfig:       c.client.DERPTLSConfig(),
 		DERPForceWebSockets: connInfo.DERPForceWebSockets,
 		Logger:              options.Logger,
@@ -318,6 +328,9 @@ func (c *Client) DialAgent(dialCtx context.Context, agentID uuid.UUID, options *
 	}
 	defer func() {
 		if err != nil {
+			// getHeaders may block on a refresh while holding the DERP lock.
+			// Cancel its context before Close waits for that lock.
+			cancel()
 			_ = conn.Close()
 		}
 	}()

docs/reference/cli/index.md

@@ -151,7 +151,7 @@ Additional HTTP headers added to all requests. Provide as key=value. Can be spec
 | Type        | <code>string</code>                |
 | Environment | <code>$CODER_HEADER_COMMAND</code> |
 
-An external command that outputs additional HTTP headers added to all requests. The command must output each header as `key=value` on its own line.
+An external command that outputs additional HTTP headers added to all requests. The command must output each header as `key=value` on its own line. If a header value is a JWT, it will be refreshed based on the value in its `exp` field.
 
 ### --force-tty
 

enterprise/cli/testdata/coder_--help.golden

@@ -70,7 +70,8 @@ variables or flags.
       --header-command string, $CODER_HEADER_COMMAND
           An external command that outputs additional HTTP headers added to all
           requests. The command must output each header as `key=value` on its
-          own line.
+          own line. If a header value is a JWT, it will be refreshed based on
+          the value in its `exp` field.
 
       --no-feature-warning bool, $CODER_NO_FEATURE_WARNING
           Suppress warnings about unlicensed features.

scaletest/loadtestutil/client.go

@@ -1,6 +1,7 @@
 package loadtestutil
 
 import (
+	"context"
 	"maps"
 	"net/http"
 
@@ -9,43 +10,62 @@ import (
 	"github.com/coder/coder/v2/codersdk"
 )
 
-// DupClientCopyingHeaders duplicates the Client, but with an independent underlying HTTP transport, so that it will not
-// share connections with the client being duplicated. It copies any headers already on the existing transport as
-// [codersdk.HeaderTransport] and add the headers in the argument.
+// DupClientCopyingHeaders duplicates the client with an independent underlying
+// HTTP transport while retaining its header providers. Additional headers take
+// precedence over headers from the original client.
 func DupClientCopyingHeaders(client *codersdk.Client, header http.Header) (*codersdk.Client, error) {
 	nc := codersdk.New(client.URL, codersdk.WithLogger(client.Logger()))
 	nc.SessionTokenProvider = client.SessionTokenProvider
-	newHeader, t, err := extractHeaderAndInnerTransport(client.HTTPClient.Transport)
+	provider, t, err := extractHeaderAndInnerTransport(client.HTTPClient.Transport)
 	if err != nil {
 		return nil, xerrors.Errorf("extract headers: %w", err)
 	}
-	maps.Copy(newHeader, header)
-
 	nc.HTTPClient.Transport = &codersdk.HeaderTransport{
 		Transport: t.Clone(),
-		Header:    newHeader,
+		Provider: nestedHeaderProvider{
+			inner:  provider,
+			static: header.Clone(),
+		},
 	}
 	return nc, nil
 }
 
-func extractHeaderAndInnerTransport(rt http.RoundTripper) (http.Header, *http.Transport, error) {
-	if t, ok := rt.(*http.Transport); ok {
-		// base case
-		return make(http.Header), t, nil
+// nestedHeaderProvider resolves current headers before adding its static headers.
+type nestedHeaderProvider struct {
+	inner  codersdk.HeaderProvider
+	static http.Header
+}
+
+func (p nestedHeaderProvider) Headers(ctx context.Context) (http.Header, error) {
+	headers, err := p.inner.Headers(ctx)
+	if err != nil {
+		return nil, err
 	}
+	headers = headers.Clone()
+	if headers == nil {
+		headers = make(http.Header)
+	}
+	maps.Copy(headers, p.static.Clone())
+	return headers, nil
+}
+
+func extractHeaderAndInnerTransport(rt http.RoundTripper) (codersdk.HeaderProvider, *http.Transport, error) {
+	var provider codersdk.HeaderProvider = codersdk.StaticHeaderProvider{}
+	// Load-test clients have at most one header transport around the base transport.
 	if ht, ok := rt.(*codersdk.HeaderTransport); ok {
-		headers, t, err := extractHeaderAndInnerTransport(ht.Transport)
-		if err != nil {
-			return nil, nil, err
+		if ht.Provider != nil {
+			provider = ht.Provider
 		}
-		maps.Copy(headers, ht.Header)
-		return headers, t, nil
+		rt = ht.Transport
 	}
-	// unrecognized RoundTripper. Just return a default transport, since we only care about preserving headers.
-	t, ok := http.DefaultTransport.(*http.Transport)
+	// We assume only one layer of nesting for HeaderTransports.
+	t, ok := rt.(*http.Transport)
 	if !ok {
-		// unhittable, unless the Go stdlib changes.
-		return nil, nil, xerrors.New("DefaultTransport is not *http.Transport")
+		// unrecognized RoundTripper. Just return a default transport, since we only care about preserving headers.
+		t, ok = http.DefaultTransport.(*http.Transport)
+		if !ok {
+			return nil, nil, xerrors.New("DefaultTransport is not *http.Transport")
+		}
 	}
-	return make(http.Header), t, nil
+	return provider, t, nil
 }

scaletest/loadtestutil/client_test.go

@@ -1,32 +1,76 @@
 package loadtestutil_test
 
 import (
+	"context"
 	"net/http"
 	"net/url"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/scaletest/loadtestutil"
 )
 
+type mutableHeaderProvider struct {
+	headers http.Header
+	err     error
+}
+
+func (p *mutableHeaderProvider) Headers(context.Context) (http.Header, error) {
+	return p.headers, p.err
+}
+
+// Check that the copied client gets new headers when the provider changes them.
+func TestDupClientCopyingHeadersDynamic(t *testing.T) {
+	t.Parallel()
+	ctx := t.Context()
+	source := &mutableHeaderProvider{headers: http.Header{
+		"Authorization": {"Bearer first"},
+		"X-Override":    {"source"},
+		"X-Removed":     {"old"},
+	}}
+	client := codersdk.New(&url.URL{Scheme: "https", Host: "coder.example.com"})
+	client.HTTPClient.Transport = &codersdk.HeaderTransport{
+		Transport: http.DefaultTransport,
+		Provider:  source,
+	}
+	dup, err := loadtestutil.DupClientCopyingHeaders(client, http.Header{"X-Override": {"extra"}})
+	require.NoError(t, err)
+	transport, ok := dup.HTTPClient.Transport.(*codersdk.HeaderTransport)
+	require.True(t, ok)
+	headers, err := transport.Provider.Headers(ctx)
+	require.NoError(t, err)
+	require.Equal(t, "Bearer first", headers.Get("Authorization"))
+	require.Equal(t, "extra", headers.Get("X-Override"))
+	headers["Authorization"][0] = "mutated"
+	require.Equal(t, "Bearer first", source.headers.Get("Authorization"))
+
+	source.headers = http.Header{"Authorization": {"Bearer second"}}
+	headers, err = transport.Provider.Headers(ctx)
+	require.NoError(t, err)
+	require.Equal(t, "Bearer second", headers.Get("Authorization"))
+	require.Equal(t, "extra", headers.Get("X-Override"))
+	require.Empty(t, headers.Get("X-Removed"))
+
+	source.err = xerrors.New("credentials unavailable")
+	headers, err = transport.Provider.Headers(ctx)
+	require.ErrorIs(t, err, source.err)
+	require.Nil(t, headers)
+}
+
 func TestDupClientCopyingHeaders(t *testing.T) {
 	t.Parallel()
 	httpClient := &http.Client{
 		Transport: &codersdk.HeaderTransport{
-			Transport: &codersdk.HeaderTransport{
-				Transport: http.DefaultTransport,
-				Header: map[string][]string{
-					"X-Coder-Test":  {"foo"},
-					"X-Coder-Test3": {"socks"},
-					"X-Coder-Test5": {"ninjas"},
-				},
-			},
-			Header: map[string][]string{
+			Transport: http.DefaultTransport,
+			Provider: codersdk.StaticHeaderProvider{Header: http.Header{
 				"X-Coder-Test":  {"bar"},
 				"X-Coder-Test2": {"baz"},
-			},
+				"X-Coder-Test3": {"socks"},
+				"X-Coder-Test5": {"ninjas"},
+			}},
 		},
 	}
 	serverURL, err := url.Parse("http://coder.example.com")
@@ -43,10 +87,12 @@ func TestDupClientCopyingHeaders(t *testing.T) {
 	require.Equal(t, "test-token", dup.SessionToken())
 	ht, ok := dup.HTTPClient.Transport.(*codersdk.HeaderTransport)
 	require.True(t, ok)
-	require.Equal(t, "bar", ht.Header.Get("X-Coder-Test"))
-	require.Equal(t, "baz", ht.Header.Get("X-Coder-Test2"))
-	require.Equal(t, "clocks", ht.Header.Get("X-Coder-Test3"))
-	require.Equal(t, "bears", ht.Header.Get("X-Coder-Test4"))
-	require.Equal(t, "ninjas", ht.Header.Get("X-Coder-Test5"))
+	headers, err := ht.Provider.Headers(t.Context())
+	require.NoError(t, err)
+	require.Equal(t, "bar", headers.Get("X-Coder-Test"))
+	require.Equal(t, "baz", headers.Get("X-Coder-Test2"))
+	require.Equal(t, "clocks", headers.Get("X-Coder-Test3"))
+	require.Equal(t, "bears", headers.Get("X-Coder-Test4"))
+	require.Equal(t, "ninjas", headers.Get("X-Coder-Test5"))
 	require.NotEqual(t, http.DefaultTransport, ht.Transport)
 }

tailnet/conn.go

@@ -93,6 +93,8 @@ type Options struct {
 	Addresses  []netip.Prefix
 	DERPMap    *tailcfg.DERPMap
 	DERPHeader *http.Header
+	// DERPGetHeaders overrides DERPHeader on each connection attempt.
+	DERPGetHeaders func() http.Header
 	// DERPTLSConfig is an optional TLS config for DERP connections.
 	DERPTLSConfig *tls.Config
 	// DERPForceWebSockets determines whether websockets is always used for DERP
@@ -247,6 +249,7 @@ func NewConn(options *Options) (conn *Conn, err error) {
 	if options.DERPHeader != nil {
 		magicConn.SetDERPHeader(options.DERPHeader.Clone())
 	}
+	magicConn.SetDERPGetHeaders(options.DERPGetHeaders)
 	if options.DERPTLSConfig != nil {
 		magicConn.SetDERPTLSConfig(options.DERPTLSConfig)
 	}

tailnet/conn_test.go

@@ -4,8 +4,11 @@ import (
 	"context"
 	"errors"
 	"net"
+	"net/http"
+	"net/http/httptest"
 	"net/netip"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -13,6 +16,10 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
+	"tailscale.com/derp"
+	"tailscale.com/derp/derphttp"
+	"tailscale.com/tailcfg"
+	"tailscale.com/types/key"
 
 	"cdr.dev/slog/v3"
 	"github.com/coder/coder/v2/tailnet"
@@ -25,6 +32,67 @@ func TestMain(m *testing.M) {
 	goleak.VerifyTestMain(m, testutil.GoleakOptions...)
 }
 
+func TestDERPGetHeadersReconnect(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitLong)
+	logger := testutil.Logger(t)
+	server := derp.NewServer(key.NewNode(), tailnet.Logger(logger))
+	t.Cleanup(func() { _ = server.Close() })
+	handler := derphttp.Handler(server)
+	received := make(chan string, 10)
+	connections := make(chan net.Conn, 10)
+	httpServer := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("Upgrade") != "" {
+			select {
+			case received <- r.Header.Get("X-Token"):
+			case <-ctx.Done():
+				return
+			}
+		}
+		handler.ServeHTTP(w, r)
+	}))
+	httpServer.Config.ConnState = func(conn net.Conn, state http.ConnState) {
+		if state == http.StateHijacked {
+			select {
+			case connections <- conn:
+			case <-ctx.Done():
+			}
+		}
+	}
+	httpServer.StartTLS()
+	t.Cleanup(httpServer.Close)
+	addr, ok := httpServer.Listener.Addr().(*net.TCPAddr)
+	require.True(t, ok)
+	derpMap := &tailcfg.DERPMap{Regions: map[int]*tailcfg.DERPRegion{
+		1: {
+			RegionID:   1,
+			RegionCode: "test",
+			Nodes: []*tailcfg.DERPNode{{
+				Name: "test", RegionID: 1, IPv4: "127.0.0.1", IPv6: "none",
+				STUNPort: -1, DERPPort: addr.Port, InsecureForTests: true,
+			}},
+		},
+	}}
+	var token atomic.Pointer[string]
+	token.Store(new("first"))
+	conn, err := tailnet.NewConn(&tailnet.Options{
+		Addresses: []netip.Prefix{tailnet.TailscaleServicePrefix.RandomPrefix()},
+		Logger:    logger,
+		DERPMap:   derpMap,
+		DERPGetHeaders: func() http.Header {
+			return http.Header{"X-Token": {*token.Load()}}
+		},
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, conn.Close()) })
+	require.Equal(t, "first", testutil.TryReceive(ctx, t, received))
+	connection := testutil.TryReceive(ctx, t, connections)
+
+	token.Store(new("second"))
+	require.NoError(t, connection.Close())
+	require.Equal(t, "second", testutil.TryReceive(ctx, t, received))
+}
+
 func TestTailnet(t *testing.T) {
 	t.Parallel()
 	derpMap, _ := tailnettest.RunDERPAndSTUN(t)

vpn/client.go

@@ -107,7 +107,7 @@ func (*client) NewConn(initCtx context.Context, serverURL *url.URL, token string
 	sdk.SetSessionToken(token)
 	sdk.HTTPClient.Transport = &codersdk.HeaderTransport{
 		Transport: http.DefaultTransport,
-		Header:    headers.Clone(),
+		Provider:  codersdk.StaticHeaderProvider{Header: headers.Clone()},
 	}
 
 	// New context, separate from initCtx. We don't want to cancel the