From 55e1ec90cbaba612c869d668971ee258141fa2d5 Mon Sep 17 00:00:00 2001
From: Garrett Delfosse <garrett@coder.com>
Date: Wed, 18 Dec 2024 16:29:52 +0000
Subject: [PATCH] feat: add workspace-proxy-url flag to scaletest
 workspace-traffic

---
 cli/exp_scaletest.go                 | 44 +++++++++++++++++++++++-----
 scaletest/workspacetraffic/config.go |  4 +++
 scaletest/workspacetraffic/run.go    | 19 ++++++++----
 3 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go
index e8cdff5a75129..a7bd0f396b5aa 100644
--- a/cli/exp_scaletest.go
+++ b/cli/exp_scaletest.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"math/rand"
 	"net/http"
+	"net/url"
 	"os"
 	"os/signal"
 	"strconv"
@@ -860,13 +861,14 @@ func (r *RootCmd) scaletestCreateWorkspaces() *serpent.Command {
 
 func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 	var (
-		tickInterval     time.Duration
-		bytesPerTick     int64
-		ssh              bool
-		useHostLogin     bool
-		app              string
-		template         string
-		targetWorkspaces string
+		tickInterval      time.Duration
+		bytesPerTick      int64
+		ssh               bool
+		useHostLogin      bool
+		app               string
+		template          string
+		targetWorkspaces  string
+		workspaceProxyURL string
 
 		client          = &codersdk.Client{}
 		tracingFlags    = &scaletestTracingFlags{}
@@ -1002,6 +1004,23 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 					return xerrors.Errorf("configure workspace app: %w", err)
 				}
 
+				var webClient *codersdk.Client
+				if workspaceProxyURL != "" {
+					u, err := url.Parse(workspaceProxyURL)
+					if err != nil {
+						return xerrors.Errorf("parse workspace proxy URL: %w", err)
+					}
+
+					webClient = codersdk.New(u)
+					webClient.HTTPClient = client.HTTPClient
+					webClient.SetSessionToken(client.SessionToken())
+
+					appConfig, err = createWorkspaceAppConfig(webClient, appHost.Host, app, ws, agent)
+					if err != nil {
+						return xerrors.Errorf("configure proxy workspace app: %w", err)
+					}
+				}
+
 				// Setup our workspace agent connection.
 				config := workspacetraffic.Config{
 					AgentID:      agent.ID,
@@ -1015,6 +1034,10 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 					App:          appConfig,
 				}
 
+				if webClient != nil {
+					config.WebClient = webClient
+				}
+
 				if err := config.Validate(); err != nil {
 					return xerrors.Errorf("validate config: %w", err)
 				}
@@ -1108,6 +1131,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 			Description: "Connect as the currently logged in user.",
 			Value:       serpent.BoolOf(&useHostLogin),
 		},
+		{
+			Flag:        "workspace-proxy-url",
+			Env:         "CODER_SCALETEST_WORKSPACE_PROXY_URL",
+			Default:     "",
+			Description: "URL for workspace proxy to send web traffic to.",
+			Value:       serpent.StringOf(&workspaceProxyURL),
+		},
 	}
 
 	tracingFlags.attach(&cmd.Options)
diff --git a/scaletest/workspacetraffic/config.go b/scaletest/workspacetraffic/config.go
index 71134a454a411..6ef0760ff3013 100644
--- a/scaletest/workspacetraffic/config.go
+++ b/scaletest/workspacetraffic/config.go
@@ -5,6 +5,8 @@ import (
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/codersdk"
 )
 
 type Config struct {
@@ -33,6 +35,8 @@ type Config struct {
 	Echo bool `json:"echo"`
 
 	App AppConfig `json:"app"`
+
+	WebClient *codersdk.Client
 }
 
 func (c Config) Validate() error {
diff --git a/scaletest/workspacetraffic/run.go b/scaletest/workspacetraffic/run.go
index c683536461bbc..7d541b05b67b9 100644
--- a/scaletest/workspacetraffic/run.go
+++ b/scaletest/workspacetraffic/run.go
@@ -23,8 +23,9 @@ import (
 )
 
 type Runner struct {
-	client *codersdk.Client
-	cfg    Config
+	client    *codersdk.Client
+	webClient *codersdk.Client
+	cfg       Config
 }
 
 var (
@@ -34,9 +35,15 @@ var (
 
 // func NewRunner(client *codersdk.Client, cfg Config, metrics *Metrics) *Runner {
 func NewRunner(client *codersdk.Client, cfg Config) *Runner {
+	webClient := client
+	if cfg.WebClient != nil {
+		webClient = cfg.WebClient
+	}
+
 	return &Runner{
-		client: client,
-		cfg:    cfg,
+		client:    client,
+		webClient: webClient,
+		cfg:       cfg,
 	}
 }
 
@@ -94,7 +101,7 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
 	switch {
 	case r.cfg.App.Name != "":
 		logger.Info(ctx, "sending traffic to workspace app", slog.F("app", r.cfg.App.Name))
-		conn, err = appClientConn(ctx, r.client, r.cfg.App.URL)
+		conn, err = appClientConn(ctx, r.webClient, r.cfg.App.URL)
 		if err != nil {
 			logger.Error(ctx, "connect to workspace app", slog.Error(err))
 			return xerrors.Errorf("connect to workspace app: %w", err)
@@ -113,7 +120,7 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
 
 	default:
 		logger.Info(ctx, "connecting to workspace agent", slog.F("method", "reconnectingpty"))
-		conn, err = connectRPTY(ctx, r.client, agentID, reconnect, command)
+		conn, err = connectRPTY(ctx, r.webClient, agentID, reconnect, command)
 		if err != nil {
 			logger.Error(ctx, "connect to workspace agent via reconnectingpty", slog.Error(err))
 			return xerrors.Errorf("connect to workspace via reconnectingpty: %w", err)