fix: Use instance identity session token for git subcommands (#4884)

This broke using gitssh with instance identity!

agent/agent.go

@@ -56,7 +56,7 @@ const (
 
 type Options struct {
 	Filesystem             afero.Fs
-	ExchangeToken          func(ctx context.Context) error
+	ExchangeToken          func(ctx context.Context) (string, error)
 	Client                 Client
 	ReconnectingPTYTimeout time.Duration
 	EnvironmentVariables   map[string]string
@@ -78,6 +78,11 @@ func New(options Options) io.Closer {
 	if options.Filesystem == nil {
 		options.Filesystem = afero.NewOsFs()
 	}
+	if options.ExchangeToken == nil {
+		options.ExchangeToken = func(ctx context.Context) (string, error) {
+			return "", nil
+		}
+	}
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	server := &agent{
 		reconnectingPTYTimeout: options.ReconnectingPTYTimeout,
@@ -97,7 +102,7 @@ func New(options Options) io.Closer {
 type agent struct {
 	logger        slog.Logger
 	client        Client
-	exchangeToken func(ctx context.Context) error
+	exchangeToken func(ctx context.Context) (string, error)
 	filesystem    afero.Fs
 
 	reconnectingPTYs       sync.Map
@@ -110,8 +115,9 @@ type agent struct {
 
 	envVars map[string]string
 	// metadata is atomic because values can change after reconnection.
-	metadata  atomic.Value
-	sshServer *ssh.Server
+	metadata     atomic.Value
+	sessionToken atomic.Pointer[string]
+	sshServer    *ssh.Server
 
 	network *tailnet.Conn
 	stats   *Stats
@@ -147,14 +153,13 @@ func (a *agent) run(ctx context.Context) error {
 	// This allows the agent to refresh it's token if necessary.
 	// For instance identity this is required, since the instance
 	// may not have re-provisioned, but a new agent ID was created.
-	if a.exchangeToken != nil {
-		err := a.exchangeToken(ctx)
-		if err != nil {
-			return xerrors.Errorf("exchange token: %w", err)
-		}
+	sessionToken, err := a.exchangeToken(ctx)
+	if err != nil {
+		return xerrors.Errorf("exchange token: %w", err)
 	}
+	a.sessionToken.Store(&sessionToken)
 
-	err := a.client.PostWorkspaceAgentVersion(ctx, buildinfo.Version())
+	err = a.client.PostWorkspaceAgentVersion(ctx, buildinfo.Version())
 	if err != nil {
 		return xerrors.Errorf("update workspace agent version: %w", err)
 	}
@@ -571,6 +576,9 @@ func (a *agent) createCommand(ctx context.Context, rawCommand string, env []stri
 	unixExecutablePath := strings.ReplaceAll(executablePath, "\\", "/")
 	cmd.Env = append(cmd.Env, fmt.Sprintf(`GIT_SSH_COMMAND=%s gitssh --`, unixExecutablePath))
 
+	// Specific Coder subcommands require the agent token exposed!
+	cmd.Env = append(cmd.Env, fmt.Sprintf("CODER_AGENT_TOKEN=%s", *a.sessionToken.Load()))
+
 	// Set SSH connection environment variables (these are also set by OpenSSH
 	// and thus expected to be present by SSH clients). Since the agent does
 	// networking in-memory, trying to provide accurate values here would be

agent/agent_test.go

@@ -533,9 +533,9 @@ func TestAgent(t *testing.T) {
 		}
 		initialized := atomic.Int32{}
 		closer := agent.New(agent.Options{
-			ExchangeToken: func(ctx context.Context) error {
+			ExchangeToken: func(ctx context.Context) (string, error) {
 				initialized.Add(1)
-				return nil
+				return "", nil
 			},
 			Client: client,
 			Logger: slogtest.Make(t, nil).Leveled(slog.LevelInfo),
@@ -567,8 +567,8 @@ func TestAgent(t *testing.T) {
 		}
 		filesystem := afero.NewMemMapFs()
 		closer := agent.New(agent.Options{
-			ExchangeToken: func(ctx context.Context) error {
-				return nil
+			ExchangeToken: func(ctx context.Context) (string, error) {
+				return "", nil
 			},
 			Client:     client,
 			Logger:     slogtest.Make(t, nil).Leveled(slog.LevelInfo),

cli/agent.go

@@ -156,22 +156,19 @@ func workspaceAgent() *cobra.Command {
 			closer := agent.New(agent.Options{
 				Client: client,
 				Logger: logger,
-				ExchangeToken: func(ctx context.Context) error {
+				ExchangeToken: func(ctx context.Context) (string, error) {
 					if exchangeToken == nil {
-						return nil
+						return client.SessionToken, nil
 					}
 					resp, err := exchangeToken(ctx)
 					if err != nil {
-						return err
+						return "", err
 					}
 					client.SessionToken = resp.SessionToken
-					return nil
+					return "", nil
 				},
 				EnvironmentVariables: map[string]string{
-					// Override the "CODER_AGENT_TOKEN" variable in all
-					// shells so "gitssh" and "gitaskpass" works!
-					"CODER_AGENT_TOKEN": client.SessionToken,
-					"GIT_ASKPASS":       executablePath,
+					"GIT_ASKPASS": executablePath,
 				},
 			})
 			<-cmd.Context().Done()