diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go index a844a7e8c6..4580ff3e1b 100644 --- a/cli/exp_scaletest.go +++ b/cli/exp_scaletest.go @@ -864,6 +864,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command { tickInterval time.Duration bytesPerTick int64 ssh bool + disableDirect bool useHostLogin bool app string template string @@ -1023,15 +1024,16 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command { // Setup our workspace agent connection. config := workspacetraffic.Config{ - AgentID: agent.ID, - BytesPerTick: bytesPerTick, - Duration: strategy.timeout, - TickInterval: tickInterval, - ReadMetrics: metrics.ReadMetrics(ws.OwnerName, ws.Name, agent.Name), - WriteMetrics: metrics.WriteMetrics(ws.OwnerName, ws.Name, agent.Name), - SSH: ssh, - Echo: ssh, - App: appConfig, + AgentID: agent.ID, + BytesPerTick: bytesPerTick, + Duration: strategy.timeout, + TickInterval: tickInterval, + ReadMetrics: metrics.ReadMetrics(ws.OwnerName, ws.Name, agent.Name), + WriteMetrics: metrics.WriteMetrics(ws.OwnerName, ws.Name, agent.Name), + SSH: ssh, + DisableDirect: disableDirect, + Echo: ssh, + App: appConfig, } if webClient != nil { @@ -1117,6 +1119,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command { Description: "Send traffic over SSH, cannot be used with --app.", Value: serpent.BoolOf(&ssh), }, + { + Flag: "disable-direct", + Env: "CODER_SCALETEST_WORKSPACE_TRAFFIC_DISABLE_DIRECT_CONNECTIONS", + Default: "false", + Description: "Disable direct connections for SSH traffic to workspaces. Does nothing if `--ssh` is not also set.", + Value: serpent.BoolOf(&disableDirect), + }, { Flag: "app", Env: "CODER_SCALETEST_WORKSPACE_TRAFFIC_APP", diff --git a/scaletest/workspacetraffic/config.go b/scaletest/workspacetraffic/config.go index 6ef0760ff3..0948d35ea7 100644 --- a/scaletest/workspacetraffic/config.go +++ b/scaletest/workspacetraffic/config.go @@ -28,6 +28,9 @@ type Config struct { SSH bool `json:"ssh"` + // Ignored unless SSH is true. + DisableDirect bool `json:"ssh_disable_direct"` + // Echo controls whether the agent should echo the data it receives. // If false, the agent will discard the data. Note that setting this // to true will double the amount of data read from the agent for diff --git a/scaletest/workspacetraffic/conn.go b/scaletest/workspacetraffic/conn.go index 7640203e6c..17cbc7c501 100644 --- a/scaletest/workspacetraffic/conn.go +++ b/scaletest/workspacetraffic/conn.go @@ -144,7 +144,7 @@ func (c *rptyConn) Close() (err error) { } //nolint:revive // Ignore requestPTY control flag. -func connectSSH(ctx context.Context, client *codersdk.Client, agentID uuid.UUID, cmd string, requestPTY bool) (rwc *countReadWriteCloser, err error) { +func connectSSH(ctx context.Context, client *codersdk.Client, agentID uuid.UUID, cmd string, requestPTY bool, blockEndpoints bool) (rwc *countReadWriteCloser, err error) { var closers []func() error defer func() { if err != nil { @@ -156,7 +156,9 @@ func connectSSH(ctx context.Context, client *codersdk.Client, agentID uuid.UUID, } }() - agentConn, err := workspacesdk.New(client).DialAgent(ctx, agentID, &workspacesdk.DialAgentOptions{}) + agentConn, err := workspacesdk.New(client).DialAgent(ctx, agentID, &workspacesdk.DialAgentOptions{ + BlockEndpoints: blockEndpoints, + }) if err != nil { return nil, xerrors.Errorf("dial workspace agent: %w", err) } diff --git a/scaletest/workspacetraffic/run.go b/scaletest/workspacetraffic/run.go index cad6a9d51c..7dd7cb6803 100644 --- a/scaletest/workspacetraffic/run.go +++ b/scaletest/workspacetraffic/run.go @@ -111,7 +111,7 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error) // If echo is enabled, disable PTY to avoid double echo and // reduce CPU usage. requestPTY := !r.cfg.Echo - conn, err = connectSSH(ctx, r.client, agentID, command, requestPTY) + conn, err = connectSSH(ctx, r.client, agentID, command, requestPTY, r.cfg.DisableDirect) if err != nil { logger.Error(ctx, "connect to workspace agent via ssh", slog.Error(err)) return xerrors.Errorf("connect to workspace via ssh: %w", err)