feat(agent): add script data dir for binaries and files

The agent is extended with a `--script-data-dir` flag, defaulting to the OS temp dir. This dir is used for storing `coder-script/bin` and `coder-script/[script uuid]`. The former is a place for all scripts to place executable binaries that will be available by other scripts, SSH sessions, etc. The latter is a place for the script to store files. Since we default to OS temp dir, files are ephemeral by default. In the future, we may consider adding new env vars or changing the default storage location. Workspace startup speed could potentially benefit from scripts being able to skip steps that require downloading software. We may also extend this with more env variables (e.g. persistent storage in HOME). Fixes #11131
coder · mafredri · Feb 20, 2024 · Feb 19, 2024 · Feb 19, 2024 · Feb 19, 2024
commit ebab733c930832c301dd0b076e53db353ee679e2
diff --git a/agent/agent.go b/agent/agent.go
@@ -66,6 +66,7 @@ type Options struct {
 	Filesystem                   afero.Fs
 	LogDir                       string
 	TempDir                      string
+	ScriptDataDir                string
 	ExchangeToken                func(ctx context.Context) (string, error)
 	Client                       Client
 	ReconnectingPTYTimeout       time.Duration
@@ -115,6 +116,12 @@ func New(options Options) Agent {
 		}
 		options.LogDir = options.TempDir
 	}
