@@ -117,7 +117,7 @@ func (s *scaletestTracingFlags) provider(ctx context.Context) (trace.TracerProvi
117
117
}
118
118
119
119
var closeTracingOnce sync.Once
120
- return tracerProvider , func (ctx context.Context ) error {
120
+ return tracerProvider , func (_ context.Context ) error {
121
121
var err error
122
122
closeTracingOnce .Do (func () {
123
123
// Allow time to upload traces even if ctx is canceled
@@ -430,7 +430,7 @@ func (r *RootCmd) scaletestCleanup() *serpent.Command {
430
430
}
431
431
432
432
cliui .Infof (inv .Stdout , "Fetching scaletest workspaces..." )
433
- workspaces , err := getScaletestWorkspaces (ctx , client , template )
433
+ workspaces , _ , err := getScaletestWorkspaces (ctx , client , "" , template )
434
434
if err != nil {
435
435
return err
436
436
}
@@ -863,6 +863,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
863
863
tickInterval time.Duration
864
864
bytesPerTick int64
865
865
ssh bool
866
+ useHostLogin bool
866
867
app string
867
868
template string
868
869
targetWorkspaces string
@@ -926,10 +927,18 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
926
927
return xerrors .Errorf ("get app host: %w" , err )
927
928
}
928
929
929
- workspaces , err := getScaletestWorkspaces (inv .Context (), client , template )
930
+ var owner string
931
+ if useHostLogin {
932
+ owner = codersdk .Me
933
+ }
934
+
935
+ workspaces , numSkipped , err := getScaletestWorkspaces (inv .Context (), client , owner , template )
930
936
if err != nil {
931
937
return err
932
938
}
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 \t Set --use-host-login to only target workspaces you own." , numSkipped )
941
+ }
933
942
934
943
if targetWorkspaceEnd == 0 {
935
944
targetWorkspaceEnd = len (workspaces )
@@ -1092,6 +1101,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
1092
1101
Description : "Send WebSocket traffic to a workspace app (proxied via coderd), cannot be used with --ssh." ,
1093
1102
Value : serpent .StringOf (& app ),
1094
1103
},
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 (& useHostLogin ),
1110
+ },
1095
1111
}
1096
1112
1097
1113
tracingFlags .attach (& cmd .Options )
@@ -1378,22 +1394,35 @@ func isScaleTestWorkspace(workspace codersdk.Workspace) bool {
1378
1394
strings .HasPrefix (workspace .Name , "scaletest-" )
1379
1395
}
1380
1396
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 ) {
1382
1398
var (
1383
1399
pageNumber = 0
1384
1400
limit = 100
1385
1401
workspaces []codersdk.Workspace
1402
+ skipped int
1386
1403
)
1387
1404
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
+
1388
1416
for {
1389
1417
page , err := client .Workspaces (ctx , codersdk.WorkspaceFilter {
1390
1418
Name : "scaletest-" ,
1391
1419
Template : template ,
1420
+ Owner : owner ,
1392
1421
Offset : pageNumber * limit ,
1393
1422
Limit : limit ,
1394
1423
})
1395
1424
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 )
1397
1426
}
1398
1427
1399
1428
pageNumber ++
@@ -1403,13 +1432,18 @@ func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client, templa
1403
1432
1404
1433
pageWorkspaces := make ([]codersdk.Workspace , 0 , len (page .Workspaces ))
1405
1434
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
1408
1441
}
1442
+ pageWorkspaces = append (pageWorkspaces , w )
1409
1443
}
1410
1444
workspaces = append (workspaces , pageWorkspaces ... )
1411
1445
}
1412
- return workspaces , nil
1446
+ return workspaces , skipped , nil
1413
1447
}
1414
1448
1415
1449
func getScaletestUsers (ctx context.Context , client * codersdk.Client ) ([]codersdk.User , error ) {
0 commit comments