fix: start packet capture immediately on speedtest

This commit is contained in:
Spike Curtis 2024-05-02 13:05:07 +04:00
parent 93d8812284
commit 76f3743a05
3 changed files with 24 additions and 15 deletions

View File

@ -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":

View File

@ -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)

View File

@ -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