From 76f3743a05af89c2f09d1197902cde3bcfb327a8 Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Thu, 2 May 2024 13:05:07 +0400
Subject: [PATCH] fix: start packet capture immediately on speedtest

---
 cli/speedtest.go                      | 30 +++++++++++++--------------
 codersdk/workspacesdk/workspacesdk.go |  5 +++++
 tailnet/conn.go                       |  4 ++++
 3 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/cli/speedtest.go b/cli/speedtest.go
index f4053d72c6283..e88872019e2d2 100644
--- a/cli/speedtest.go
+++ b/cli/speedtest.go
@@ -60,10 +60,22 @@ func (r *RootCmd) speedtest() *serpent.Command {
 			if r.disableDirect {
 				_, _ = fmt.Fprintln(inv.Stderr, "Direct connections disabled.")
 			}
+			opts := &workspacesdk.DialAgentOptions{
+				Logger: logger,
+			}
+			if pcapFile != "" {
+				s := capture.New()
+				opts.CaptureHook = s.LogPacket
+				f, err := os.OpenFile(pcapFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
+				if err != nil {
+					return err
+				}
+				defer f.Close()
+				unregister := s.RegisterOutput(f)
+				defer unregister()
+			}
 			conn, err := workspacesdk.New(client).
-				DialAgent(ctx, workspaceAgent.ID, &workspacesdk.DialAgentOptions{
-					Logger: logger,
-				})
+				DialAgent(ctx, workspaceAgent.ID, opts)
 			if err != nil {
 				return err
 			}
@@ -102,18 +114,6 @@ func (r *RootCmd) speedtest() *serpent.Command {
 				conn.AwaitReachable(ctx)
 			}
 
-			if pcapFile != "" {
-				s := capture.New()
-				conn.InstallCaptureHook(s.LogPacket)
-				f, err := os.OpenFile(pcapFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
-				if err != nil {
-					return err
-				}
-				defer f.Close()
-				unregister := s.RegisterOutput(f)
-				defer unregister()
-			}
-
 			var tsDir tsspeedtest.Direction
 			switch direction {
 			case "up":
diff --git a/codersdk/workspacesdk/workspacesdk.go b/codersdk/workspacesdk/workspacesdk.go
index 8af29f8f76f1f..f1e3bd67ea3dc 100644
--- a/codersdk/workspacesdk/workspacesdk.go
+++ b/codersdk/workspacesdk/workspacesdk.go
@@ -16,6 +16,7 @@ import (
 	"golang.org/x/xerrors"
 	"nhooyr.io/websocket"
 	"tailscale.com/tailcfg"
+	"tailscale.com/wgengine/capture"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/codersdk"
@@ -176,6 +177,9 @@ type DialAgentOptions struct {
 	// BlockEndpoints forced a direct connection through DERP. The Client may
 	// have DisableDirect set which will override this value.
 	BlockEndpoints bool
+	// CaptureHook is a callback that captures Disco packets and packets sent
+	// into the tailnet tunnel.
+	CaptureHook capture.Callback
 }
 
 func (c *Client) DialAgent(dialCtx context.Context, agentID uuid.UUID, options *DialAgentOptions) (agentConn *AgentConn, err error) {
@@ -203,6 +207,7 @@ func (c *Client) DialAgent(dialCtx context.Context, agentID uuid.UUID, options *
 		DERPForceWebSockets: connInfo.DERPForceWebSockets,
 		Logger:              options.Logger,
 		BlockEndpoints:      c.client.DisableDirectConnections || options.BlockEndpoints,
+		CaptureHook:         options.CaptureHook,
 	})
 	if err != nil {
 		return nil, xerrors.Errorf("create tailnet: %w", err)
diff --git a/tailnet/conn.go b/tailnet/conn.go
index d4d58c7cc9231..1f4c543c933a1 100644
--- a/tailnet/conn.go
+++ b/tailnet/conn.go
@@ -93,6 +93,9 @@ type Options struct {
 	BlockEndpoints bool
 	Logger         slog.Logger
 	ListenPort     uint16
+	// CaptureHook is a callback that captures Disco packets and packets sent
+	// into the tailnet tunnel.
+	CaptureHook capture.Callback
 }
 
 // NodeID creates a Tailscale NodeID from the last 8 bytes of a UUID. It ensures
@@ -158,6 +161,7 @@ func NewConn(options *Options) (conn *Conn, err error) {
 			wireguardEngine.Close()
 		}
 	}()
+	wireguardEngine.InstallCaptureHook(options.CaptureHook)
 	dialer.UseNetstackForIP = func(ip netip.Addr) bool {
 		_, ok := wireguardEngine.PeerForIP(ip)
 		return ok