tstest/natlab/vnet: add pcap support

Updates #13038

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

tstest/integration/nat/nat_test.go

@@ -33,7 +33,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
@@ -142,6 +145,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,21 @@ 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
+		}
+		nw, err := pcapgo.NewNgWriter(pcf, layers.LinkTypeEthernet)
+		if err != nil {
+			return err
+		}
+		pw := &pcapWriter{
+			f: pcf,
+			w: nw,
+		}
+		s.pcapWriter = pw
+	}
 	for _, conf := range c.networks {
 		if conf.err != nil {
 			return conf.err
@@ -206,14 +229,21 @@ 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()],
 		}
+		if s.pcapWriter != nil {
+			s.pcapWriter.w.AddInterface(pcapgo.NgInterface{
+				Name:     fmt.Sprintf("node%d", n.id),
+				LinkType: layers.LinkTypeEthernet,
+			})
+		}
 		conf.n = n
 		if _, ok := s.nodeByMAC[n.mac]; ok {
 			return fmt.Errorf("two nodes have the same MAC %v", n.mac)

tstest/natlab/vnet/pcap.go

@@ -0,0 +1,39 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package vnet
+
+import (
+	"io"
+	"os"
+	"sync"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/pcapgo"
+)
+
+type pcapWriter struct {
+	f *os.File
+
+	mu sync.Mutex
+	w  *pcapgo.NgWriter
+}
+
+func (p *pcapWriter) WritePacket(ci gopacket.CaptureInfo, data []byte) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.w == nil {
+		return io.ErrClosedPipe
+	}
+	return p.w.WritePacket(ci, data)
+}
+
+func (p *pcapWriter) Close() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.w != nil {
+		p.w.Flush()
+		p.w = nil
+	}
+	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
+	control    *testcontrol.Server
-	derps   []*derpServer
+	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 shutdown := s.shuttingDown.Swap(true); !shutdown {
+		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
@@ -686,6 +717,17 @@ func (s *Server) ServeUnixConn(uc *net.UnixConn, proto Protocol) {
 				continue
 			}
 		}
+		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))
+		}
 		netw.HandleEthernetPacket(ep)
 	}
 }
@@ -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
 	}