let cmd args control skb_track

jschwinger233 · Mar 21, 2023 · ab2910b · ab2910b
1 parent b6d0b25
commit ab2910b
Show file tree

Hide file tree

Showing 12 changed files with 62 additions and 51 deletions.
diff --git a/config.go b/config.go
@@ -13,6 +13,7 @@ import (
 type Config struct {
 	Iface         string
 	PerfOutput    bool
+	SkbTrack      bool
 	SkbFilename   string
 	PcapFilename  string
 	PcapFilterExp string
@@ -23,11 +24,12 @@ var (
 	bpfObjects bpf.BpfObjects
 )
 
-func initConfig() {
+func mustInitConfig() {
 	flag.StringVarP(&config.Iface, "interface", "i", "lo", "interface to capture")
 	flag.BoolVarP(&config.PerfOutput, "perf-output", "", false, "use bpf_perf_event_output to lift payload size limit")
 	flag.StringVarP(&config.SkbFilename, "skb-filename", "s", "skbdump.skb", "output skb filename")
 	flag.StringVarP(&config.PcapFilename, "pcap-filename", "w", "skbdump.pcap", "output pcap filename")
+	flag.BoolVarP(&config.SkbTrack, "skb-track", "t", false, "track skb by address")
 	flag.Parse()
 	config.PcapFilterExp = strings.Join(flag.Args(), " ")
 

diff --git a/internal/bpf/bpf.go b/internal/bpf/bpf.go
@@ -32,8 +32,17 @@ type Skb struct {
 	Data []byte
 }
 
+type LoadOptions struct {
+	Filter    []bpf.Instruction
+	BpfConfig BpfConfig
+}
+
+type BpfConfig struct {
+	SkbTrack bool
+}
+
 type BpfObjects interface {
-	Load(cbpf []bpf.Instruction) error
+	Load(LoadOptions) error
 	IngressFilter() *ebpf.Program
 	EgressFilter() *ebpf.Program
 	PollSkb(context.Context) (<-chan Skb, error)

diff --git a/internal/bpf/headers/skbdump.h b/internal/bpf/headers/skbdump.h
@@ -61,3 +61,12 @@ struct skb_meta {
 	__u32	cb[5];
 };
 #endif
+
+#ifndef SKBDUMP_CONFIG_DEFINED
+#define SKBDUMP_CONFIG_DEFINED
+struct skbdump_config {
+	bool skb_track;
+};
+
+static volatile const struct skbdump_config SKBDUMP_CONFIG = {};
+#endif
diff --git a/internal/bpf/pcap_filter.go b/internal/bpf/pcap_filter.go
@@ -41,6 +41,10 @@ func mustFindUpDevice() string {
 }
 
 func MustPcapCompile(expr string) (insts []bpf.Instruction) {
+	if len(expr) == 0 {
+		return
+	}
+
 	buf := (*C.char)(C.calloc(C.PCAP_ERRBUF_SIZE, 1))
 	defer C.free(unsafe.Pointer(buf))
 

diff --git a/internal/bpf/perf/bpf.go b/internal/bpf/perf/bpf.go
@@ -72,8 +72,8 @@ func (o *PerfBpfObjects) setFilter(cbpfFilter []bpf.Instruction) (err error) {
 	return
 }
 
-func (o *PerfBpfObjects) Load(cbpfFilter []bpf.Instruction) (err error) {
-	if err = o.setFilter(cbpfFilter); err != nil {
+func (o *PerfBpfObjects) Load(opts internalbpf.LoadOptions) (err error) {
+	if err = o.setFilter(opts.Filter); err != nil {
 		return
 	}
 	if err = errors.WithStack(o.spec.LoadAndAssign(o.objs, nil)); err != nil {

diff --git a/internal/bpf/perf/skbdump_bpfel_x86.o b/internal/bpf/perf/skbdump_bpfel_x86.o
diff --git a/internal/bpf/queue/bpf.go b/internal/bpf/queue/bpf.go
@@ -15,29 +15,34 @@ import (
 //go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -no-strip -target native -type skb_meta -type skb_data Skbdump ./skbdump.c -- -I../headers -I. -Wall
 
 const (
-	/* bpf bytecode
-				00000000000467a8 <on_ingress>:
+	/*
+				00000000000467c8 <on_ingress>:
 				; {
-					36085:	r6 = r1
-					36086:	r7 = 0
+					36089:	r6 = r1
 				; 	bpf_skb_pull_data(skb, 0);
-					36087:	r2 = 0
-					36088:	call 39
+					36090:	r2 = 0
+					36091:	call 39
 				; 	__u64 skb_addr = (__u64)(void *)skb;
-					36089:	*(u64 *)(r10 - 104) = r6
-					36090:	r2 = r10
-					36091:	r2 += -104
-				; 	if (bpf_map_lookup_elem(&skb_address, &skb_addr))
-					36092:	r1 = 0 ll
-					36094:	call 1
-	GotoIndex ->	36095:	if r0 != 0 goto +11 <LBB1501_3>
+					36092:	*(u64 *)(r10 - 104) = r6
+				; 	if (SKBDUMP_CONFIG.skb_track && bpf_map_lookup_elem(&skb_address, &skb_addr))
+					36093:	r1 = 0 ll
+					36095:	r1 = *(u8 *)(r1 + 0)
+					36096:	if r1 == 0 goto +6 <LBB1501_2>
+					36097:	r2 = r10
+					36098:	r2 += -104
+				; 	if (SKBDUMP_CONFIG.skb_track && bpf_map_lookup_elem(&skb_address, &skb_addr))
+					36099:	r1 = 0 ll
+					36101:	call 1
+		GotoIndex ->		36102:	if r0 != 0 goto +11 <LBB1501_4>
+
+				0000000000046838 <LBB1501_2>:
 				; 	if (!pcap_filter((void *)(long)skb->data, (void *)(long)skb->data_end))
-					36096:	r1 = *(u32 *)(r6 + 80)
-					36097:	r2 = *(u32 *)(r6 + 76)
-	FilterIndex ->	36098:	if r2 >= r1 goto +71 <LBB1501_6>
+					36103:	r1 = *(u32 *)(r6 + 80)
+					36104:	r2 = *(u32 *)(r6 + 76)
+		FilterIndex ->		36105:	if r2 >= r1 goto +72 <LBB1501_7>
 	*/
-	FilterIndex = 12
-	GotoIndex   = 9
+	FilterIndex = 14
+	GotoIndex   = 11
 )
 
 type QueueBpfObjects struct {
@@ -85,8 +90,11 @@ func (o *QueueBpfObjects) setFilter(cbpfFilter []bpf.Instruction) (err error) {
 	return
 }
 
-func (o *QueueBpfObjects) Load(cbpfFilter []bpf.Instruction) (err error) {
-	if err = o.setFilter(cbpfFilter); err != nil {
+func (o *QueueBpfObjects) Load(opts internalbpf.LoadOptions) (err error) {
+	if err = errors.WithStack(o.spec.RewriteConstants(map[string]interface{}{"SKBDUMP_CONFIG": opts.BpfConfig})); err != nil {
+		return
+	}
+	if err = o.setFilter(opts.Filter); err != nil {
 		return
 	}
 	if err = errors.WithStack(o.spec.LoadAndAssign(o.objs, nil)); err != nil {

diff --git a/internal/bpf/queue/skbdump.c b/internal/bpf/queue/skbdump.c
@@ -39,7 +39,7 @@ void handle_skb(struct __sk_buff *skb, bool ingress)
 	bpf_skb_pull_data(skb, 0);
 
 	__u64 skb_addr = (__u64)(void *)skb;
-	if (bpf_map_lookup_elem(&skb_address, &skb_addr))
+	if (SKBDUMP_CONFIG.skb_track && bpf_map_lookup_elem(&skb_address, &skb_addr))
 		goto cont;
 
 	if (!pcap_filter((void *)(long)skb->data, (void *)(long)skb->data_end))

diff --git a/internal/bpf/queue/skbdump.pcap b/internal/bpf/queue/skbdump.pcap
diff --git a/internal/bpf/queue/skbdump.skb b/internal/bpf/queue/skbdump.skb
diff --git a/internal/bpf/queue/skbdump_bpfel_x86.o b/internal/bpf/queue/skbdump_bpfel_x86.o
diff --git a/main.go b/main.go
@@ -20,7 +20,7 @@ import (
 )
 
 func init() {
-	initConfig()
+	mustInitConfig()
 	if err := rlimit.RemoveMemlock(); err != nil {
 		log.Fatalf("Failed to remove rlimit memlock: %v", err)
 	}
@@ -37,7 +37,10 @@ func main() {
 	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
 	defer stop()
 
-	if err = bpfObjects.Load(bpf.MustPcapCompile(config.PcapFilterExp)); err != nil {
+	if err = bpfObjects.Load(bpf.LoadOptions{
+		Filter:    bpf.MustPcapCompile(config.PcapFilterExp),
+		BpfConfig: bpf.BpfConfig{SkbTrack: config.SkbTrack},
+	}); err != nil {
 		return
 	}