add more tests

mozillazg · Oct 13, 2024 · 8d477a3 · 8d477a3
1 parent e079e39
commit 8d477a3
Show file tree

Hide file tree

Showing 9 changed files with 240 additions and 35 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -78,6 +78,12 @@ jobs:
           command: |
             sudo bash testdata/test_netns.sh ./ptcpdump
 
+      - run:
+          name: e2e (test netns newly)
+          command: |
+            sudo bash testdata/test_netns_newly_normal.sh ./ptcpdump
+            sudo bash testdata/test_netns_newly_exec.sh ./ptcpdump
+
       - run:
           name: e2e (test nat)
           command: |

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -230,6 +230,16 @@ jobs:
             set -ex
             bash /host/testdata/test_netns.sh /host/ptcpdump/ptcpdump
 
+      - name: Test netns newly
+        # if: ${{ (!startsWith(matrix.kernel, '5.4')) && (!startsWith(matrix.kernel, '4.')) }}
+        uses: cilium/little-vm-helper@97c89f004bd0ab4caeacfe92ebc956e13e362e6b # v0.0.19
+        with:
+          provision: 'false'
+          cmd: |
+            set -ex
+            bash /host/testdata/test_netns_newly_normal.sh /host/ptcpdump/ptcpdump
+            bash /host/testdata/test_netns_newly_exec.sh /host/ptcpdump/ptcpdump
+
       - name: Test run sub program
         uses: cilium/little-vm-helper@97c89f004bd0ab4caeacfe92ebc956e13e362e6b # v0.0.19
         with:

diff --git a/cmd/writer.go b/cmd/writer.go
@@ -3,7 +3,6 @@ package cmd
 import (
 	"fmt"
 	"io"
-	"math"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -78,40 +77,19 @@ func newPcapNgWriter(w io.Writer, pcache *metadata.ProcessCache, opts *Options)
 	}
 
 	// to avoid "Interface id 9 not present in section (have only 7 interfaces)"
-	var maxIndex int
+	maxIndex := 1
 	for _, dev := range devices.Devs() {
 		if dev.Ifindex > maxIndex {
 			maxIndex = dev.Ifindex
 		}
 	}
-	interfaces := make([]pcapgo.NgInterface, maxIndex+1)
+	interfaces := make([]pcapgo.NgInterface, maxIndex)
 	for _, dev := range devices.Devs() {
-		comment := ""
-		if dev.NetNs != nil {
-			comment = fmt.Sprintf("netNsInode: %d, netNsPath: %s", dev.NetNs.Inode(), dev.NetNs.Path())
-		}
-		interfaces[dev.Ifindex] = pcapgo.NgInterface{
-			Index:      dev.Ifindex,
-			Name:       dev.Name,
-			Comment:    comment,
-			Filter:     opts.pcapFilter,
-			OS:         runtime.GOOS,
-			LinkType:   layers.LinkTypeEthernet,
-			SnapLength: uint32(math.MaxUint16),
-			//TimestampResolution: 9,
-		}
+		interfaces[dev.Ifindex] = metadata.NewNgInterface(dev, opts.pcapFilter)
 	}
 	for i, iface := range interfaces {
 		if iface.Index == 0 {
-			interfaces[i] = pcapgo.NgInterface{
-				Index:      i,
-				Name:       fmt.Sprintf("dummy-%d", iface.Index),
-				Filter:     opts.pcapFilter,
-				OS:         runtime.GOOS,
-				LinkType:   layers.LinkTypeEthernet,
-				SnapLength: uint32(math.MaxUint16),
-				//TimestampResolution: 9,
-			}
+			interfaces[i] = metadata.NewDummyNgInterface(i)
 		}
 	}
 
@@ -137,7 +115,8 @@ func newPcapNgWriter(w io.Writer, pcache *metadata.ProcessCache, opts *Options)
 		return nil, fmt.Errorf("writing pcapNg header: %w", err)
 	}
 
