Skip to content

Commit c63f569

Browse files
mafredrimafredri
authored
refactor(agent/agentssh): move envs to agent and add agentssh config struct (#12204)
This commit refactors where custom environment variables are set in the workspace and decouples agent specific configs from the `agentssh.Server`. To reproduce all functionality, `agentssh.Config` is introduced. The custom environment variables are now configured in `agent/agent.go` and the agent retains control of the final state. This will allow for easier extension in the future and keep other modules decoupled.
1 parent 817cc78 commit c63f569

File tree

9 files changed

+268
-137
lines changed

9 files changed

+268
-137
lines changed

‎agent/agent.go

Lines changed: 89 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ func New(options Options) Agent {
146146
logger: options.Logger,
147147
closeCancel: cancelFunc,
148148
closed: make(chan struct{}),
149-
envVars: options.EnvironmentVariables,
149+
environmentVariables: options.EnvironmentVariables,
150150
client: options.Client,
151151
exchangeToken: options.ExchangeToken,
152152
filesystem: options.Filesystem,
@@ -169,6 +169,8 @@ func New(options Options) Agent {
169169
prometheusRegistry: prometheusRegistry,
170170
metrics: newAgentMetrics(prometheusRegistry),
171171
}
172+
a.serviceBanner.Store(new(codersdk.ServiceBannerConfig))
173+
a.sessionToken.Store(new(string))
172174
a.init(ctx)
173175
return a
174176
}
@@ -196,7 +198,7 @@ type agent struct {
196198
closeMutex sync.Mutex
197199
closed chan struct{}
198200

199-
envVars map[string]string
201+
environmentVariables map[string]string
200202

201203
manifest atomic.Pointer[agentsdk.Manifest] // manifest is atomic because values can change after reconnection.
202204
reportMetadataInterval time.Duration
@@ -235,14 +237,16 @@ func (a *agent) TailnetConn() *tailnet.Conn {
235237
}
236238

237239
func (a *agent) init(ctx context.Context) {
238-
sshSrv, err := agentssh.NewServer(ctx, a.logger.Named("ssh-server"), a.prometheusRegistry, a.filesystem, a.sshMaxTimeout, "")
240+
sshSrv, err := agentssh.NewServer(ctx, a.logger.Named("ssh-server"), a.prometheusRegistry, a.filesystem, &agentssh.Config{
241+
MaxTimeout: a.sshMaxTimeout,
242+
MOTDFile: func() string { return a.manifest.Load().MOTDFile },
243+
ServiceBanner: func() *codersdk.ServiceBannerConfig { return a.serviceBanner.Load() },
244+
UpdateEnv: a.updateCommandEnv,
245+
WorkingDirectory: func() string { return a.manifest.Load().Directory },
246+
})
239247
if err != nil {
240248
panic(err)
241249
}
242-
sshSrv.Env = a.envVars
243-
sshSrv.AgentToken = func() string { return *a.sessionToken.Load() }
244-
sshSrv.Manifest = &a.manifest
245-
sshSrv.ServiceBanner = &a.serviceBanner
246250
a.sshServer = sshSrv
247251
a.scriptRunner = agentscripts.New(agentscripts.Options{
248252
LogDir: a.logDir,
@@ -879,6 +883,83 @@ func (a *agent) run(ctx context.Context) error {
879883
return eg.Wait()
880884
}
881885

886+
// updateCommandEnv updates the provided command environment with the
887+
// following set of environment variables:
888+
// - Predefined workspace environment variables
889+
// - Environment variables currently set (overriding predefined)
890+
// - Environment variables passed via the agent manifest (overriding predefined and current)
891+
// - Agent-level environment variables (overriding all)
892+
func (a *agent) updateCommandEnv(current []string) (updated []string, err error) {
893+
manifest := a.manifest.Load()
894+
if manifest == nil {
895+
return nil, xerrors.Errorf("no manifest")
896+
}
897+
898+
executablePath, err := os.Executable()
899+
if err != nil {
900+
return nil, xerrors.Errorf("getting os executable: %w", err)
901+
}
902+
unixExecutablePath := strings.ReplaceAll(executablePath, "\\", "/")
903+
904+
// Define environment variables that should be set for all commands,
905+
// and then merge them with the current environment.
906+
envs := map[string]string{
907+
// Set env vars indicating we're inside a Coder workspace.
908+
"CODER": "true",
909+
"CODER_WORKSPACE_NAME": manifest.WorkspaceName,
910+
"CODER_WORKSPACE_AGENT_NAME": manifest.AgentName,
911+
912+
// Specific Coder subcommands require the agent token exposed!
913+
"CODER_AGENT_TOKEN": *a.sessionToken.Load(),
914+
915+
// Git on Windows resolves with UNIX-style paths.
916+
// If using backslashes, it's unable to find the executable.
917+
"GIT_SSH_COMMAND": fmt.Sprintf("%s gitssh --", unixExecutablePath),
918+
// Hide Coder message on code-server's "Getting Started" page
919+
"CS_DISABLE_GETTING_STARTED_OVERRIDE": "true",
920+
}
921+
922+
// This adds the ports dialog to code-server that enables
923+
// proxying a port dynamically.
924+
// If this is empty string, do not set anything. Code-server auto defaults
925+
// using its basepath to construct a path based port proxy.
926+
if manifest.VSCodePortProxyURI != "" {
927+
envs["VSCODE_PROXY_URI"] = manifest.VSCodePortProxyURI
928+
}
929+
930+
// Allow any of the current env to override what we defined above.
931+
for _, env := range current {
932+
parts := strings.SplitN(env, "=", 2)
933+
if len(parts) != 2 {
934+
continue
935+
}
936+
if _, ok := envs[parts[0]]; !ok {
937+
envs[parts[0]] = parts[1]
938+
}
939+
}
940+
941+
// Load environment variables passed via the agent manifest.
942+
// These override all variables we manually specify.
943+
for k, v := range manifest.EnvironmentVariables {
944+
// Expanding environment variables allows for customization
945+
// of the $PATH, among other variables. Customers can prepend
946+
// or append to the $PATH, so allowing expand is required!
947+
envs[k] = os.ExpandEnv(v)
948+
}
949+
950+
// Agent-level environment variables should take over all. This is
951+
// used for setting agent-specific variables like CODER_AGENT_TOKEN
952+
// and GIT_ASKPASS.
953+
for k, v := range a.environmentVariables {
954+
envs[k] = v
955+
}
956+
957+
for k, v := range envs {
958+
updated = append(updated, fmt.Sprintf("%s=%s", k, v))
959+
}
960+
return updated, nil
961+
}
962+
882963
func (a *agent) wireguardAddresses(agentID uuid.UUID) []netip.Prefix {
883964
if len(a.addresses) == 0 {
884965
return []netip.Prefix{
@@ -1314,7 +1395,7 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
13141395
}
13151396
}()
13161397

1317-
if val := a.envVars[EnvProcPrioMgmt]; val == "" || runtime.GOOS != "linux" {
1398+
if val := a.environmentVariables[EnvProcPrioMgmt]; val == "" || runtime.GOOS != "linux" {
13181399
a.logger.Debug(ctx, "process priority not enabled, agent will not manage process niceness/oom_score_adj ",
13191400
slog.F("env_var", EnvProcPrioMgmt),
13201401
slog.F("value", val),

‎agent/agent_test.go

Lines changed: 83 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"bytes"
66
"context"
77
"encoding/json"
8+
"errors"
89
"fmt"
910
"io"
1011
"math/rand"
@@ -281,6 +282,83 @@ func TestAgent_SessionExec(t *testing.T) {
281282
require.Equal(t, "test", strings.TrimSpace(string(output)))
282283
}
283284

285+
//nolint:tparallel // Sub tests need to run sequentially.
286+
func TestAgent_Session_EnvironmentVariables(t *testing.T) {
287+
t.Parallel()
288+
289+
manifest := agentsdk.Manifest{
290+
EnvironmentVariables: map[string]string{
291+
"MY_MANIFEST": "true",
292+
"MY_OVERRIDE": "false",
293+
"MY_SESSION_MANIFEST": "false",
294+
},
295+
}
296+
banner := codersdk.ServiceBannerConfig{}
297+
session := setupSSHSession(t, manifest, banner, nil, func(_ *agenttest.Client, opts *agent.Options) {
298+
opts.EnvironmentVariables["MY_OVERRIDE"] = "true"
299+
})
300+
301+
err := session.Setenv("MY_SESSION_MANIFEST", "true")
302+
require.NoError(t, err)
303+
err = session.Setenv("MY_SESSION", "true")
304+
require.NoError(t, err)
305+
306+
command := "sh"
307+
echoEnv := func(t *testing.T, w io.Writer, env string) {
308+
if runtime.GOOS == "windows" {
309+
_, err := fmt.Fprintf(w, "echo %%%s%%\r\n", env)
310+
require.NoError(t, err)
311+
} else {
312+
_, err := fmt.Fprintf(w, "echo $%s\n", env)
313+
require.NoError(t, err)
314+
}
315+
}
316+
if runtime.GOOS == "windows" {
317+
command = "cmd.exe"
318+
}
319+
stdin, err := session.StdinPipe()
320+
require.NoError(t, err)
321+
defer stdin.Close()
322+
stdout, err := session.StdoutPipe()
323+
require.NoError(t, err)
324+
325+
err = session.Start(command)
326+
require.NoError(t, err)
327+
328+
// Context is fine here since we're not doing a parallel subtest.
329+
ctx := testutil.Context(t, testutil.WaitLong)
330+
go func() {
331+
<-ctx.Done()
332+
_ = session.Close()
333+
}()
334+
335+
s := bufio.NewScanner(stdout)
336+
337+
//nolint:paralleltest // These tests need to run sequentially.
338+
for k, partialV := range map[string]string{
339+
"CODER": "true", // From the agent.
340+
"MY_MANIFEST": "true", // From the manifest.
341+
"MY_OVERRIDE": "true", // From the agent environment variables option, overrides manifest.
342+
"MY_SESSION_MANIFEST": "false", // From the manifest, overrides session env.
343+
"MY_SESSION": "true", // From the session.
344+
} {
345+
t.Run(k, func(t *testing.T) {
346+
echoEnv(t, stdin, k)
347+
// Windows is unreliable, so keep scanning until we find a match.
348+
for s.Scan() {
349+
got := strings.TrimSpace(s.Text())
350+
t.Logf("%s=%s", k, got)
351+
if strings.Contains(got, partialV) {
352+
break
353+
}
354+
}
355+
if err := s.Err(); !errors.Is(err, io.EOF) {
356+
require.NoError(t, err)
357+
}
358+
})
359+
}
360+
}
361+
284362
func TestAgent_GitSSH(t *testing.T) {
285363
t.Parallel()
286364
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)
@@ -1991,15 +2069,17 @@ func setupSSHSession(
19912069
manifest agentsdk.Manifest,
19922070
serviceBanner codersdk.ServiceBannerConfig,
19932071
prepareFS func(fs afero.Fs),
2072+
opts ...func(*agenttest.Client, *agent.Options),
19942073
) *ssh.Session {
19952074
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
19962075
defer cancel()
1997-
//nolint:dogsled
1998-
conn, _, _, fs, _ := setupAgent(t, manifest, 0, func(c *agenttest.Client, _ *agent.Options) {
2076+
opts = append(opts, func(c *agenttest.Client, o *agent.Options) {
19992077
c.SetServiceBannerFunc(func() (codersdk.ServiceBannerConfig, error) {
20002078
return serviceBanner, nil
20012079
})
20022080
})
2081+
//nolint:dogsled
2082+
conn, _, _, fs, _ := setupAgent(t, manifest, 0, opts...)
20032083
if prepareFS != nil {
20042084
prepareFS(fs)
20052085
}
@@ -2057,6 +2137,7 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
20572137
Filesystem: fs,
20582138
Logger: logger.Named("agent"),
20592139
ReconnectingPTYTimeout: ptyTimeout,
2140+
EnvironmentVariables: map[string]string{},
20602141
}
20612142

20622143
for _, opt := range opts {

‎agent/agentscripts/agentscripts_test.go

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import (
88
"github.com/prometheus/client_golang/prometheus"
99
"github.com/spf13/afero"
1010
"github.com/stretchr/testify/require"
11-
"go.uber.org/atomic"
1211
"go.uber.org/goleak"
1312

1413
"cdr.dev/slog/sloggers/slogtest"
@@ -72,10 +71,8 @@ func setup(t *testing.T, patchLogs func(ctx context.Context, req agentsdk.PatchL
7271
}
7372
fs := afero.NewMemMapFs()
7473
logger := slogtest.Make(t, nil)
75-
s, err := agentssh.NewServer(context.Background(), logger, prometheus.NewRegistry(), fs, 0, "")
74+
s, err := agentssh.NewServer(context.Background(), logger, prometheus.NewRegistry(), fs, nil)
7675
require.NoError(t, err)
77-
s.AgentToken = func() string { return "" }
78-
s.Manifest = atomic.NewPointer(&agentsdk.Manifest{})
7976
t.Cleanup(func() {
8077
_ = s.Close()
8178
})

0 commit comments

Comments
 (0)