chore(bpf): Use TCX where kernel support allows (#187)

* chore(bpf): Use TCX where kernel support allows * ci: increase timeout from 10 to 15
mozillazg · Nov 17, 2024 · 0ef1419 · 0ef1419
1 parent a4ca7e9
commit 0ef1419
Show file tree

Hide file tree

Showing 14 changed files with 107 additions and 5 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -70,7 +70,7 @@ jobs:
           - 'bpf-20241101.013334'
           # renovate: datasource=docker depName=quay.io/lvh-images/kernel-images
           - 'bpf-next-20241101.013334'
-    timeout-minutes: 10
+    timeout-minutes: 15
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4

diff --git a/bpf/bpf.go b/bpf/bpf.go
@@ -15,6 +15,7 @@ import (
 	"github.com/jschwinger233/elibpcap"
 	"github.com/mozillazg/ptcpdump/internal/log"
 	"github.com/mozillazg/ptcpdump/internal/types"
+	"github.com/mozillazg/ptcpdump/internal/utils"
 	"golang.org/x/sys/unix"
 )
 
@@ -32,6 +33,7 @@ type BPF struct {
 	closeFuncs []func()
 
 	skipAttachCgroup bool
+	skipTcx          bool
 	isLegacyKernel   bool
 	report           *types.CountReport
 }
@@ -86,6 +88,7 @@ func NewBPF() (*BPF, error) {
 		report:           &types.CountReport{},
 		isLegacyKernel:   legacyKernel,
 		skipAttachCgroup: skipAttachCgroup,
+		skipTcx:          !supportTcx(),
 	}
 
 	return bf, nil
@@ -113,10 +116,15 @@ func (b *BPF) Load(opts Options) error {
 	}
 
 	if opts.pcapFilter != "" {
-		for _, progName := range []string{"tc_ingress", "tc_egress"} {
+		for _, progName := range []string{"tc_ingress", "tc_egress", "tcx_ingress", "tcx_egress"} {
 			prog, ok := b.spec.Programs[progName]
 			if !ok {
-				return fmt.Errorf("program %s not found", progName)
+				log.Infof("program %s not found", progName)
+				continue
+			}
+			if prog == nil {
+				log.Infof("program %s is nil", progName)
+				continue
 			}
 			prog.Instructions, err = elibpcap.Inject(
 				opts.pcapFilter,
@@ -372,6 +380,54 @@ func (b *BPF) AttachTracepoints() error {
 }
 
 func (b *BPF) AttachTcHooks(ifindex int, egress, ingress bool) ([]func(), error) {
+	closers, err := b.attachTcxHooks(ifindex, egress, ingress)
+	if err != nil {
+		log.Infof("attach tcx failed, fallback to tc: %+v", err)
+		utils.RunClosers(closers)
+		closers, err = b.attachTcHooks(ifindex, egress, ingress)
+	}
+	return closers, err
+}
+
+func (b *BPF) attachTcxHooks(ifindex int, egress, ingress bool) ([]func(), error) {
+	var closeFuncs []func()
+
+	if b.skipTcx || b.objs.TcxEgress == nil || b.objs.TcxIngress == nil {
+		return closeFuncs, errors.New("tcx programs not found")
+	}
+
+	if egress {
+		lk, err := link.AttachTCX(link.TCXOptions{
+			Interface: ifindex,
+			Program:   b.objs.TcxEgress,
+			Attach:    ebpf.AttachTCXEgress,
+		})
+		if err != nil {
+			return closeFuncs, fmt.Errorf("attach tcx/egress hooks: %w", err)
+		}
+		closeFuncs = append(closeFuncs, func() {
+			lk.Close()
+		})
+	}
+
+	if ingress {
+		lk, err := link.AttachTCX(link.TCXOptions{
+			Interface: ifindex,
+			Program:   b.objs.TcxIngress,
+			Attach:    ebpf.AttachTCXIngress,
+		})
+		if err != nil {
+			return closeFuncs, fmt.Errorf("attach tcx/ingress hooks: %w", err)
+		}
+		closeFuncs = append(closeFuncs, func() {
+			lk.Close()
+		})
+	}
+
+	return closeFuncs, nil
+}
+
+func (b *BPF) attachTcHooks(ifindex int, egress, ingress bool) ([]func(), error) {
 	var closeFuncs []func()
 	closeFunc, err := ensureTcQdisc(ifindex)
 	if err != nil {

diff --git a/bpf/bpf_arm64_bpfel.go b/bpf/bpf_arm64_bpfel.go
diff --git a/bpf/bpf_arm64_bpfel.o b/bpf/bpf_arm64_bpfel.o
diff --git a/bpf/bpf_legacy.go b/bpf/bpf_legacy.go
@@ -11,7 +11,7 @@ import (
 )
 
 // $TARGET is set by the Makefile
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -no-strip -no-global-types -target $TARGET bpf_legacy ./ptcpdump.c -- -I./headers -I./headers/$TARGET -I. -Wall -DLEGACY_KERNEL -DNO_CGROUP_PROG -DNO_TRACING
+//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -no-strip -no-global-types -target $TARGET bpf_legacy ./ptcpdump.c -- -I./headers -I./headers/$TARGET -I. -Wall -DLEGACY_KERNEL -DNO_CGROUP_PROG -DNO_TRACING -DNO_TCX
 
 func supportCgroupSock() bool {
 	if err := features.HaveProgramHelper(ebpf.CGroupSock, asm.FnGetSocketCookie); err != nil {
@@ -46,6 +46,14 @@ func isLegacyKernel() (bool, error) {
 	return false, nil
 }
 
+func supportTcx() bool {
+	versionCode, _ := features.LinuxVersionCode()
+	if versionCode <= 0 {
+		return false
+	}
+	return versionCode >= kernelVersion(6, 6, 0)
+}
+
 func loadBpfWithData(b []byte) (*ebpf.CollectionSpec, error) {
 	reader := bytes.NewReader(b)
 	spec, err := ebpf.LoadCollectionSpecFromReader(reader)

diff --git a/bpf/bpf_legacy_arm64_bpfel.o b/bpf/bpf_legacy_arm64_bpfel.o
diff --git a/bpf/bpf_legacy_x86_bpfel.o b/bpf/bpf_legacy_x86_bpfel.o
diff --git a/bpf/bpf_no_tracing.go b/bpf/bpf_no_tracing.go
@@ -7,7 +7,7 @@ import (
 )
 
 // $TARGET is set by the Makefile
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -no-strip -no-global-types -target $TARGET bpf_no_tracing ./ptcpdump.c -- -I./headers -I./headers/$TARGET -I. -Wall -DNO_TRACING
+//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -no-strip -no-global-types -target $TARGET bpf_no_tracing ./ptcpdump.c -- -I./headers -I./headers/$TARGET -I. -Wall -DNO_TRACING -DNO_TCX
 
 func supportTracing() bool {
 	if err := features.HaveProgramType(ebpf.Tracing); err != nil {

diff --git a/bpf/bpf_no_tracing_arm64_bpfel.o b/bpf/bpf_no_tracing_arm64_bpfel.o
diff --git a/bpf/bpf_no_tracing_x86_bpfel.o b/bpf/bpf_no_tracing_x86_bpfel.o
diff --git a/bpf/bpf_x86_bpfel.go b/bpf/bpf_x86_bpfel.go
diff --git a/bpf/bpf_x86_bpfel.o b/bpf/bpf_x86_bpfel.o
diff --git a/bpf/ptcpdump.c b/bpf/ptcpdump.c
@@ -23,6 +23,7 @@
 #define IPPROTO_UDP 17    /* User Datagram Protocol		*/
 #define IPPROTO_SCTP 132  /* Stream Control Transport Protocol	*/
 #define TC_ACT_UNSPEC (-1)
+#define TCX_NEXT (-1)
 #define AF_INET 2
 #define AF_INET6 10
 #define INGRESS_PACKET 1
@@ -1085,8 +1086,24 @@ int tc_ingress(struct __sk_buff *skb) {
     return TC_ACT_UNSPEC;
 }
 
+#ifndef NO_TCX
+SEC("tcx/ingress")
+int tcx_ingress(struct __sk_buff *skb) {
+    handle_tc(skb, false);
+    return TCX_NEXT;
+}
+#endif
+
 SEC("tc")
 int tc_egress(struct __sk_buff *skb) {
     handle_tc(skb, true);
     return TC_ACT_UNSPEC;
 }
+
+#ifndef NO_TCX
+SEC("tcx/egress")
+int tcx_egress(struct __sk_buff *skb) {
+    handle_tc(skb, true);
+    return TCX_NEXT;
+}
+#endif
diff --git a/internal/utils/error.go b/internal/utils/error.go
@@ -11,3 +11,12 @@ func UnwrapErr(err error) error {
 		}
 	}
 }
+
+func RunClosers(funcs []func()) {
+	for i := len(funcs) - 1; i >= 0; i-- {
+		f := funcs[i]
+		if f != nil {
+			f()
+		}
+	}
+}