-	return writer.NewPcapNGWriter(pcapNgWriter, pcache), nil
+	wt := writer.NewPcapNGWriter(pcapNgWriter, pcache, interfaces).WithPcapFilter(opts.pcapFilter)
+	return wt, nil
 }
 
 func newPcapWriter(w io.Writer) (*writer.PcapWriter, error) {

diff --git a/internal/capturer/capturer.go b/internal/capturer/capturer.go
@@ -379,9 +379,10 @@ func (c *Capturer) handleNewDevEvents() {
 		c.opts.DeviceCache.Add(dev.NetnsId, dev.Ifindex, utils.GoString(dev.Name[:]))
 
 		device, _ := c.opts.DeviceCache.GetByIfindex(int(dev.Ifindex), dev.NetnsId)
+		c.addNewDevToWriter(device)
 		log.Infof("start attach tc hooks to %s, triggered by events", device.String())
 		if err := c.attachTcHooks(device); err != nil {
-			log.Errorf("attach tc hooks failed: %s", err)
+			log.Infof("attach tc hooks failed: %s", err)
 		}
 	}
 }
@@ -395,8 +396,17 @@ func (c *Capturer) handleDevChangeEvents() {
 		c.opts.DeviceCache.Add(dev.NetnsId, dev.Ifindex, utils.GoString(dev.Name[:]))
 
 		device, _ := c.opts.DeviceCache.GetByIfindex(int(dev.Ifindex), dev.NetnsId)
+		c.addNewDevToWriter(device)
 		if err := c.attachTcHooks(device); err != nil {
-			log.Errorf("attach tc hooks failed: %s", err)
+			log.Infof("attach tc hooks failed: %s", err)
+		}
+	}
+}
+
+func (c *Capturer) addNewDevToWriter(dev types.Device) {
+	for _, w := range c.opts.Writers {
+		if pw, ok := w.(*writer.PcapNGWriter); ok {
+			pw.AddDev(dev)
 		}
 	}
 }

diff --git a/internal/metadata/device.go b/internal/metadata/device.go
@@ -4,9 +4,13 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/gopacket/gopacket/layers"
+	"github.com/gopacket/gopacket/pcapgo"
 	"github.com/mozillazg/ptcpdump/internal/log"
 	"github.com/mozillazg/ptcpdump/internal/types"
+	"math"
 	"net"
+	"runtime"
 	"sync"
 )
 
@@ -188,3 +192,31 @@ func (d *DeviceCache) getAllLinks(inode uint32) ([]net.Interface, error) {
 
 	return d.allLinks[inode], nil
 }
