Skip to content

Commit

Permalink
feat(output): Add -C and -W flags for output file rotation based …
Browse files Browse the repository at this point in the history
…on size and count
  • Loading branch information
mozillazg committed Feb 16, 2025
1 parent 8342c51 commit d7735a9
Show file tree
Hide file tree
Showing 13 changed files with 529 additions and 53 deletions.
12 changes: 12 additions & 0 deletions .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,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: |
Expand Down
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,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'
```

Expand Down Expand Up @@ -343,7 +344,6 @@ Flags:
| --pod-name *pod_name.namespace* | ||
| -f, --follow-forks | ||
| -- *command [args]* | ||
| --oneline | ||
| --netns *path_to_net_ns* | ||
| --print |||
| -c *count* |||
Expand All @@ -359,7 +359,7 @@ Flags:
| -vvv |||
| -B *bufer_size*, --buffer-size=*buffer_size* || |
| --count |||
| -C *file_size || |
| -C *file_size || |
| -d || |
| -dd || |
| -ddd || |
Expand Down Expand Up @@ -398,7 +398,7 @@ Flags:
| -u || |
| -U, --packet-buffered || |
| -V *file* || |
| -W *filecont* || |
| -W *filecont* || |
| -y *datalinktype*, --linktype=*datalinktype* || |
| -z *postrotate-command* || |
| -Z *user*, --relinquish-privileges=*user* || |
Expand Down
5 changes: 3 additions & 2 deletions README.zh-CN.md
Original file line number Diff line number Diff line change
Expand Up @@ -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'
```

Expand Down Expand Up @@ -361,7 +362,7 @@ Flags:
| -vvv |||
| -B *bufer_size*, --buffer-size=*buffer_size* || |
| --count |||
| -C *file_size || |
| -C *file_size || |
| -d || |
| -dd || |
| -ddd || |
Expand Down Expand Up @@ -400,7 +401,7 @@ Flags:
| -u || |
| -U, --packet-buffered || |
| -V *file* || |
| -W *filecont* || |
| -W *filecont* || |
| -y *datalinktype*, --linktype=*datalinktype* || |
| -z *postrotate-command* || |
| -Z *user*, --relinquish-privileges=*user* || |
Expand Down
13 changes: 13 additions & 0 deletions cmd/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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,
}
}
9 changes: 9 additions & 0 deletions cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}

Expand Down
75 changes: 41 additions & 34 deletions cmd/writer.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand All @@ -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)
Expand All @@ -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)
}
84 changes: 84 additions & 0 deletions internal/types/flag.go
Original file line number Diff line number Diff line change
@@ -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) }
49 changes: 49 additions & 0 deletions internal/types/flag_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
})
}
}
Loading

0 comments on commit d7735a9

Please sign in to comment.