+	if options.ScriptDataDir == "" {
+		if options.TempDir != os.TempDir() {
+			options.Logger.Debug(context.Background(), "script data dir not set, using temp dir", slog.F("temp_dir", options.TempDir))
+		}
+		options.ScriptDataDir = options.TempDir
+	}
 	if options.ExchangeToken == nil {
 		options.ExchangeToken = func(ctx context.Context) (string, error) {
 			return "", nil
@@ -152,6 +159,7 @@ func New(options Options) Agent {
 		filesystem:                   options.Filesystem,
 		logDir:                       options.LogDir,
 		tempDir:                      options.TempDir,
+		scriptDataDir:                options.ScriptDataDir,
 		lifecycleUpdate:              make(chan struct{}, 1),
 		lifecycleReported:            make(chan codersdk.WorkspaceAgentLifecycle, 1),
 		lifecycleStates:              []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},
@@ -183,6 +191,7 @@ type agent struct {
 	filesystem        afero.Fs
 	logDir            string
 	tempDir           string
+	scriptDataDir     string
 	// ignorePorts tells the api handler which ports to ignore when
 	// listing all listening ports. This is helpful to hide ports that
 	// are used by the agent, that the user does not care about.
@@ -250,6 +259,7 @@ func (a *agent) init(ctx context.Context) {
 	a.sshServer = sshSrv
 	a.scriptRunner = agentscripts.New(agentscripts.Options{
 		LogDir:     a.logDir,
+		DataDir:    a.scriptDataDir,
 		Logger:     a.logger,
 		SSHServer:  sshSrv,
 		Filesystem: a.filesystem,
@@ -954,6 +964,13 @@ func (a *agent) updateCommandEnv(current []string) (updated []string, err error)
 		envs[k] = v
 	}
 
+	// Prepend the agent script bin directory to the PATH
+	// (this is where Coder modules place their binaries).
+	if _, ok := envs["PATH"]; !ok {
+		envs["PATH"] = os.Getenv("PATH")
+	}
+	envs["PATH"] = a.scriptRunner.ScriptBinDir() + ":" + envs["PATH"]
+
 	for k, v := range envs {
 		updated = append(updated, fmt.Sprintf("%s=%s", k, v))
 	}

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -286,6 +286,11 @@ func TestAgent_SessionExec(t *testing.T) {
 func TestAgent_Session_EnvironmentVariables(t *testing.T) {
 	t.Parallel()
 
+	tmpdir := t.TempDir()
+
+	// Defined by the coder script runner.
+	scriptBinDir := filepath.Join(tmpdir, "coder-script-data", "bin")
+
 	manifest := agentsdk.Manifest{
 		EnvironmentVariables: map[string]string{
 			"MY_MANIFEST":         "true",
@@ -294,7 +299,10 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) {
 		},
 	}
 	banner := codersdk.ServiceBannerConfig{}
+	fs := afero.NewMemMapFs()
 	session := setupSSHSession(t, manifest, banner, nil, func(_ *agenttest.Client, opts *agent.Options) {
+		opts.Filesystem = fs
+		opts.ScriptDataDir = tmpdir
 		opts.EnvironmentVariables["MY_OVERRIDE"] = "true"
 	})
 
@@ -341,6 +349,7 @@ func TestAgent_Session_EnvironmentVariables(t *testing.T) {
 		"MY_OVERRIDE":         "true",  // From the agent environment variables option, overrides manifest.
 		"MY_SESSION_MANIFEST": "false", // From the manifest, overrides session env.
 		"MY_SESSION":          "true",  // From the session.
+		"PATH":                scriptBinDir + ":",
 	} {
 		t.Run(k, func(t *testing.T) {
 			echoEnv(t, stdin, k)

diff --git a/agent/agentscripts/agentscripts.go b/agent/agentscripts/agentscripts.go
@@ -43,6 +43,7 @@ var (
 
 // Options are a set of options for the runner.
 type Options struct {
+	DataDir    string
 	LogDir     string
 	Logger     slog.Logger
 	SSHServer  *agentssh.Server
@@ -59,6 +60,7 @@ func New(opts Options) *Runner {
 		cronCtxCancel: cronCtxCancel,
 		cron:          cron.New(cron.WithParser(parser)),
 		closed:        make(chan struct{}),
+		dataDir:       filepath.Join(opts.DataDir, "coder-script-data"),
 		scriptsExecuted: prometheus.NewCounterVec(prometheus.CounterOpts{
 			Namespace: "agent",
 			Subsystem: "scripts",
@@ -78,13 +80,20 @@ type Runner struct {
 	cron          *cron.Cron
 	initialized   atomic.Bool
 	scripts       []codersdk.WorkspaceAgentScript
+	dataDir       string
 
 	// scriptsExecuted includes all scripts executed by the workspace agent. Agents
 	// execute startup scripts, and scripts on a cron schedule. Both will increment
 	// this counter.
 	scriptsExecuted *prometheus.CounterVec
 }
 
+// ScriptBinDir returns the directory where scripts can store executable
+// binaries.
+func (r *Runner) ScriptBinDir() string {
+	return filepath.Join(r.dataDir, "bin")
+}
+
 func (r *Runner) RegisterMetrics(reg prometheus.Registerer) {
 	if reg == nil {
 		// If no registry, do nothing.
@@ -104,6 +113,11 @@ func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript) error {
 	r.scripts = scripts
 	r.Logger.Info(r.cronCtx, "initializing agent scripts", slog.F("script_count", len(scripts)), slog.F("log_dir", r.LogDir))
 
+	err := r.Filesystem.MkdirAll(r.ScriptBinDir(), 0o700)
+	if err != nil {
+		return xerrors.Errorf("create script bin dir: %w", err)
+	}
+
 	for _, script := range scripts {
 		if script.Cron == "" {
 			continue
@@ -208,7 +222,18 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript)
 	if !filepath.IsAbs(logPath) {
 		logPath = filepath.Join(r.LogDir, logPath)
 	}
-	logger := r.Logger.With(slog.F("log_path", logPath))
+
+	scriptDataDir := filepath.Join(r.dataDir, script.LogSourceID.String())
+	err := r.Filesystem.MkdirAll(scriptDataDir, 0o700)
+	if err != nil {
+		return xerrors.Errorf("%s script: create script temp dir: %w", scriptDataDir, err)
+	}
+
+	logger := r.Logger.With(
+		slog.F("log_source_id", script.LogSourceID),
+		slog.F("log_path", logPath),
+		slog.F("script_data_dir", scriptDataDir),
+	)
 	logger.Info(ctx, "running agent script", slog.F("script", script.Script))
 
 	fileWriter, err := r.Filesystem.OpenFile(logPath, os.O_CREATE|os.O_RDWR, 0o600)
@@ -238,6 +263,13 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript)
 	cmd.WaitDelay = 10 * time.Second
 	cmd.Cancel = cmdCancel(cmd)
 
+	// Expose env vars that can be used in the script for storing data
+	// and binaries. In the future, we may want to expose more env vars
+	// for the script to use, like CODER_SCRIPT_DATA_DIR for persistent
+	// storage.
+	cmd.Env = append(cmd.Env, "CODER_SCRIPT_DATA_DIR="+scriptDataDir)
+	cmd.Env = append(cmd.Env, "CODER_SCRIPT_BIN_DIR="+r.ScriptBinDir())
+
 	send, flushAndClose := agentsdk.LogsSender(script.LogSourceID, r.PatchLogs, logger)
 	// If ctx is canceled here (or in a writer below), we may be
 	// discarding logs, but that's okay because we're shutting down

diff --git a/cli/agent.go b/cli/agent.go
@@ -40,6 +40,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 	var (
 		auth                string
 		logDir              string
+		scriptDataDir       string
 		pprofAddress        string
 		noReap              bool
 		sshMaxTimeout       time.Duration
@@ -289,6 +290,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				Client:            client,
 				Logger:            logger,
 				LogDir:            logDir,
+				ScriptDataDir:     scriptDataDir,
 				TailnetListenPort: uint16(tailnetListenPort),
 				ExchangeToken: func(ctx context.Context) (string, error) {
 					if exchangeToken == nil {
@@ -339,6 +341,13 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 			Env:         "CODER_AGENT_LOG_DIR",
 			Value:       clibase.StringOf(&logDir),
 		},
+		{
+			Flag:        "script-data-dir",
+			Default:     os.TempDir(),
+			Description: "Specify the location for storing script data.",
+			Env:         "CODER_AGENT_SCRIPT_DIR",
+			Value:       clibase.StringOf(&scriptDataDir),
+		},
 		{
 			Flag:        "pprof-address",
 			Default:     "127.0.0.1:6060",