From 3de737fdc8b0383e9001f03bd05b741e993483ce Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Thu, 2 May 2024 19:44:32 +0400
Subject: [PATCH] fix: start packet capture immediately on speedtest (#13128)

I initially made this change when hacking wgengine to also capture wireguard packets going into the magicsock, so that we could capture the initial wireguard handshake.

I don't think we should ship that additional capture logic, but... it seems generally useful to capture packets from the get go on speedtest, so that you can see disco and pings before the TCP speedtest session starts.
---
 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 f4053d72c6..e88872019e 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 8af29f8f76..f1e3bd67ea 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 d4d58c7cc9..1f4c543c93 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