fix: start packet capture immediately on speedtest

spikecurtis · spikecurtis · commit 76f3743a05af · 2024-05-02T13:05:07.000+04:00
diff --git 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
@@ -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
@@ -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
