Skip to content

Commit 61ed855

Browse files
johnstcnjohnstcn
committed
fix(cli): scaletest: check for DisableOwnerWorkspaceExec
1 parent d4adfa3 commit 61ed855

File tree

2 files changed

+97
-8
lines changed

2 files changed

+97
-8
lines changed

‎cli/exp_scaletest.go

Lines changed: 42 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ func (s *scaletestTracingFlags) provider(ctx context.Context) (trace.TracerProvi
117117
}
118118

119119
var closeTracingOnce sync.Once
120-
return tracerProvider, func(ctx context.Context) error {
120+
return tracerProvider, func(_ context.Context) error {
121121
var err error
122122
closeTracingOnce.Do(func() {
123123
// Allow time to upload traces even if ctx is canceled
@@ -430,7 +430,7 @@ func (r *RootCmd) scaletestCleanup() *serpent.Command {
430430
}
431431

432432
cliui.Infof(inv.Stdout, "Fetching scaletest workspaces...")
433-
workspaces, err := getScaletestWorkspaces(ctx, client, template)
433+
workspaces, _, err := getScaletestWorkspaces(ctx, client, "", template)
434434
if err != nil {
435435
return err
436436
}
@@ -863,6 +863,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
863863
tickInterval time.Duration
864864
bytesPerTick int64
865865
ssh bool
866+
useHostUser bool
866867
app string
867868
template string
868869
targetWorkspaces string
@@ -926,10 +927,18 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
926927
return xerrors.Errorf("get app host: %w", err)
927928
}
928929

929-
workspaces, err := getScaletestWorkspaces(inv.Context(), client, template)
930+
var owner string
931+
if useHostUser {
932+
owner = codersdk.Me
933+
}
934+
935+
workspaces, numSkipped, err := getScaletestWorkspaces(inv.Context(), client, owner, template)
930936
if err != nil {
931937
return err
932938
}
939+
if numSkipped > 0 {
940+
cliui.Warnf(inv.Stdout, "CODER_DISABLE_OWNER_WORKSPACE_ACCESS is set on the deployment.\n\t%d workspace(s) were skipped due to ownership mismatch.\n\tSet --use-host-login to only target workspaces you own.", numSkipped)
941+
}
933942

934943
if targetWorkspaceEnd == 0 {
935944
targetWorkspaceEnd = len(workspaces)
@@ -1092,6 +1101,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
10921101
Description: "Send WebSocket traffic to a workspace app (proxied via coderd), cannot be used with --ssh.",
10931102
Value: serpent.StringOf(&app),
10941103
},
1104+
{
1105+
Flag: "use-host-login",
1106+
Env: "CODER_SCALETEST_USE_HOST_LOGIN",
1107+
Default: "false",
1108+
Description: "Connect as the currently logged in user.",
1109+
Value: serpent.BoolOf(&useHostUser),
1110+
},
10951111
}
10961112

10971113
tracingFlags.attach(&cmd.Options)
@@ -1378,22 +1394,35 @@ func isScaleTestWorkspace(workspace codersdk.Workspace) bool {
13781394
strings.HasPrefix(workspace.Name, "scaletest-")
13791395
}
13801396

1381-
func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client, template string) ([]codersdk.Workspace, error) {
1397+
func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client, owner, template string) ([]codersdk.Workspace, int, error) {
13821398
var (
13831399
pageNumber = 0
13841400
limit = 100
13851401
workspaces []codersdk.Workspace
1402+
skipped int
13861403
)
13871404

1405+
me, err := client.User(ctx, codersdk.Me)
1406+
if err != nil {
1407+
return nil, 0, xerrors.Errorf("check logged-in user")
1408+
}
1409+
1410+
dv, err := client.DeploymentConfig(ctx)
1411+
if err != nil {
1412+
return nil, 0, xerrors.Errorf("fetch deployment config: %w", err)
1413+
}
1414+
noOwnerAccess := dv.Values != nil && dv.Values.DisableOwnerWorkspaceExec.Value()
1415+
13881416
for {
13891417
page, err := client.Workspaces(ctx, codersdk.WorkspaceFilter{
13901418
Name: "scaletest-",
13911419
Template: template,
1420+
Owner: owner,
13921421
Offset: pageNumber * limit,
13931422
Limit: limit,
13941423
})
13951424
if err != nil {
1396-
return nil, xerrors.Errorf("fetch scaletest workspaces page %d: %w", pageNumber, err)
1425+
return nil, 0, xerrors.Errorf("fetch scaletest workspaces page %d: %w", pageNumber, err)
13971426
}
13981427

13991428
pageNumber++
@@ -1403,13 +1432,18 @@ func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client, templa
14031432

14041433
pageWorkspaces := make([]codersdk.Workspace, 0, len(page.Workspaces))
14051434
for _, w := range page.Workspaces {
1406-
if isScaleTestWorkspace(w) {
1407-
pageWorkspaces = append(pageWorkspaces, w)
1435+
if !isScaleTestWorkspace(w) {
1436+
continue
1437+
}
1438+
if noOwnerAccess && w.OwnerID != me.ID {
1439+
skipped++
1440+
continue
14081441
}
1442+
pageWorkspaces = append(pageWorkspaces, w)
14091443
}
14101444
workspaces = append(workspaces, pageWorkspaces...)
14111445
}
1412-
return workspaces, nil
1446+
return workspaces, skipped, nil
14131447
}
14141448

14151449
func getScaletestUsers(ctx context.Context, client *codersdk.Client) ([]codersdk.User, error) {

‎cli/exp_scaletest_test.go

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cli_test
22

33
import (
4+
"bytes"
45
"context"
56
"path/filepath"
67
"testing"
@@ -11,6 +12,7 @@ import (
1112

1213
"github.com/coder/coder/v2/cli/clitest"
1314
"github.com/coder/coder/v2/coderd/coderdtest"
15+
"github.com/coder/coder/v2/codersdk"
1416
"github.com/coder/coder/v2/pty/ptytest"
1517
"github.com/coder/coder/v2/testutil"
1618
)
@@ -116,6 +118,59 @@ func TestScaleTestWorkspaceTraffic_Template(t *testing.T) {
116118
require.ErrorContains(t, err, "could not find template \"doesnotexist\" in any organization")
117119
}
118120

121+
// This test validates that the scaletest CLI filters out workspaces not owned
122+
// when disable owner workspace access is set.
123+
// nolint:paralleltest // DisableOwnerWorkspaceExec is not safe for parallel tests.
124+
func TestScaleTestWorkspaceTraffic_UseHostLogin(t *testing.T) {
125+
ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitMedium)
126+
defer cancelFunc()
127+
128+
log := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
129+
client := coderdtest.New(t, &coderdtest.Options{
130+
Logger: &log,
131+
IncludeProvisionerDaemon: true,
132+
DeploymentValues: coderdtest.DeploymentValues(t, func(dv *codersdk.DeploymentValues) {
133+
dv.DisableOwnerWorkspaceExec = true
134+
}),
135+
})
136+
owner := coderdtest.CreateFirstUser(t, client)
137+
tv := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, nil)
138+
_ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, tv.ID)
139+
tpl := coderdtest.CreateTemplate(t, client, owner.OrganizationID, tv.ID)
140+
// Create a workspace owned by a different user
141+
memberClient, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
142+
_ = coderdtest.CreateWorkspace(t, memberClient, tpl.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
143+
cwr.Name = "scaletest-workspace"
144+
})
145+
146+
// Test without --use-host-login first.
147+
inv, root := clitest.New(t, "exp", "scaletest", "workspace-traffic",
148+
"--template", tpl.Name,
149+
)
150+
// nolint:gocritic // We are intentionally testing this as the owner.
151+
clitest.SetupConfig(t, client, root)
152+
var stdoutBuf bytes.Buffer
153+
inv.Stdout = &stdoutBuf
154+
155+
err := inv.WithContext(ctx).Run()
156+
require.ErrorContains(t, err, "no scaletest workspaces exist")
157+
require.Contains(t, stdoutBuf.String(), `1 workspace(s) were skipped`)
158+
159+
// Test once again with --use-host-login.
160+
inv, root = clitest.New(t, "exp", "scaletest", "workspace-traffic",
161+
"--template", tpl.Name,
162+
"--use-host-login",
163+
)
164+
// nolint:gocritic // We are intentionally testing this as the owner.
165+
clitest.SetupConfig(t, client, root)
166+
stdoutBuf.Reset()
167+
inv.Stdout = &stdoutBuf
168+
169+
err = inv.WithContext(ctx).Run()
170+
require.ErrorContains(t, err, "no scaletest workspaces exist")
171+
require.NotContains(t, stdoutBuf.String(), `1 workspace(s) were skipped`)
172+
}
173+
119174
// This test just validates that the CLI command accepts its known arguments.
120175
func TestScaleTestWorkspaceTraffic_TargetWorkspaces(t *testing.T) {
121176
t.Parallel()

0 commit comments

Comments
 (0)