tstest/natlab/vnet: add pcap support

Signed-off-by: Maisem Ali <maisem@tailscale.com>

tstest/integration/nat/nat_test.go

@@ -30,7 +30,10 @@ import (
 	"tailscale.com/tstest/natlab/vnet"
 )
 
-var logTailscaled = flag.Bool("log-tailscaled", false, "log tailscaled output")
+var (
+	logTailscaled = flag.Bool("log-tailscaled", false, "log tailscaled output")
+	pcapFile      = flag.String("pcap", "", "write pcap to file")
+)
 
 type natTest struct {
 	tb      testing.TB
@@ -139,6 +142,7 @@ func (nt *natTest) runTest(node1, node2 addNodeFunc) pingRoute {
 	t := nt.tb
 
 	var c vnet.Config
+	c.SetPCAPFile(*pcapFile)
 	nodes := []*vnet.Node{
 		node1(&c),
 		node2(&c),

tstest/natlab/vnet/conf.go

@@ -8,8 +8,11 @@ import (
 	"fmt"
 	"log"
 	"net/netip"
+	"os"
 	"slices"
 
+	"github.com/google/gopacket/layers"
+	"github.com/google/gopacket/pcapgo"
 	"tailscale.com/types/logger"
 	"tailscale.com/util/set"
 )
@@ -27,6 +30,11 @@ import (
 type Config struct {
 	nodes    []*Node
 	networks []*Network
+	pcapFile string
+}
+
+func (c *Config) SetPCAPFile(file string) {
+	c.pcapFile = file
 }
 
 func (c *Config) NumNodes() int {
@@ -183,6 +191,18 @@ func (n *Network) AddService(s NetworkService) {
 // there were any configuration issues.
 func (s *Server) initFromConfig(c *Config) error {
 	netOfConf := map[*Network]*network{}
+	if c.pcapFile != "" {
+		pcf, err := os.OpenFile(c.pcapFile, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
+		if err != nil {
+			return err
+		}
+		pw := &pcapWriter{
+			f: pcf,
+			w: pcapgo.NewWriter(pcf),
+		}
+		pw.w.WriteFileHeader(65536, layers.LinkTypeEthernet)
+		s.pcapWriter = pw
+	}
 	for _, conf := range c.networks {
 		if conf.err != nil {
 			return conf.err
@@ -206,12 +226,13 @@ func (s *Server) initFromConfig(c *Config) error {
 		}
 		s.networkByWAN[conf.wanIP] = n
 	}
-	for _, conf := range c.nodes {
+	for i, conf := range c.nodes {
 		if conf.err != nil {
 			return conf.err
 		}
 		n := &node{
 			mac: conf.mac,
+			id:  i + 1,
 			net: netOfConf[conf.Network()],
 		}
 		conf.n = n

tstest/natlab/vnet/pcap.go

@@ -0,0 +1,26 @@
+package vnet
+
+import (
+	"os"
+	"sync"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/pcapgo"
+)
+
+type pcapWriter struct {
+	f *os.File
+
+	mu sync.Mutex
+	w  *pcapgo.Writer
+}
+
+func (p *pcapWriter) WritePacket(ci gopacket.CaptureInfo, data []byte) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.w.WritePacket(ci, data)
+}
+
+func (p *pcapWriter) Close() error {
+	return p.f.Close()
+}

tstest/natlab/vnet/vnet.go

@@ -31,6 +31,7 @@ import (
 	"os/exec"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/google/gopacket"
@@ -58,6 +59,7 @@ import (
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
 	"tailscale.com/util/mak"
+	"tailscale.com/util/must"
 	"tailscale.com/util/set"
 )
 
@@ -444,6 +446,7 @@ func (n *network) MACOfIP(ip netip.Addr) (_ MAC, ok bool) {
 
 type node struct {
 	mac   MAC
+	id    int
 	net   *network
 	lanIP netip.Addr // must be in net.lanIP prefix + unique in net
 }
@@ -474,6 +477,8 @@ func newDERPServer() *derpServer {
 type Server struct {
 	shutdownCtx    context.Context
 	shutdownCancel context.CancelFunc
+	shuttingDown   atomic.Bool
+	wg             sync.WaitGroup
 	blendReality   bool
 
 	derpIPs set.Set[netip.Addr]
@@ -483,8 +488,9 @@ type Server struct {
 	networks     set.Set[*network]
 	networkByWAN map[netip.Addr]*network
 
-	control *testcontrol.Server
-	derps   []*derpServer
+	control    *testcontrol.Server
+	derps      []*derpServer
+	pcapWriter *pcapWriter
 
 	mu              sync.Mutex
 	agentConnWaiter map[*node]chan<- struct{} // signaled after added to set
@@ -562,7 +568,13 @@ func New(c *Config) (*Server, error) {
 }
 
 func (s *Server) Close() {
-	s.shutdownCancel()
+	if !s.shuttingDown.Swap(true) {
+		s.shutdownCancel()
+		if s.pcapWriter != nil {
+			s.pcapWriter.Close()
+		}
+	}
+	s.wg.Wait()
 }
 
 func (s *Server) HWAddr(mac MAC) net.HardwareAddr {
@@ -601,11 +613,20 @@ const (
 
 // serveConn serves a single connection from a client.
 func (s *Server) ServeUnixConn(uc *net.UnixConn, proto Protocol) {
+	if s.shuttingDown.Load() {
+		return
+	}
+	s.wg.Add(1)
+	defer s.wg.Done()
+	context.AfterFunc(s.shutdownCtx, func() {
+		uc.SetDeadline(time.Now())
+	})
 	log.Printf("Got conn %T %p", uc, uc)
 	defer uc.Close()
 
 	bw := bufio.NewWriterSize(uc, 2<<10)
 	var writeMu sync.Mutex
+	var srcNode *node
 	writePkt := func(pkt []byte) {
 		if pkt == nil {
 			return
@@ -626,10 +647,20 @@ func (s *Server) ServeUnixConn(uc *net.UnixConn, proto Protocol) {
 		if err := bw.Flush(); err != nil {
 			log.Printf("Flush: %v", err)
 		}
+		if s.pcapWriter != nil {
+			ci := gopacket.CaptureInfo{
+				Timestamp:     time.Now(),
+				CaptureLength: len(pkt),
+				Length:        len(pkt),
+			}
+			if srcNode != nil {
+				ci.InterfaceIndex = srcNode.id
+			}
+			must.Do(s.pcapWriter.WritePacket(ci, pkt))
+		}
 	}
 
 	buf := make([]byte, 16<<10)
-	var srcNode *node
 	var netw *network // non-nil after first packet
 	for {
 		var packetRaw []byte
@@ -661,6 +692,17 @@ func (s *Server) ServeUnixConn(uc *net.UnixConn, proto Protocol) {
 			}
 			packetRaw = buf[4 : 4+n] // raw ethernet frame
 		}
+		if s.pcapWriter != nil {
+			ci := gopacket.CaptureInfo{
+				Timestamp:     time.Now(),
+				CaptureLength: len(packetRaw),
+				Length:        len(packetRaw),
+			}
+			if srcNode != nil {
+				ci.InterfaceIndex = srcNode.id
+			}
+			must.Do(s.pcapWriter.WritePacket(ci, packetRaw))
+		}
 
 		packet := gopacket.NewPacket(packetRaw, layers.LayerTypeEthernet, gopacket.Lazy)
 		le, ok := packet.LinkLayer().(*layers.Ethernet)
@@ -840,7 +882,6 @@ func (n *network) WriteUDPPacketNoNAT(p UDPPacket) {
 // IP may be the router's IP, or an internet (routed) IP.
 func (n *network) HandleEthernetIPv4PacketForRouter(ep EthernetPacket) {
 	packet := ep.gp
-	writePkt := n.writeEth
 
 	v4, ok := packet.Layer(layers.LayerTypeIPv4).(*layers.IPv4)
 	if !ok {
@@ -857,7 +898,7 @@ func (n *network) HandleEthernetIPv4PacketForRouter(ep EthernetPacket) {
 			n.logf("createDHCPResponse: %v", err)
 			return
 		}
-		writePkt(res)
+		n.writeEth(res)
 		return
 	}
 
@@ -874,7 +915,7 @@ func (n *network) HandleEthernetIPv4PacketForRouter(ep EthernetPacket) {
 			n.logf("createDNSResponse: %v", err)
 			return
 		}
-		writePkt(res)
+		n.writeEth(res)
 		return
 	}