feat: Use environment variables and startup script in agent (#1147)

These values were ignored. Environment variables are applied to
new sessions, and are refreshed on reconnect. This is cool because
a workspace could be updated with new environment variables without
requiring a complete start/stop.

The startup script is only ran once regardless of changes, which
feels like the expected behavior.

.vscode/settings.json

@@ -16,6 +16,7 @@
     "gographviz",
     "goleak",
     "gossh",
+    "gsyslog",
     "hashicorp",
     "hclsyntax",
     "httpmw",

agent/agent.go

@@ -11,9 +11,13 @@ import (
 	"os"
 	"os/exec"
 	"os/user"
+	"runtime"
 	"sync"
 	"time"
 
+	gsyslog "github.com/hashicorp/go-syslog"
+	"go.uber.org/atomic"
+
 	"cdr.dev/slog"
 	"github.com/coder/coder/agent/usershell"
 	"github.com/coder/coder/peer"
@@ -29,10 +33,11 @@ import (
 )
 
 type Options struct {
-	Logger slog.Logger
+	EnvironmentVariables map[string]string
+	StartupScript        string
 }
 
-type Dialer func(ctx context.Context, logger slog.Logger) (*peerbroker.Listener, error)
+type Dialer func(ctx context.Context, logger slog.Logger) (*Options, *peerbroker.Listener, error)
 
 func New(dialer Dialer, logger slog.Logger) io.Closer {
 	ctx, cancelFunc := context.WithCancel(context.Background())
@@ -55,16 +60,21 @@ type agent struct {
 	closeMutex    sync.Mutex
 	closed        chan struct{}
 
-	sshServer *ssh.Server
+	// Environment variables sent by Coder to inject for shell sessions.
+	// This is atomic because values can change after reconnect.
+	envVars       atomic.Value
+	startupScript atomic.Bool
+	sshServer     *ssh.Server
 }
 
 func (a *agent) run(ctx context.Context) {
+	var options *Options
 	var peerListener *peerbroker.Listener
 	var err error
 	// An exponential back-off occurs when the connection is failing to dial.
 	// This is to prevent server spam in case of a coderd outage.
 	for retrier := retry.New(50*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
-		peerListener, err = a.dialer(ctx, a.logger)
+		options, peerListener, err = a.dialer(ctx, a.logger)
 		if err != nil {
 			if errors.Is(err, context.Canceled) {
 				return
@@ -83,6 +93,20 @@ func (a *agent) run(ctx context.Context) {
 		return
 	default:
 	}
+	a.envVars.Store(options.EnvironmentVariables)
+
+	if a.startupScript.CAS(false, true) {
+		// The startup script has not ran yet!
+		go func() {
+			err := a.runStartupScript(ctx, options.StartupScript)
+			if errors.Is(err, context.Canceled) {
+				return
+			}
+			if err != nil {
+				a.logger.Warn(ctx, "agent script failed", slog.Error(err))
+			}
+		}()
+	}
 
 	for {
 		conn, err := peerListener.Accept()
@@ -101,6 +125,48 @@ func (a *agent) run(ctx context.Context) {
 	}
 }
 
+func (*agent) runStartupScript(ctx context.Context, script string) error {
+	if script == "" {
+		return nil
+	}
+	currentUser, err := user.Current()
+	if err != nil {
+		return xerrors.Errorf("get current user: %w", err)
+	}
+	username := currentUser.Username
+
+	shell, err := usershell.Get(username)
+	if err != nil {
+		return xerrors.Errorf("get user shell: %w", err)
+	}
+
+	var writer io.WriteCloser
+	// Attempt to use the syslog to write startup information.
+	writer, err = gsyslog.NewLogger(gsyslog.LOG_INFO, "USER", "coder-startup-script")
+	if err != nil {
+		// If the syslog isn't supported or cannot be created, use a text file in temp.
+		writer, err = os.CreateTemp("", "coder-startup-script.txt")
+		if err != nil {
+			return xerrors.Errorf("open startup script log file: %w", err)
+		}
+	}
+	defer func() {
+		_ = writer.Close()
+	}()
+	caller := "-c"
+	if runtime.GOOS == "windows" {
+		caller = "/c"
+	}
+	cmd := exec.CommandContext(ctx, shell, caller, script)
+	cmd.Stdout = writer
+	cmd.Stderr = writer
+	err = cmd.Run()
+	if err != nil {
+		return xerrors.Errorf("run: %w", err)
+	}
+	return nil
+}
+
 func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
 	go func() {
 		select {
@@ -230,8 +296,24 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
 
 	// OpenSSH executes all commands with the users current shell.
 	// We replicate that behavior for IDE support.
-	cmd := exec.CommandContext(session.Context(), shell, "-c", command)
+	caller := "-c"
+	if runtime.GOOS == "windows" {
+		caller = "/c"
+	}
+	cmd := exec.CommandContext(session.Context(), shell, caller, command)
 	cmd.Env = append(os.Environ(), session.Environ()...)
+
+	// Load environment variables passed via the agent.
+	envVars := a.envVars.Load()
+	if envVars != nil {
+		envVarMap, ok := envVars.(map[string]string)
+		if ok {
+			for key, value := range envVarMap {
+				cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", key, value))
+			}
+		}
+	}
+
 	executablePath, err := os.Executable()
 	if err != nil {
 		return xerrors.Errorf("getting os executable: %w", err)

agent/agent_test.go

@@ -12,12 +12,15 @@ import (
 	"strconv"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/pion/webrtc/v3"
 	"github.com/pkg/sftp"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/text/encoding/unicode"
+	"golang.org/x/text/transform"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
@@ -37,7 +40,7 @@ func TestAgent(t *testing.T) {
 	t.Parallel()
 	t.Run("SessionExec", func(t *testing.T) {
 		t.Parallel()
-		session := setupSSHSession(t)
+		session := setupSSHSession(t, nil)
 
 		command := "echo test"
 		if runtime.GOOS == "windows" {
@@ -50,7 +53,7 @@ func TestAgent(t *testing.T) {
 
 	t.Run("GitSSH", func(t *testing.T) {
 		t.Parallel()
-		session := setupSSHSession(t)
+		session := setupSSHSession(t, nil)
 		command := "sh -c 'echo $GIT_SSH_COMMAND'"
 		if runtime.GOOS == "windows" {
 			command = "cmd.exe /c echo %GIT_SSH_COMMAND%"
@@ -68,7 +71,7 @@ func TestAgent(t *testing.T) {
 			// it seems like it could be either.
 			t.Skip("ConPTY appears to be inconsistent on Windows.")
 		}
-		session := setupSSHSession(t)
+		session := setupSSHSession(t, nil)
 		command := "bash"
 		if runtime.GOOS == "windows" {
 			command = "cmd.exe"
@@ -128,7 +131,7 @@ func TestAgent(t *testing.T) {
 
 	t.Run("SFTP", func(t *testing.T) {
 		t.Parallel()
-		sshClient, err := setupAgent(t).SSHClient()
+		sshClient, err := setupAgent(t, nil).SSHClient()
 		require.NoError(t, err)
 		client, err := sftp.NewClient(sshClient)
 		require.NoError(t, err)
@@ -140,10 +143,52 @@ func TestAgent(t *testing.T) {
 		_, err = os.Stat(tempFile)
 		require.NoError(t, err)
 	})
+
+	t.Run("EnvironmentVariables", func(t *testing.T) {
+		t.Parallel()
+		key := "EXAMPLE"
+		value := "value"
+		session := setupSSHSession(t, &agent.Options{
+			EnvironmentVariables: map[string]string{
+				key: value,
+			},
+		})
+		command := "sh -c 'echo $" + key + "'"
+		if runtime.GOOS == "windows" {
+			command = "cmd.exe /c echo %" + key + "%"
+		}
+		output, err := session.Output(command)
+		require.NoError(t, err)
+		require.Equal(t, value, strings.TrimSpace(string(output)))
+	})
+
+	t.Run("StartupScript", func(t *testing.T) {
+		t.Parallel()
+		tempPath := filepath.Join(os.TempDir(), "content.txt")
+		content := "somethingnice"
+		setupAgent(t, &agent.Options{
+			StartupScript: "echo " + content + " > " + tempPath,
+		})
+		var gotContent string
+		require.Eventually(t, func() bool {
+			content, err := os.ReadFile(tempPath)
+			if err != nil {
+				return false
+			}
+			if runtime.GOOS == "windows" {
+				// Windows uses UTF16! 🪟🪟🪟
+				content, _, err = transform.Bytes(unicode.UTF16(unicode.LittleEndian, unicode.UseBOM).NewDecoder(), content)
+				require.NoError(t, err)
+			}
+			gotContent = string(content)
+			return true
+		}, 15*time.Second, 100*time.Millisecond)
+		require.Equal(t, content, strings.TrimSpace(gotContent))
+	})
 }
 
 func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
-	agentConn := setupAgent(t)
+	agentConn := setupAgent(t, nil)
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 	go func() {
@@ -171,18 +216,22 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
 	return exec.Command("ssh", args...)
 }
 
-func setupSSHSession(t *testing.T) *ssh.Session {
-	sshClient, err := setupAgent(t).SSHClient()
+func setupSSHSession(t *testing.T, options *agent.Options) *ssh.Session {
+	sshClient, err := setupAgent(t, options).SSHClient()
 	require.NoError(t, err)
 	session, err := sshClient.NewSession()
 	require.NoError(t, err)
 	return session
 }
 
-func setupAgent(t *testing.T) *agent.Conn {
+func setupAgent(t *testing.T, options *agent.Options) *agent.Conn {
+	if options == nil {
+		options = &agent.Options{}
+	}
 	client, server := provisionersdk.TransportPipe()
-	closer := agent.New(func(ctx context.Context, logger slog.Logger) (*peerbroker.Listener, error) {
-		return peerbroker.Listen(server, nil)
+	closer := agent.New(func(ctx context.Context, logger slog.Logger) (*agent.Options, *peerbroker.Listener, error) {
+		listener, err := peerbroker.Listen(server, nil)
+		return options, listener, err
 	}, slogtest.Make(t, nil).Leveled(slog.LevelDebug))
 	t.Cleanup(func() {
 		_ = client.Close()

cli/gitssh.go

@@ -7,10 +7,11 @@ import (
 	"os/exec"
 	"strings"
 
-	"github.com/coder/coder/cli/cliui"
-	"github.com/coder/coder/codersdk"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/cli/cliui"
+	"github.com/coder/coder/codersdk"
 )
 
 func gitssh() *cobra.Command {

coderd/coderd.go

@@ -197,7 +197,8 @@ func New(options *Options) (http.Handler, func()) {
 			r.Post("/google-instance-identity", api.postWorkspaceAuthGoogleInstanceIdentity)
 			r.Route("/me", func(r chi.Router) {
 				r.Use(httpmw.ExtractWorkspaceAgent(options.Database))
-				r.Get("/", api.workspaceAgentListen)
+				r.Get("/", api.workspaceAgentMe)
+				r.Get("/listen", api.workspaceAgentListen)
 				r.Get("/gitsshkey", api.agentGitSSHKey)
 				r.Get("/turn", api.workspaceAgentTurn)
 				r.Get("/iceservers", api.workspaceAgentICEServers)

coderd/workspaceagents.go

@@ -88,6 +88,18 @@ func (api *api) workspaceAgentDial(rw http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func (api *api) workspaceAgentMe(rw http.ResponseWriter, r *http.Request) {
+	agent := httpmw.WorkspaceAgent(r)
+	apiAgent, err := convertWorkspaceAgent(agent, api.AgentConnectionUpdateFrequency)
+	if err != nil {
+		httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
+			Message: fmt.Sprintf("convert workspace agent: %s", err),
+		})
+		return
+	}
+	httpapi.Write(rw, http.StatusOK, apiAgent)
+}
+
 func (api *api) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
 	api.websocketWaitMutex.Lock()
 	api.websocketWaitGroup.Add(1)

coderd/workspaceagents_test.go

@@ -102,6 +102,8 @@ func TestWorkspaceAgentListen(t *testing.T) {
 	})
 	_, err = conn.Ping()
 	require.NoError(t, err)
+	_, err = agentClient.WorkspaceAgent(context.Background(), codersdk.Me)
+	require.NoError(t, err)
 }
 
 func TestWorkspaceAgentTURN(t *testing.T) {

codersdk/workspaceagents.go

@@ -178,14 +178,14 @@ func (c *Client) AuthWorkspaceAzureInstanceIdentity(ctx context.Context) (Worksp
 
 // ListenWorkspaceAgent connects as a workspace agent identifying with the session token.
 // On each inbound connection request, connection info is fetched.
-func (c *Client) ListenWorkspaceAgent(ctx context.Context, logger slog.Logger) (*peerbroker.Listener, error) {
-	serverURL, err := c.URL.Parse("/api/v2/workspaceagents/me")
+func (c *Client) ListenWorkspaceAgent(ctx context.Context, logger slog.Logger) (*agent.Options, *peerbroker.Listener, error) {
+	serverURL, err := c.URL.Parse("/api/v2/workspaceagents/me/listen")
 	if err != nil {
-		return nil, xerrors.Errorf("parse url: %w", err)
+		return nil, nil, xerrors.Errorf("parse url: %w", err)
 	}
 	jar, err := cookiejar.New(nil)
 	if err != nil {
-		return nil, xerrors.Errorf("create cookie jar: %w", err)
+		return nil, nil, xerrors.Errorf("create cookie jar: %w", err)
 	}
 	jar.SetCookies(serverURL, []*http.Cookie{{
 		Name:  httpmw.AuthCookie,
@@ -201,17 +201,17 @@ func (c *Client) ListenWorkspaceAgent(ctx context.Context, logger slog.Logger) (
 	})
 	if err != nil {
 		if res == nil {
-			return nil, err
+			return nil, nil, err
 		}
-		return nil, readBodyAsError(res)
+		return nil, nil, readBodyAsError(res)
 	}
 	config := yamux.DefaultConfig()
 	config.LogOutput = io.Discard
 	session, err := yamux.Client(websocket.NetConn(ctx, conn, websocket.MessageBinary), config)
 	if err != nil {
-		return nil, xerrors.Errorf("multiplex client: %w", err)
+		return nil, nil, xerrors.Errorf("multiplex client: %w", err)
 	}
-	return peerbroker.Listen(session, func(ctx context.Context) ([]webrtc.ICEServer, *peer.ConnOptions, error) {
+	listener, err := peerbroker.Listen(session, func(ctx context.Context) ([]webrtc.ICEServer, *peer.ConnOptions, error) {
 		// This can be cached if it adds to latency too much.
 		res, err := c.request(ctx, http.MethodGet, "/api/v2/workspaceagents/me/iceservers", nil)
 		if err != nil {
@@ -237,6 +237,17 @@ func (c *Client) ListenWorkspaceAgent(ctx context.Context, logger slog.Logger) (
 			Logger:        logger,
 		}, nil
 	})
+	if err != nil {
+		return nil, nil, xerrors.Errorf("listen peerbroker: %w", err)
+	}
+	workspaceAgent, err := c.WorkspaceAgent(ctx, Me)
+	if err != nil {
+		return nil, nil, xerrors.Errorf("get workspace agent: %w", err)
+	}
+	return &agent.Options{
+		EnvironmentVariables: workspaceAgent.EnvironmentVariables,
+		StartupScript:        workspaceAgent.StartupScript,
+	}, listener, err
 }
 
 // DialWorkspaceAgent creates a connection to the specified resource.
@@ -313,7 +324,7 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, opti
 
 // WorkspaceAgent returns an agent by ID.
 func (c *Client) WorkspaceAgent(ctx context.Context, id uuid.UUID) (WorkspaceAgent, error) {
-	res, err := c.request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspaceagents/%s", id), nil)
+	res, err := c.request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspaceagents/%s", uuidOrMe(id)), nil)
 	if err != nil {
 		return WorkspaceAgent{}, err
 	}

go.mod

@@ -66,6 +66,7 @@ require (
 	github.com/golang-migrate/migrate/v4 v4.15.1
 	github.com/google/go-github/v43 v43.0.1-0.20220414155304-00e42332e405
 	github.com/google/uuid v1.3.0
+	github.com/hashicorp/go-syslog v1.0.0
 	github.com/hashicorp/go-version v1.4.0
 	github.com/hashicorp/hc-install v0.3.1
 	github.com/hashicorp/hcl/v2 v2.12.0
@@ -107,6 +108,7 @@ require (
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 	golang.org/x/sys v0.0.0-20220412211240-33da011f77ad
 	golang.org/x/term v0.0.0-20210927222741-03fcf44c2211
+	golang.org/x/text v0.3.7
 	golang.org/x/xerrors v0.0.0-20220411194840-2f41105eb62f
 	google.golang.org/api v0.75.0
 	google.golang.org/protobuf v1.28.0
@@ -227,7 +229,6 @@ require (
 	github.com/zclconf/go-cty v1.10.0 // indirect
 	github.com/zeebo/errs v1.2.2 // indirect
 	go.opencensus.io v0.23.0 // indirect
-	golang.org/x/text v0.3.7 // indirect
 	golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto v0.0.0-20220421151946-72621c1f0bd3 // indirect

go.sum

@@ -881,6 +881,7 @@ github.com/hashicorp/go-rootcerts v1.0.0/go.mod h1:K6zTfqpRlCUIjkwsN4Z+hiSfzSTQa
 github.com/hashicorp/go-rootcerts v1.0.2/go.mod h1:pqUvnprVnM5bf7AOirdbb01K4ccR319Vf4pU3K5EGc8=
 github.com/hashicorp/go-sockaddr v1.0.0/go.mod h1:7Xibr9yA9JjQq1JpNB2Vw7kxv8xerXegt+ozgdvDeDU=
 github.com/hashicorp/go-sockaddr v1.0.2/go.mod h1:rB4wwRAUzs07qva3c5SdrY/NEtAUjGlgmH/UkBUC97A=
+github.com/hashicorp/go-syslog v1.0.0 h1:KaodqZuhUoZereWVIYmpUgZysurB1kBLX2j0MwMrUAE=
 github.com/hashicorp/go-syslog v1.0.0/go.mod h1:qPfqrKkXGihmCqbJM2mZgkZGvKG1dFdvsLplgctolz4=
 github.com/hashicorp/go-uuid v1.0.0/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro=
 github.com/hashicorp/go-uuid v1.0.1/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro=