tstest/natlab/vnet: capture wan interfaces too

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

tstest/natlab/vnet/conf.go

@@ -14,6 +14,7 @@ import (
 	"github.com/google/gopacket/layers"
 	"github.com/google/gopacket/pcapgo"
 	"tailscale.com/types/logger"
+	"tailscale.com/util/must"
 	"tailscale.com/util/set"
 )
 
@@ -206,7 +207,7 @@ func (s *Server) initFromConfig(c *Config) error {
 		}
 		s.pcapWriter = pw
 	}
-	for _, conf := range c.networks {
+	for i, conf := range c.networks {
 		if conf.err != nil {
 			return conf.err
 		}
@@ -228,6 +229,10 @@ func (s *Server) initFromConfig(c *Config) error {
 			return fmt.Errorf("two networks have the same WAN IP %v; Anycast not (yet?) supported", conf.wanIP)
 		}
 		s.networkByWAN[conf.wanIP] = n
+		n.interfaceID = must.Get(s.pcapWriter.AddInterface(pcapgo.NgInterface{
+			Name:     fmt.Sprintf("network%d", i+1),
+			LinkType: layers.LinkTypeIPv4,
+		}))
 	}
 	for i, conf := range c.nodes {
 		if conf.err != nil {
@@ -235,15 +240,12 @@ func (s *Server) initFromConfig(c *Config) error {
 		}
 		n := &node{
 			mac: conf.mac,
-			id:  i + 1,
 			net: netOfConf[conf.Network()],
 		}
-		if s.pcapWriter != nil {
+		n.interfaceID = must.Get(s.pcapWriter.AddInterface(pcapgo.NgInterface{
-			s.pcapWriter.w.AddInterface(pcapgo.NgInterface{
+			Name:     fmt.Sprintf("node%d", i+1),
-				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

@@ -17,6 +17,9 @@ type pcapWriter struct {
 }
 
 func (p *pcapWriter) WritePacket(ci gopacket.CaptureInfo, data []byte) error {
+	if p == nil {
+		return nil
+	}
 	p.mu.Lock()
 	defer p.mu.Unlock()
 	if p.w == nil {
@@ -25,7 +28,19 @@ func (p *pcapWriter) WritePacket(ci gopacket.CaptureInfo, data []byte) error {
 	return p.w.WritePacket(ci, data)
 }
 
+func (p *pcapWriter) AddInterface(i pcapgo.NgInterface) (int, error) {
+	if p == nil {
+		return 0, nil
+	}
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.w.AddInterface(i)
+}
+
 func (p *pcapWriter) Close() error {
+	if p == nil {
+		return nil
+	}
 	p.mu.Lock()
 	defer p.mu.Unlock()
 	if p.w != nil {

tstest/natlab/vnet/vnet.go

@@ -407,6 +407,7 @@ type network struct {
 	s           *Server
 	mac         MAC
 	portmap     bool
+	interfaceID int
 	wanIP       netip.Addr
 	lanIP       netip.Prefix // with host bits set (e.g. 192.168.2.1/24)
 	nodesByIP   map[netip.Addr]*node
@@ -446,7 +447,7 @@ func (n *network) MACOfIP(ip netip.Addr) (_ MAC, ok bool) {
 
 type node struct {
 	mac         MAC
-	id    int
+	interfaceID int
 	net         *network
 	lanIP       netip.Addr // must be in net.lanIP prefix + unique in net
 }
@@ -570,10 +571,8 @@ func New(c *Config) (*Server, error) {
 func (s *Server) Close() {
 	if shutdown := s.shuttingDown.Swap(true); !shutdown {
 		s.shutdownCancel()
-		if s.pcapWriter != nil {
 		s.pcapWriter.Close()
 	}
-	}
 	s.wg.Wait()
 }
 
@@ -647,18 +646,16 @@ 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
+			ci.InterfaceIndex = srcNode.interfaceID
 		}
 		must.Do(s.pcapWriter.WritePacket(ci, pkt))
 	}
-	}
 
 	buf := make([]byte, 16<<10)
 	var netw *network // non-nil after first packet
@@ -717,17 +714,13 @@ 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),
-			}
+			InterfaceIndex: srcNode.interfaceID,
-			if srcNode != nil {
-				ci.InterfaceIndex = srcNode.id
 		}
 		must.Do(s.pcapWriter.WritePacket(ci, packetRaw))
-		}
 		netw.HandleEthernetPacket(ep)
 	}
 }
@@ -831,6 +824,17 @@ func (n *network) HandleUDPPacket(p UDPPacket) {
 		return
 	}
 	p.Dst = dst
+	buff, err := n.serializedUDPPacket(p.Src, p.Dst, p.Payload, nil)
+	if err != nil {
+		n.logf("serializing UDP packet: %v", err)
+		return
+	}
+	n.s.pcapWriter.WritePacket(gopacket.CaptureInfo{
+		Timestamp:      time.Now(),
+		CaptureLength:  len(buff),
+		Length:         len(buff),
+		InterfaceIndex: n.interfaceID,
+	}, buff)
 	n.WriteUDPPacketNoNAT(p)
 }
 
@@ -853,6 +857,15 @@ func (n *network) WriteUDPPacketNoNAT(p UDPPacket) {
 		DstMAC:       node.mac.HWAddr(),
 		EthernetType: layers.EthernetTypeIPv4,
 	}
+	ethRaw, err := n.serializedUDPPacket(src, dst, p.Payload, eth)
+	if err != nil {
+		n.logf("serializing UDP packet: %v", err)
+		return
+	}
+	n.writeEth(ethRaw)
+}
+
+func (n *network) serializedUDPPacket(src, dst netip.AddrPort, payload []byte, eth *layers.Ethernet) ([]byte, error) {
 	ip := &layers.IPv4{
 		Version:  4,
 		TTL:      64,
@@ -868,12 +881,14 @@ func (n *network) WriteUDPPacketNoNAT(p UDPPacket) {
 
 	buffer := gopacket.NewSerializeBuffer()
 	options := gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: true}
-	if err := gopacket.SerializeLayers(buffer, options, eth, ip, udp, gopacket.Payload(p.Payload)); err != nil {
+	layers := []gopacket.SerializableLayer{eth, ip, udp, gopacket.Payload(payload)}
-		n.logf("serializing UDP: %v", err)
+	if eth == nil {
-		return
+		layers = layers[1:]
 	}
-	ethRaw := buffer.Bytes()
+	if err := gopacket.SerializeLayers(buffer, options, layers...); err != nil {
-	n.writeEth(ethRaw)
+		return nil, fmt.Errorf("serializing UDP: %v", err)
+	}
+	return buffer.Bytes(), nil
 }
 
 // HandleEthernetIPv4PacketForRouter handles an IPv4 packet that is
@@ -936,6 +951,18 @@ func (n *network) HandleEthernetIPv4PacketForRouter(ep EthernetPacket) {
 		_ = src0
 		//log.Printf("XXX UDP out %v=>%v to %v", src0, src, dst)
 
+		buf, err := n.serializedUDPPacket(src, dst, udp.Payload, nil)
+		if err != nil {
+			n.logf("serializing UDP packet: %v", err)
+			return
+		}
+		n.s.pcapWriter.WritePacket(gopacket.CaptureInfo{
+			Timestamp:      time.Now(),
+			CaptureLength:  len(buf),
+			Length:         len(buf),
+			InterfaceIndex: n.interfaceID,
+		}, buf)
+
 		n.s.routeUDPPacket(UDPPacket{
 			Src:     src,
 			Dst:     dst,