feat: add option to speedtest to dump a pcap of network traffic (#11848)

coadler · web-flow · commit bc14e926d8eb · 2024-01-29T09:57:31.000-06:00
diff --git a/cli/speedtest.go b/cli/speedtest.go
@@ -3,11 +3,13 @@ package cli
 import (
 	"context"
 	"fmt"
+	"os"
 	"time"
 
 	"github.com/jedib0t/go-pretty/v6/table"
 	"golang.org/x/xerrors"
 	tsspeedtest "tailscale.com/net/speedtest"
+	"tailscale.com/wgengine/capture"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
@@ -21,6 +23,7 @@ func (r *RootCmd) speedtest() *clibase.Cmd {
 		direct    bool
 		duration  time.Duration
 		direction string
+		pcapFile  string
 	)
 	client := new(codersdk.Client)
 	cmd := &clibase.Cmd{
@@ -63,6 +66,7 @@ func (r *RootCmd) speedtest() *clibase.Cmd {
 				return err
 			}
 			defer conn.Close()
+
 			if direct {
 				ticker := time.NewTicker(time.Second)
 				defer ticker.Stop()
@@ -95,6 +99,19 @@ func (r *RootCmd) speedtest() *clibase.Cmd {
 			} else {
 				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":
@@ -146,6 +163,12 @@ func (r *RootCmd) speedtest() *clibase.Cmd {
 			Default:       tsspeedtest.DefaultDuration.String(),
 			Value:         clibase.DurationOf(&duration),
 		},
+		{
+			Description: "Specifies a file to write a network capture to.",
+			Flag:        "pcap-file",
+			Default:     "",
+			Value:       clibase.StringOf(&pcapFile),
+		},
 	}
 	return cmd
 }
diff --git a/cli/testdata/coder_speedtest_--help.golden b/cli/testdata/coder_speedtest_--help.golden
@@ -14,6 +14,9 @@ OPTIONS:
           Specifies whether to run in reverse mode where the client receives and
           the server sends.
 
+      --pcap-file string
+          Specifies a file to write a network capture to.
+
   -t, --time duration (default: 5s)
           Specifies the duration to monitor traffic.
 
diff --git a/docs/cli/speedtest.md b/docs/cli/speedtest.md
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -30,6 +30,7 @@ import (
 	tslogger "tailscale.com/types/logger"
 	"tailscale.com/types/netlogtype"
 	"tailscale.com/wgengine"
+	"tailscale.com/wgengine/capture"
 	"tailscale.com/wgengine/magicsock"
 	"tailscale.com/wgengine/netstack"
 	"tailscale.com/wgengine/router"
@@ -310,6 +311,12 @@ type Conn struct {
 	trafficStats *connstats.Statistics
 }
 
+func (c *Conn) InstallCaptureHook(f capture.Callback) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+	c.wireguardEngine.InstallCaptureHook(f)
+}
+
 func (c *Conn) MagicsockSetDebugLoggingEnabled(enabled bool) {
 	c.magicConn.SetDebugLoggingEnabled(enabled)
 }