+
+func NewNgInterface(dev types.Device, filter string) pcapgo.NgInterface {
+	comment := ""
+	if dev.NetNs != nil {
+		comment = fmt.Sprintf("netNsInode: %d, netNsPath: %s", dev.NetNs.Inode(), dev.NetNs.Path())
+	}
+	return pcapgo.NgInterface{
+		Index:      dev.Ifindex,
+		Name:       dev.Name,
+		Comment:    comment,
+		Filter:     filter,
+		OS:         runtime.GOOS,
+		LinkType:   layers.LinkTypeEthernet,
+		SnapLength: uint32(math.MaxUint16),
+		//TimestampResolution: 9,
+	}
+}
+
+func NewDummyNgInterface(index int) pcapgo.NgInterface {
+	return pcapgo.NgInterface{
+		Index:      index,
+		Name:       fmt.Sprintf("dummy-%d", index),
+		OS:         runtime.GOOS,
+		LinkType:   layers.LinkTypeEthernet,
+		SnapLength: uint32(math.MaxUint16),
+		//TimestampResolution: 9,
+	}
+}
diff --git a/internal/writer/pcapng.go b/internal/writer/pcapng.go
@@ -3,6 +3,7 @@ package writer
 import (
 	"bytes"
 	"fmt"
+	"github.com/mozillazg/ptcpdump/internal/types"
 	"sync"
 
 	"github.com/gopacket/gopacket"
@@ -13,16 +14,19 @@ import (
 )
 
 type PcapNGWriter struct {
-	pw     *pcapgo.NgWriter
-	pcache *metadata.ProcessCache
+	pw         *pcapgo.NgWriter
+	pcache     *metadata.ProcessCache
+	interfaces []pcapgo.NgInterface
+	pcapFilter string
 
 	noBuffer bool
 	lock     sync.Mutex
 	keylogs  bytes.Buffer
 }
 
-func NewPcapNGWriter(pw *pcapgo.NgWriter, pcache *metadata.ProcessCache) *PcapNGWriter {
-	return &PcapNGWriter{pw: pw, pcache: pcache, lock: sync.Mutex{}}
+func NewPcapNGWriter(pw *pcapgo.NgWriter, pcache *metadata.ProcessCache,
+	interfaces []pcapgo.NgInterface) *PcapNGWriter {
+	return &PcapNGWriter{pw: pw, pcache: pcache, interfaces: interfaces, lock: sync.Mutex{}}
 }
 
 func (w *PcapNGWriter) Write(e *event.Packet) error {
@@ -60,7 +64,7 @@ func (w *PcapNGWriter) Write(e *event.Packet) error {
 		)
 	}
 
-	if err := w.writeTLSKeyLog(); err != nil {
+	if err := w.writeTLSKeyLogs(); err != nil {
 		return err
 	}
 
@@ -83,7 +87,7 @@ func (w *PcapNGWriter) WriteTLSKeyLog(line string) error {
 	return nil
 }
 
-func (w *PcapNGWriter) writeTLSKeyLog() error {
+func (w *PcapNGWriter) writeTLSKeyLogs() error {
 	w.lock.Lock()
 	defer w.lock.Unlock()
 
@@ -101,6 +105,35 @@ func (w *PcapNGWriter) writeTLSKeyLog() error {
 	return nil
 }
 
+func (w *PcapNGWriter) AddDev(dev types.Device) {
+	w.lock.Lock()
+	defer w.lock.Unlock()
+
+	log.Infof("new dev: %+v, currLen: %d", dev, len(w.interfaces))
+	if len(w.interfaces) > dev.Ifindex {
+		return
+	}
+
+	for i := len(w.interfaces); i <= dev.Ifindex; i++ {
+		var intf pcapgo.NgInterface
+		if i == dev.Ifindex {
+			intf = metadata.NewNgInterface(dev, w.pcapFilter)
+		} else {
+			intf = metadata.NewDummyNgInterface(i)
+		}
+		log.Debugf("add interface: %+v", intf)
+		if _, err := w.pw.AddInterface(intf); err != nil {
+			log.Errorf("error adding interface %s: %+v", intf.Name, err)
+		}
+		w.interfaces = append(w.interfaces, intf)
+	}
+}
+
+func (w *PcapNGWriter) WithPcapFilter(filter string) *PcapNGWriter {
+	w.pcapFilter = filter
+	return w
+}
+
 func (w *PcapNGWriter) Flush() error {
 	return w.pw.Flush()
 }

diff --git a/testdata/create_netns.sh b/testdata/create_netns.sh
@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+
+set -ex
+
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" > /dev/null && pwd )"
+NETNS1="$1"
+VETH1="$2"
+NETNS2="$3"
+VETH2="$4"
+
+function setup_netns() {
+    ip netns add ${NETNS1} || true
+    ip netns add ${NETNS2} || true
+
+    ip link add ${VETH1} type veth peer name ${VETH2} || true
+    ip link set ${VETH1} netns ${NETNS1} || true
+    ip link set ${VETH2} netns ${NETNS2} || true
+
+    ip -n ${NETNS1} addr add 192.168.64.1/24 dev ${VETH1} || true
+    ip -n ${NETNS2} addr add 192.168.64.2/24 dev ${VETH2} || true
+
+    ip -n ${NETNS1} link set ${VETH1} up
+    ip -n ${NETNS2} link set ${VETH2} up
+}
+
+function main() {
+  setup_netns
+}
+
+main
diff --git a/testdata/test_netns_newly_exec.sh b/testdata/test_netns_newly_exec.sh
@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+
+set -ex
+
+CMD="$1"
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" > /dev/null && pwd )"
+FILE_PREFIX="/tmp/ptcpdump"
+CREATE_NS_SCRIPT="${SCRIPT_DIR}/create_netns.sh"
+FNAME="${FILE_PREFIX}_netns_newly.pcapng"
+LNAME="${FILE_PREFIX}_netns_newly.log"
+RNAME="${FILE_PREFIX}_netns_newly.read.txt"
+NETNS1="netns30"
+VETH1="veth30"
+NETNS2="netns31"
+VETH2="veth31"
+
+function cleanup() {
+    ip netns exec ${NETNS1} ip link delete ${VETH1} || true
+    ip netns exec ${NETNS2} ip link delete ${VETH2} || true
+
+    ip netns delete ${NETNS1}
+    ip netns delete ${NETNS2}
+}
+
+trap cleanup EXIT
+
+function test_ptcpdump_exec() {
+  local LNAME="${LNAME}.exec"
+  curl 1.1.1.1 &>/dev/null || true &
+  ip netns exec ${NETNS2}  sh -c "echo -e 'HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n' | nc -l 8000" &
+
+  timeout 30s ${CMD} -c 4 -i any --netns newly -v --print -w "${FNAME}" \
+      'tcp' \
+      -- sh -c "bash ${CREATE_NS_SCRIPT} $NETNS1 $VETH1 $NETNS2 $VETH2 && ip netns exec ${NETNS1} curl http://192.168.64.2:8000"  | tee "${LNAME}"
+
+  cat "${LNAME}"
+  cat "${LNAME}" | grep '192.168.64.2.* > 192.168.64.1'
+  cat "${LNAME}" | grep '192.168.64.1.* > 192.168.64.2'
+}
+
+function main() {
+  test_ptcpdump_exec
+}
+
+main
diff --git a/testdata/test_netns_newly_normal.sh b/testdata/test_netns_newly_normal.sh
@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+
+set -ex
+
+CMD="$1"
+FILE_PREFIX="/tmp/ptcpdump"
+FNAME="${FILE_PREFIX}_netns_newly.pcapng"
+LNAME="${FILE_PREFIX}_netns_newly.log"
+RNAME="${FILE_PREFIX}_netns_newly.read.txt"
+NETNS1="netns20"
+VETH1="veth20"
+NETNS2="netns21"
+VETH2="veth21"
+
+function setup_netns() {
+    ip netns add ${NETNS1} || true
+    ip netns add ${NETNS2} || true
+
+    ip link add ${VETH1} type veth peer name ${VETH2} || true
+    ip link set ${VETH1} netns ${NETNS1} || true
+    ip link set ${VETH2} netns ${NETNS2} || true
+
+    ip -n ${NETNS1} addr add 192.168.64.1/24 dev ${VETH1} || true
+    ip -n ${NETNS2} addr add 192.168.64.2/24 dev ${VETH2} || true
+
+    ip -n ${NETNS1} link set ${VETH1} up
+    ip -n ${NETNS2} link set ${VETH2} up
+}
+
+function cleanup() {
+    ip netns exec ${NETNS1} ip link delete ${VETH1} || true
+    ip netns exec ${NETNS2} ip link delete ${VETH2} || true
+
+    ip netns delete ${NETNS1}
+    ip netns delete ${NETNS2}
+}
+
+trap cleanup EXIT
+
+function test_ptcpdump_normal() {
+  local LNAME="${LNAME}.normal"
+  timeout 30s ${CMD} -c 4 -i any --netns newly -v --print -w "${FNAME}" \
+      'icmp' | tee "${LNAME}" &
+  sleep 10
+  setup_netns
+
+  ping -c 5 1.1.1.1 &>/dev/null || true &
+  ip netns exec ${NETNS1} ping -c 2 192.168.64.2 &>/dev/null || true
+  wait
+
+  cat "${LNAME}"
+  cat "${LNAME}" | grep -F '192.168.64.1 > 192.168.64.2: ICMP'
+  cat "${LNAME}" | grep -F '192.168.64.2 > 192.168.64.1: ICMP'
+}
+
+function main() {
+  test_ptcpdump_normal
+}
+
+main