fix(support): also sanitize agent environment (#12615)

johnstcn · web-flow · commit 9ff0bafcee31 · 2024-03-15T20:19:35.000Z
diff --git a/support/support.go b/support/support.go
@@ -298,6 +298,7 @@ func AgentInfo(ctx context.Context, client *codersdk.Client, log slog.Logger, ag
 		if err != nil {
 			return xerrors.Errorf("fetch workspace agent: %w", err)
 		}
+		sanitizeEnv(agt.EnvironmentVariables)
 		a.Agent = &agt
 		return nil
 	})
diff --git a/support/support_test.go b/support/support_test.go
@@ -73,6 +73,7 @@ func TestRun(t *testing.T) {
 		assertNotNilNotEmpty(t, bun.Workspace.TemplateFileBase64, "workspace template file should be present")
 		require.NotNil(t, bun.Workspace.Parameters, "workspace parameters should be present")
 		assertNotNilNotEmpty(t, bun.Agent.Agent, "agent should be present")
+		assertSanitizedAgent(t, *bun.Agent.Agent)
 		assertNotNilNotEmpty(t, bun.Agent.ListeningPorts, "agent listening ports should be present")
 		assertNotNilNotEmpty(t, bun.Agent.Logs, "agent logs should be present")
 		assertNotNilNotEmpty(t, bun.Agent.AgentMagicsockHTML, "agent magicsock should be present")
@@ -163,13 +164,18 @@ func assertSanitizedWorkspace(t *testing.T, ws codersdk.Workspace) {
 	t.Helper()
 	for _, res := range ws.LatestBuild.Resources {
 		for _, agt := range res.Agents {
-			for k, v := range agt.EnvironmentVariables {
-				assert.Equal(t, "***REDACTED***", v, "environment variable %q not sanitized", k)
-			}
+			assertSanitizedAgent(t, agt)
 		}
 	}
 }
 
+func assertSanitizedAgent(t *testing.T, agt codersdk.WorkspaceAgent) {
+	t.Helper()
+	for k, v := range agt.EnvironmentVariables {
+		assert.Equal(t, "***REDACTED***", v, "agent %q environment variable %q not sanitized", agt.Name, k)
+	}
+}
+
 func setupWorkspaceAndAgent(ctx context.Context, t *testing.T, client *codersdk.Client, db database.Store, user codersdk.CreateFirstUserResponse) (codersdk.Workspace, codersdk.WorkspaceAgent) {
 	// This is a valid zip file
 	zipBytes := make([]byte, 22)

Original file line number	Diff line number	Diff line change
`@@ -298,6 +298,7 @@ func AgentInfo(ctx context.Context, client *codersdk.Client, log slog.Logger, ag`
`298`	`298`	`if err != nil {`
`299`	`299`	`return xerrors.Errorf("fetch workspace agent: %w", err)`
`300`	`300`	`}`
	`301`	`+ sanitizeEnv(agt.EnvironmentVariables)`
`301`	`302`	`a.Agent = &agt`
`302`	`303`	`return nil`
`303`	`304`	`})`