Merge pull request #9 from jschwinger233/gray/skb-address

add --skb-track to track NAT-ed skb

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(), " ")
 

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)

internal/bpf/headers/skbdump.h

@@ -6,14 +6,18 @@
 #define MAX_QUEUE_SIZE 10000
 #endif
 
-#ifndef ETH_HLEN
-#define ETH_HLEN 14
-#endif
-
 #ifndef MAX_DATA_SIZE
 #define MAX_DATA_SIZE 1500
 #endif
 
+#ifndef MAX_TRACK_SIZE
+#define MAX_TRACK_SIZE 1000
+#endif
+
+#ifndef ETH_HLEN
+#define ETH_HLEN 14
+#endif
+
 #ifndef ETH_P_IP
 #define ETH_P_IP 0x0800 /* Internet Protocol packet	*/
 #endif
@@ -57,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

internal/bpf/inject.go

@@ -0,0 +1,13 @@
+package bpf
+
+import (
+	"github.com/cilium/ebpf/asm"
+)
+
+func InjectInstructions(old, new []asm.Instruction, start, end int, gotoIndices []int) []asm.Instruction {
+	old = append(old[:start], append(new, old[end:]...)...)
+	for _, gotoIdx := range gotoIndices {
+		old[gotoIdx].Offset += int16(len(new) - (end - start))
+	}
+	return old
+}

internal/bpf/pcap_filter.go

@@ -2,8 +2,11 @@ package bpf
 
 import (
 	"log"
+	"net"
 	"unsafe"
 
+	"github.com/pkg/errors"
+	"github.com/vishvananda/netlink"
 	"golang.org/x/net/bpf"
 )
 
@@ -22,11 +25,30 @@ const (
 	MAXIMUM_SNAPLEN          = 262144
 )
 
+func mustFindUpDevice() string {
+	links, err := netlink.LinkList()
+	if err != nil {
+		log.Fatalf("%+v", errors.WithStack(err))
+	}
+
+	for _, link := range links {
+		if link.Attrs().Flags&net.FlagUp != 0 {
+			return link.Attrs().Name
+		}
+	}
+	log.Fatal("cannot find an up device to call pcap_open_live")
+	return ""
+}
+
 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))
 
-	cptr := C.pcap_open_live(C.CString("lo"), C.int(MAXIMUM_SNAPLEN), C.int(0), C.int(0), buf)
+	cptr := C.pcap_open_live(C.CString(mustFindUpDevice()), C.int(MAXIMUM_SNAPLEN), C.int(0), C.int(0), buf)
 	if cptr == nil {
 		log.Fatalf("failed to pcap_open_live: %+v\n", C.GoString(buf))
 	}

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 {

internal/bpf/queue/bpf.go

@@ -14,6 +14,37 @@ 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 (
+	/*
+				00000000000467c8 <on_ingress>:
+				; {
+					36089:	r6 = r1
+				; 	bpf_skb_pull_data(skb, 0);
+					36090:	r2 = 0
+					36091:	call 39
+				; 	__u64 skb_addr = (__u64)(void *)skb;
+					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))
+					36103:	r1 = *(u32 *)(r6 + 80)
+					36104:	r2 = *(u32 *)(r6 + 76)
+		FilterIndex ->		36105:	if r2 >= r1 goto +72 <LBB1501_7>
+	*/
+	FilterIndex = 14
+	GotoIndex   = 11
+)
+
 type QueueBpfObjects struct {
 	spec *ebpf.CollectionSpec
 	objs *SkbdumpObjects
@@ -53,18 +84,17 @@ func (o *QueueBpfObjects) setFilter(cbpfFilter []bpf.Instruction) (err error) {
 		asm.Return().WithSymbol("return"),           // else return TC_ACT_OK
 		asm.Mov.Imm(asm.R0, 0).WithSymbol("continue"),
 	)
-	ingressInsts := o.spec.Programs["on_ingress"].Instructions
-	ingressInsts = append(ingressInsts[:6], append(ebpfFilter, ingressInsts[7:]...)...)
-	o.spec.Programs["on_ingress"].Instructions = ingressInsts
 
-	egressInsts := o.spec.Programs["on_egress"].Instructions
-	egressInsts = append(egressInsts[:6], append(ebpfFilter, egressInsts[7:]...)...)
-	o.spec.Programs["on_egress"].Instructions = egressInsts
+	o.spec.Programs["on_ingress"].Instructions = internalbpf.InjectInstructions(o.spec.Programs["on_ingress"].Instructions, ebpfFilter, FilterIndex, FilterIndex+1, []int{GotoIndex})
+	o.spec.Programs["on_egress"].Instructions = internalbpf.InjectInstructions(o.spec.Programs["on_egress"].Instructions, ebpfFilter, FilterIndex, FilterIndex+1, []int{GotoIndex})
 	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 {

internal/bpf/queue/skbdump.c

@@ -5,6 +5,7 @@
 #include "skbdump.h"
 #include "skb_data.h"
 
+const static bool TRUE = true;
 
 char __license[] SEC("license") = "Dual MIT/GPL";
 
@@ -18,6 +19,13 @@ struct bpf_map_def SEC("maps") meta_queue = {
 	.max_entries = MAX_QUEUE_SIZE,
 };
 
+struct bpf_map_def SEC("maps") skb_address = {
+	.type = BPF_MAP_TYPE_HASH,
+	.key_size = sizeof(__u64),
+	.value_size = sizeof(bool),
+	.max_entries = MAX_TRACK_SIZE,
+};
+
 static __always_inline
 bool pcap_filter(void *data, void* data_end)
 {
@@ -27,12 +35,19 @@ bool pcap_filter(void *data, void* data_end)
 static __always_inline
 void handle_skb(struct __sk_buff *skb, bool ingress)
 {
+	struct skb_meta meta;
 	bpf_skb_pull_data(skb, 0);
 
+	__u64 skb_addr = (__u64)(void *)skb;
+	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))
 		return;
 
-	struct skb_meta meta = {};
+	bpf_map_update_elem(&skb_address, &skb_addr, &TRUE, BPF_ANY);
+
+cont:
 	__builtin_memset(&meta, 0, sizeof(meta));
 	meta.is_ingress = ingress;
 	meta.time_ns = bpf_ktime_get_ns();

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
 	}