From 680f2d2445de9026e20f67a3d14a2bb889190f4c Mon Sep 17 00:00:00 2001
From: mozillazg <mozillazg101@gmail.com>
Date: Sun, 28 Apr 2024 22:24:18 +0800
Subject: [PATCH] add new flag: -Q/--direction

---
 README.md      |  1 +
 bpf/bpf.go     | 32 +++++++++++++++++++++-----------
 cmd/ebpf.go    |  2 +-
 cmd/options.go | 12 ++++++++++++
 cmd/root.go    |  8 +++++++-
 5 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index a8ddbfbf..a5f47730 100644
--- a/README.md
+++ b/README.md
@@ -42,6 +42,7 @@ Examples:
 Expression: see "man 7 pcap-filter"
 
 Flags:
+  -Q, --direction string     Choose send/receive direction for which packets should be captured. Possible values are 'in', 'out' and 'inout' (default "inout")
   -f, --follow-forks         Include child processes when filter by process
   -h, --help                 help for ptcpdump
   -i, --interface strings    Interfaces to capture (default [lo])
diff --git a/bpf/bpf.go b/bpf/bpf.go
index 868c03f5..36484b28 100644
--- a/bpf/bpf.go
+++ b/bpf/bpf.go
@@ -162,7 +162,7 @@ func (b *BPF) AttachTracepoints() error {
 	return nil
 }
 
-func (b *BPF) AttachTcHooks(ifindex int) error {
+func (b *BPF) AttachTcHooks(ifindex int, egress, ingress bool) error {
 	closeFunc, err := ensureTcQdisc(ifindex)
 	if err != nil {
 		if closeFunc != nil {
@@ -171,20 +171,30 @@ func (b *BPF) AttachTcHooks(ifindex int) error {
 		return xerrors.Errorf("attach tc hooks: %w", err)
 	}
 
-	c1, err := attachTcHook(ifindex, b.objs.TcEgress, false)
-	if err != nil {
-		closeFunc()
-		return xerrors.Errorf("attach tc hooks: %w", err)
+	if egress {
+		c1, err := attachTcHook(ifindex, b.objs.TcEgress, false)
+		if err != nil {
+			if c1 != nil {
+				c1()
+			}
+			closeFunc()
+			return xerrors.Errorf("attach tc hooks: %w", err)
+		}
+		b.closeFuncs = append(b.closeFuncs, c1)
 	}
 
-	c2, err := attachTcHook(ifindex, b.objs.TcIngress, true)
-	if err != nil {
-		c1()
-		closeFunc()
-		return xerrors.Errorf("attach tc hooks: %w", err)
+	if ingress {
+		c2, err := attachTcHook(ifindex, b.objs.TcIngress, true)
+		if err != nil {
+			if c2 != nil {
+				c2()
+			}
+			closeFunc()
+			return xerrors.Errorf("attach tc hooks: %w", err)
+		}
+		b.closeFuncs = append(b.closeFuncs, c2)
 	}
 
-	b.closeFuncs = append(b.closeFuncs, closeFunc, c1, c2)
 	return nil
 }
 
diff --git a/cmd/ebpf.go b/cmd/ebpf.go
index ec15fdd5..9b9ed0a9 100644
--- a/cmd/ebpf.go
+++ b/cmd/ebpf.go
@@ -29,7 +29,7 @@ func attachHooks(opts Options) (map[int]dev.Device, *bpf.BPF, error) {
 		return devices, bf, err
 	}
 	for _, iface := range devices {
-		if err := bf.AttachTcHooks(iface.Ifindex); err != nil {
+		if err := bf.AttachTcHooks(iface.Ifindex, opts.DirectionOut(), opts.DirectionIn()); err != nil {
 			return devices, bf, err
 		}
 	}
diff --git a/cmd/options.go b/cmd/options.go
index bb8a3271..33479df1 100644
--- a/cmd/options.go
+++ b/cmd/options.go
@@ -11,8 +11,20 @@ type Options struct {
 	version        bool
 	print          bool
 	maxPacketCount uint
+	direction      string
 }
 
 func (o Options) WritePath() string {
 	return o.writeFilePath
 }
+
+func (o Options) DirectionIn() bool {
+	return o.DirectionInOut() || o.direction == "in"
+}
+func (o Options) DirectionOut() bool {
+	return o.DirectionInOut() || o.direction == "out"
+}
+
+func (o Options) DirectionInOut() bool {
+	return o.direction == "inout"
+}
diff --git a/cmd/root.go b/cmd/root.go
index 9b8bb528..eb025523 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -53,7 +53,10 @@ func init() {
 		"Print the ptcpdump and libpcap version strings and exit")
 	rootCmd.Flags().BoolVar(&opts.print, "print", false,
 		"Print parsed packet output, even if the raw packets are being saved to a file with the -w flag")
-	rootCmd.Flags().UintVarP(&opts.maxPacketCount, "receive-count", "c", 0, "Exit after receiving count packets")
+	rootCmd.Flags().UintVarP(&opts.maxPacketCount, "receive-count", "c", 0,
+		"Exit after receiving count packets")
+	rootCmd.Flags().StringVarP(&opts.direction, "direction", "Q",
+		"inout", "Choose send/receive direction for which packets should be captured. Possible values are 'in', 'out' and 'inout'")
 }
 
 func Execute() error {
@@ -82,6 +85,9 @@ func run(cmd *cobra.Command, args []string) error {
 
 	devices, bf, err := attachHooks(opts)
 	if err != nil {
+		if bf != nil {
+			bf.Close()
+		}
 		return err
 	}
 	defer bf.Close()