diff --git a/agent/agent.go b/agent/agent.go index 350962720d..a0d215ebd0 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -377,7 +377,7 @@ func (a *agent) init() { a.containerAPI = agentcontainers.NewAPI(a.logger.Named("containers"), containerAPIOpts...) a.filesAPI = agentfiles.NewAPI(a.logger.Named("files"), a.filesystem) - a.processAPI = agentproc.NewAPI(a.logger.Named("processes"), a.execer) + a.processAPI = agentproc.NewAPI(a.logger.Named("processes"), a.execer, a.updateCommandEnv) a.reconnectingPTYServer = reconnectingpty.NewServer( a.logger.Named("reconnecting-pty"), diff --git a/agent/agentproc/api.go b/agent/agentproc/api.go index e6ac27524d..5a8c6d42d1 100644 --- a/agent/agentproc/api.go +++ b/agent/agentproc/api.go @@ -22,10 +22,10 @@ type API struct { } // NewAPI creates a new process API handler. -func NewAPI(logger slog.Logger, execer agentexec.Execer) *API { +func NewAPI(logger slog.Logger, execer agentexec.Execer, updateEnv func(current []string) (updated []string, err error)) *API { return &API{ logger: logger, - manager: newManager(logger, execer), + manager: newManager(logger, execer, updateEnv), } } diff --git a/agent/agentproc/api_test.go b/agent/agentproc/api_test.go index 75df874451..dc204c9076 100644 --- a/agent/agentproc/api_test.go +++ b/agent/agentproc/api_test.go @@ -88,11 +88,18 @@ func postSignal(t *testing.T, handler http.Handler, id string, req workspacesdk. // execer, returning the handler and API. func newTestAPI(t *testing.T) http.Handler { t.Helper() + return newTestAPIWithUpdateEnv(t, nil) +} + +// newTestAPIWithUpdateEnv creates a new API with an optional +// updateEnv hook for testing environment injection. +func newTestAPIWithUpdateEnv(t *testing.T, updateEnv func([]string) ([]string, error)) http.Handler { + t.Helper() logger := slogtest.Make(t, &slogtest.Options{ IgnoreErrors: true, }).Leveled(slog.LevelDebug) - api := agentproc.NewAPI(logger, agentexec.DefaultExecer) + api := agentproc.NewAPI(logger, agentexec.DefaultExecer, updateEnv) t.Cleanup(func() { _ = api.Close() }) @@ -257,6 +264,54 @@ func TestStartProcess(t *testing.T) { require.Equal(t, 0, *resp.ExitCode) require.Contains(t, strings.TrimSpace(resp.Output), envVal) }) + + t.Run("UpdateEnvHook", func(t *testing.T) { + t.Parallel() + + envKey := fmt.Sprintf("TEST_UPDATE_ENV_%d", time.Now().UnixNano()) + envVal := "injected_by_hook" + + handler := newTestAPIWithUpdateEnv(t, func(current []string) ([]string, error) { + return append(current, fmt.Sprintf("%s=%s", envKey, envVal)), nil + }) + + // The process should see the variable even though it + // was not passed in req.Env. + id := startAndGetID(t, handler, workspacesdk.StartProcessRequest{ + Command: fmt.Sprintf("printenv %s", envKey), + }) + + resp := waitForExit(t, handler, id) + require.NotNil(t, resp.ExitCode) + require.Equal(t, 0, *resp.ExitCode) + require.Contains(t, strings.TrimSpace(resp.Output), envVal) + }) + + t.Run("UpdateEnvHookOverriddenByReqEnv", func(t *testing.T) { + t.Parallel() + + envKey := fmt.Sprintf("TEST_OVERRIDE_%d", time.Now().UnixNano()) + hookVal := "from_hook" + reqVal := "from_request" + + handler := newTestAPIWithUpdateEnv(t, func(current []string) ([]string, error) { + return append(current, fmt.Sprintf("%s=%s", envKey, hookVal)), nil + }) + + // req.Env should take precedence over the hook. + id := startAndGetID(t, handler, workspacesdk.StartProcessRequest{ + Command: fmt.Sprintf("printenv %s", envKey), + Env: map[string]string{envKey: reqVal}, + }) + + resp := waitForExit(t, handler, id) + require.NotNil(t, resp.ExitCode) + require.Equal(t, 0, *resp.ExitCode) + // When duplicate env vars exist, shells use the last + // value. Since req.Env is appended after the hook, + // the request value wins. + require.Contains(t, strings.TrimSpace(resp.Output), reqVal) + }) } func TestListProcesses(t *testing.T) { diff --git a/agent/agentproc/process.go b/agent/agentproc/process.go index 3059a40bc0..c797bb633e 100644 --- a/agent/agentproc/process.go +++ b/agent/agentproc/process.go @@ -65,21 +65,23 @@ func (p *process) output() (string, *workspacesdk.ProcessTruncation) { // manager tracks processes spawned by the agent. type manager struct { - mu sync.Mutex - logger slog.Logger - execer agentexec.Execer - clock quartz.Clock - procs map[string]*process - closed bool + mu sync.Mutex + logger slog.Logger + execer agentexec.Execer + clock quartz.Clock + procs map[string]*process + closed bool + updateEnv func(current []string) (updated []string, err error) } // newManager creates a new process manager. -func newManager(logger slog.Logger, execer agentexec.Execer) *manager { +func newManager(logger slog.Logger, execer agentexec.Execer, updateEnv func(current []string) (updated []string, err error)) *manager { return &manager{ - logger: logger, - execer: execer, - clock: quartz.NewReal(), - procs: make(map[string]*process), + logger: logger, + execer: execer, + clock: quartz.NewReal(), + procs: make(map[string]*process), + updateEnv: updateEnv, } } @@ -116,13 +118,31 @@ func (m *manager) start(req workspacesdk.StartProcessRequest) (*process, error) cmd.Stdout = buf cmd.Stderr = buf - if len(req.Env) > 0 { - cmd.Env = os.Environ() - for k, v := range req.Env { - cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v)) + // Build the process environment. If the manager has an + // updateEnv hook (provided by the agent), use it to get the + // full agent environment including GIT_ASKPASS, CODER_* vars, + // etc. Otherwise fall back to the current process env. + baseEnv := os.Environ() + if m.updateEnv != nil { + updated, err := m.updateEnv(baseEnv) + if err != nil { + m.logger.Warn( + context.Background(), + "failed to update command environment, falling back to os env", + slog.Error(err), + ) + } else { + baseEnv = updated } } + // Always set cmd.Env explicitly so that req.Env overrides + // are applied on top of the full agent environment. + cmd.Env = baseEnv + for k, v := range req.Env { + cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v)) + } + if err := cmd.Start(); err != nil { cancel() return nil, xerrors.Errorf("start process: %w", err)