chore: add --disable-direct to coder exp scaletest workspace-traffic --ssh (#19632)

Relates to https://github.com/coder/internal/issues/888

As part of our renewed connection scaletesting efforts, we want to
scaletest coder in a scenario where direct connections aren't available
(relatively common for our customers), and all concurrent connections
are relayed via DERP.

This PR adds a flag, `--disable-direct` that can be included on the
existing`coder exp scaletest workspace-traffic -ssh` to disable direct
connections.

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",

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

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)
 	}

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)