feat(output): Add -C and -W flags for output file rotation based on size and count

.github/workflows/test.yml

@@ -319,6 +319,18 @@ jobs:
             done
             exit 1
 
+      - name: Test rotate
+        uses: cilium/little-vm-helper@e87948476ca97050b1f149ab2aec379d0de19b84 # v0.0.23
+        if: ${{ (!startsWith(matrix.backend, 'cgroup-skb')) && (contains(matrix.kernel, 'next')) }} # no need test for all kernels
+        with:
+          provision: 'false'
+          cmd: |
+            set -ex
+            export PTCPDUMP_EXTRA_ARGS="${{ env.PTCPDUMP_EXTRA_ARGS }}"
+
+            bash /host/testdata/test_rotate_filesize.sh /host/ptcpdump/ptcpdump
+            bash /host/testdata/test_rotate_filesize_with_count.sh /host/ptcpdump/ptcpdump
+
       - name: build demo app
         if: ${{ (!startsWith(matrix.kernel, '5.4')) && (!startsWith(matrix.kernel, '4.')) }}
         run: |

README.md

@@ -4,6 +4,7 @@
 
 [![amd64-e2e](https://img.shields.io/github/actions/workflow/status/mozillazg/ptcpdump/test.yml?label=x86_64%20(amd64)%20e2e)](https://github.com/mozillazg/ptcpdump/actions/workflows/test.yml)
 [![arm64-e2e](https://img.shields.io/circleci/build/gh/mozillazg/ptcpdump/master?label=aarch64%20(arm64)%20e2e)](https://app.circleci.com/pipelines/github/mozillazg/ptcpdump?branch=master)
+[![Release](https://img.shields.io/github/v/release/mozillazg/ptcpdump)](https://github.com/mozillazg/ptcpdump/releases)
 English | [中文](README.zh-CN.md)
 
 
@@ -67,6 +68,7 @@ Filter like tcpdump:
 ```
 sudo ptcpdump -i eth0 tcp
 sudo ptcpdump -i eth0 -A -s 0 -n -v tcp and port 80 and host 10.10.1.1
+sudo ptcpdump -i any -s 0 -n -v -C 100MB -W 3 'tcp and port 80 and host 10.10.1.1'
 sudo ptcpdump -i eth0 'tcp[tcpflags] & (tcp-syn|tcp-fin) != 0'
 ```
 
@@ -358,7 +360,7 @@ Flags:
 | -vvv                                              | ✅       | ⭕                        |
 | -B *bufer_size*, --buffer-size=*buffer_size*      | ✅       |                          |
 | --count                                           | ✅       | ✅                        |
-| -C *file_size                                     | ✅       |                          |
+| -C *file_size                                     | ✅       | ✅                        |
 | -d                                                | ✅       |                          |
 | -dd                                               | ✅       |                          |
 | -ddd                                              | ✅       |                          |
@@ -397,7 +399,7 @@ Flags:
 | -u                                                | ✅       |                          |
 | -U, --packet-buffered                             | ✅       |                          |
 | -V *file*                                         | ✅       |                          |
-| -W *filecont*                                     | ✅       |                          |
+| -W *filecont*                                     | ✅       | ✅                        |
 | -y *datalinktype*, --linktype=*datalinktype*      | ✅       |                          |
 | -z *postrotate-command*                           | ✅       |                          |
 | -Z *user*, --relinquish-privileges=*user*         | ✅       |                          |

README.zh-CN.md

@@ -75,6 +75,7 @@ Table of Contents
 ```
 sudo ptcpdump -i eth0 tcp
 sudo ptcpdump -i eth0 -A -s 0 -n -v tcp and port 80 and host 10.10.1.1
+sudo ptcpdump -i any -s 0 -n -v -C 100MB -W 3 'tcp and port 80 and host 10.10.1.1'
 sudo ptcpdump -i eth0 'tcp[tcpflags] & (tcp-syn|tcp-fin) != 0'
 ```
 
@@ -360,7 +361,7 @@ Flags:
 | -vvv                                              | ✅       | ⭕                        |
 | -B *bufer_size*, --buffer-size=*buffer_size*      | ✅       |                          |
 | --count                                           | ✅       | ✅                        |
-| -C *file_size                                     | ✅       |                          |
+| -C *file_size                                     | ✅       | ✅                        |
 | -d                                                | ✅       |                          |
 | -dd                                               | ✅       |                          |
 | -ddd                                              | ✅       |                          |
@@ -399,7 +400,7 @@ Flags:
 | -u                                                | ✅       |                          |
 | -U, --packet-buffered                             | ✅       |                          |
 | -V *file*                                         | ✅       |                          |
-| -W *filecont*                                     | ✅       |                          |
+| -W *filecont*                                     | ✅       | ✅                        |
 | -y *datalinktype*, --linktype=*datalinktype*      | ✅       |                          |
 | -z *postrotate-command*                           | ✅       |                          |
 | -Z *user*, --relinquish-privileges=*user*         | ✅       |                          |

cmd/options.go

@@ -46,6 +46,11 @@ type Options struct {
 	direction      string
 	oneLine        bool
 
+	noBuffer      bool
+	fileCount     uint
+	fileSize      types.FlagTypeFileSize
+	fileSizeBytes int64
+
 	printPacketNumber   bool
 	dontPrintTimestamp  bool
 	onlyPrintCount      bool
@@ -184,6 +189,7 @@ func prepareOptions(opts *Options, rawArgs []string, args []string) {
 		opts.backend = string(types.NetHookBackendTc)
 	}
 
+	opts.fileSizeBytes = opts.fileSize.Bytes()
 }
 
 func getPodNameFilter(raw string) (name, ns string) {
@@ -383,3 +389,10 @@ func (o *Options) enableContainerContext() bool {
 func (o *Options) enablePodContext() bool {
 	return o.enhancedContext.PodContext()
 }
+
+func (o *Options) rotatorOption() writer.RotatorOption {
+	return writer.RotatorOption{
+		MaxFileNumber:    int(o.fileCount),
+		MaxFileSizeBytes: o.fileSizeBytes,
+	}
+}

cmd/root.go

@@ -138,6 +138,15 @@ func init() {
 			fmt.Sprintf("Possible values are %q and %q",
 				types.NetHookBackendTc, types.NetHookBackendCgroupSkb))
 
+	rootCmd.Flags().VarP(&opts.fileSize, "file-size", "C",
+		"Before writing a raw packet to a savefile, check whether the file is currently larger than file_size and, "+
+			"if so, close the current savefile and open a new one. Savefiles after the first savefile "+
+			"will have the name specified with the -w flag, "+
+			"with a number after it, starting at 1 and continuing upward.")
+	rootCmd.Flags().UintVarP(&opts.fileCount, "file-count", "W", 0,
+		"Used in conjunction with the -C option, this will limit the number of files created to the specified number, "+
+			"and begin overwriting files from the beginning, thus creating a 'rotating' buffer.")
+
 	silenceKlog()
 }
 

cmd/writer.go

@@ -16,39 +16,39 @@ import (
 
 func getWriters(opts *Options, pcache *metadata.ProcessCache) ([]writer.PacketWriter, func() error, error) {
 	var writers []writer.PacketWriter
-	var pcapFile *os.File
+	var rr writer.Rotator
 	var err error
 
 	if opts.WritePath() != "" {
 		ext := filepath.Ext(opts.WritePath())
 		switch {
 		case opts.WritePath() == "-":
-			w, err := newPcapNgWriter(os.Stdout, pcache, opts)
+			w, err := newPcapNgWriter(writer.NewStdoutRotator(), pcache, opts)
 			if err != nil {
 				return nil, nil, fmt.Errorf(": %w", err)
 			}
 			w.WithNoBuffer()
 			writers = append(writers, w)
 			break
 		case ext == extPcap:
-			pcapFile, err = os.Create(opts.WritePath())
+			rr, err = writer.NewFileRotator(opts.WritePath(), opts.rotatorOption())
 			if err != nil {
 				return nil, nil, fmt.Errorf(": %w", err)
 			}
-			w, err := newPcapWriter(pcapFile)
+			w, err := newPcapWriter(rr)
 			if err != nil {
-				return nil, pcapFile.Close, fmt.Errorf(": %w", err)
+				return nil, rr.Close, fmt.Errorf(": %w", err)
 			}
 			writers = append(writers, w)
 			break
 		default:
-			pcapFile, err = os.Create(opts.WritePath())
+			rr, err = writer.NewFileRotator(opts.WritePath(), opts.rotatorOption())
 			if err != nil {
 				return nil, nil, fmt.Errorf(": %w", err)
 			}
-			w, err := newPcapNgWriter(pcapFile, pcache, opts)
+			w, err := newPcapNgWriter(rr, pcache, opts)
 			if err != nil {
-				return nil, pcapFile.Close, fmt.Errorf(": %w", err)
+				return nil, rr.Close, fmt.Errorf(": %w", err)
 			}
 			writers = append(writers, w)
 		}
@@ -60,17 +60,17 @@ func getWriters(opts *Options, pcache *metadata.ProcessCache) ([]writer.PacketWr
 	}
 
 	closer := func() error {
-		if pcapFile != nil {
-			pcapFile.Sync()
-			return pcapFile.Close()
+		if rr != nil {
+			rr.Flush()
+			return rr.Close()
 		}
 		return nil
 	}
 
 	return writers, closer, nil
 }
 
-func newPcapNgWriter(w io.Writer, pcache *metadata.ProcessCache, opts *Options) (*writer.PcapNGWriter, error) {
+func newPcapNgWriter(rr writer.Rotator, pcache *metadata.ProcessCache, opts *Options) (*writer.PcapNGWriter, error) {
 	devices, err := opts.GetDevices()
 	if err != nil {
 		return nil, fmt.Errorf(": %w", err)
@@ -85,39 +85,46 @@ func newPcapNgWriter(w io.Writer, pcache *metadata.ProcessCache, opts *Options)
 		interfaceIds[dev.Key()] = index
 	}
 
-	pcapNgWriter, err := pcapgo.NewNgWriterInterface(w, interfaces[0], pcapgo.NgWriterOptions{
-		SectionInfo: pcapgo.NgSectionInfo{
-			Hardware:    runtime.GOARCH,
-			OS:          runtime.GOOS,
-			Application: fmt.Sprintf("ptcpdump %s", internal.Version),
-			Comment:     "ptcpdump: https://github.com/mozillazg/ptcpdump",
-		},
-	})
-	if err != nil {
-		return nil, fmt.Errorf(": %w", err)
-	}
-	for _, ifc := range interfaces[1:] {
-		_, err := pcapNgWriter.AddInterface(ifc)
+	f := func(w io.Writer) (*pcapgo.NgWriter, error) {
+		pcapNgWriter, err := pcapgo.NewNgWriterInterface(w, interfaces[0], pcapgo.NgWriterOptions{
+			SectionInfo: pcapgo.NgSectionInfo{
+				Hardware:    runtime.GOARCH,
+				OS:          runtime.GOOS,
+				Application: fmt.Sprintf("ptcpdump %s", internal.Version),
+				Comment:     "ptcpdump: https://github.com/mozillazg/ptcpdump",
+			},
+		})
 		if err != nil {
 			return nil, fmt.Errorf(": %w", err)
 		}
+		for _, ifc := range interfaces[1:] {
+			_, err := pcapNgWriter.AddInterface(ifc)
+			if err != nil {
+				return nil, fmt.Errorf(": %w", err)
+			}
+		}
+		return pcapNgWriter, nil
 	}
 
-	if err := pcapNgWriter.Flush(); err != nil {
-		return nil, fmt.Errorf("writing pcapNg header: %w", err)
+	wt, err := writer.NewPcapNGWriter(rr, f, pcache, interfaceIds)
+	if err != nil {
+		return nil, fmt.Errorf(": %w", err)
 	}
-
-	wt := writer.NewPcapNGWriter(pcapNgWriter, pcache, interfaceIds).WithPcapFilter(opts.pcapFilter)
+	wt.WithPcapFilter(opts.pcapFilter)
 	wt.WithEnhancedContext(opts.enhancedContext)
 	return wt, nil
 }
 
-func newPcapWriter(w io.Writer) (*writer.PcapWriter, error) {
-	pcapWriter := pcapgo.NewWriterNanos(w)
+func newPcapWriter(rr writer.Rotator) (*writer.PcapWriter, error) {
 
-	if err := pcapWriter.WriteFileHeader(65536, layers.LinkTypeEthernet); err != nil {
-		return nil, fmt.Errorf("writing pcap header: %w", err)
+	f := func(w io.Writer) (*pcapgo.Writer, error) {
+		pcapWriter := pcapgo.NewWriterNanos(w)
+
+		if err := pcapWriter.WriteFileHeader(65536, layers.LinkTypeEthernet); err != nil {
+			return nil, fmt.Errorf("writing pcap header: %w", err)
+		}
+		return pcapWriter, nil
 	}
 
-	return writer.NewPcapWriter(pcapWriter), nil
+	return writer.NewPcapWriter(rr, f)
 }

internal/types/flag.go

@@ -0,0 +1,84 @@
+package types
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+type FlagTypeFileSize struct {
+	val string
+	n   uint64
+}
+
+func (s *FlagTypeFileSize) Set(val string) error {
+	n, err := strconv.ParseUint(val, 10, 64)
+	if err == nil {
+		s.n = n * 1_000_000
+		return nil
+	}
+	val = strings.ToLower(val)
+	switch {
+	case strings.HasSuffix(val, "k"):
+		val = strings.TrimSuffix(val, "k")
+		n, err = strconv.ParseUint(val, 10, 64)
+		if err != nil {
+			return err
+		}
+		s.n = n * 1024
+		break
+	case strings.HasSuffix(val, "kb"):
+		val = strings.TrimSuffix(val, "kb")
+		n, err = strconv.ParseUint(val, 10, 64)
+		if err != nil {
+			return err
+		}
+		s.n = n * 1024
+		break
+	case strings.HasSuffix(val, "m"):
+		val = strings.TrimSuffix(val, "m")
+		n, err = strconv.ParseUint(val, 10, 64)
+		if err != nil {
+			return err
+		}
+		s.n = n * 1024 * 1024
+		break
+	case strings.HasSuffix(val, "mb"):
+		val = strings.TrimSuffix(val, "mb")
+		n, err = strconv.ParseUint(val, 10, 64)
+		if err != nil {
+			return err
+		}
+		s.n = n * 1024 * 1024
+		break
+	case strings.HasSuffix(val, "g"):
+		val = strings.TrimSuffix(val, "g")
+		n, err = strconv.ParseUint(val, 10, 64)
+		if err != nil {
+			return err
+		}
+		s.n = n * 1024 * 1024 * 1024
+		break
+	case strings.HasSuffix(val, "gb"):
+		val = strings.TrimSuffix(val, "gb")
+		n, err = strconv.ParseUint(val, 10, 64)
+		if err != nil {
+			return err
+		}
+		s.n = n * 1024 * 1024 * 1024
+		break
+	default:
+		return fmt.Errorf("invalid file size: %s", val)
+	}
+
+	return nil
+}
+func (s *FlagTypeFileSize) Type() string {
+	return "fileSize"
+}
+
+func (s *FlagTypeFileSize) Bytes() int64 {
+	return int64(s.n)
+}
+
+func (s *FlagTypeFileSize) String() string { return string(s.val) }

internal/types/flag_test.go

@@ -0,0 +1,49 @@
+package types
+
+import (
+	"testing"
+)
+
+func TestFlagTypeFileSize_Set(t *testing.T) {
+	tests := []struct {
+		name    string
+		val     string
+		want    uint64
+		wantErr bool
+	}{
+		{"valid bytes", "1000", 1000 * 1000000, false},
+		{"valid kilobytes", "1k", 1 * 1024, false},
+		{"valid kilobytes", "1K", 1 * 1024, false},
+		{"valid kilobytes with kb", "1kb", 1 * 1024, false},
+		{"valid kilobytes with kb", "1KB", 1 * 1024, false},
+		{"valid megabytes", "1m", 1 * 1024 * 1024, false},
+		{"valid megabytes", "1M", 1 * 1024 * 1024, false},
+		{"valid megabytes with mb", "1mb", 1 * 1024 * 1024, false},
+		{"valid megabytes with mb", "1MB", 1 * 1024 * 1024, false},
+		{"valid gigabytes", "1g", 1 * 1024 * 1024 * 1024, false},
+		{"valid gigabytes", "1G", 1 * 1024 * 1024 * 1024, false},
+		{"valid gigabytes with gb", "1gb", 1 * 1024 * 1024 * 1024, false},
+		{"valid gigabytes with gb", "1GB", 1 * 1024 * 1024 * 1024, false},
+		{"invalid format", "1tb", 0, true},
+		{"invalid number", "abc", 0, true},
+		{"empty string", "", 0, true},
+		{"negative number", "-1k", 0, true},
+		{"zero value", "0", 0, false},
+		{"whitespace around value", " 1k ", 0, true},
+		{"uppercase suffix", "1K", 1 * 1024, false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := &FlagTypeFileSize{}
+			err := s.Set(tt.val)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Set() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if s.n != tt.want {
+				t.Errorf("Set() = %v, want %v", s.n, tt.want)
+			}
+		})
+	}
+}