Script to trace packet drops from Linux's xfrm layer

packetdrop-xfrm.py

#!/usr/bin/python
# Copyright (c) Isovalent, 2021.
# Licensed under the Apache License, Version 2.0 (the "License")

from bcc import BPF
import ctypes as ct
import time
import socket
import struct

b = BPF(text="""
#ifndef KBUILD_MODNAME
# define KBUILD_MODNAME "foo"
#endif
#include <uapi/linux/ptrace.h>
#include <uapi/linux/ip.h>
#include <uapi/linux/tcp.h>
#include <uapi/linux/if_ether.h>
#include <linux/netdevice.h>
#include <net/sock.h>
#include <net/xfrm.h>

struct data_t {
    u64 stack_id;
    int ret;
    int dir;
    u32 mark;
    u16 proto;
    u32 saddr;
    u32 daddr;
    u16 sport;
    u16 dport;
    u32 seq;
    u32 ack_seq;
    struct sk_buff *skb;
    struct sec_path *sp;
    int sp_len;
    u32 sa_mark_val;
    u32 sa_spi;
    u32 sa_reqid;
    char device[IFNAMSIZ];
};

BPF_HASH(args, u64, struct data_t);
BPF_STACK_TRACE(stack_traces, 4096);
BPF_PERF_OUTPUT(events);

static inline bool
sec_path_exist(const struct sk_buff *skb) {
    return skb->active_extensions & (1 << SKB_EXT_SEC_PATH);
}

static inline struct sec_path *
get_sec_path(const struct sk_buff *skb) {
    if (sec_path_exist(skb)) {
        struct skb_ext *ext = skb->extensions;
        u64 offset = ext->offset[SKB_EXT_SEC_PATH];
        return (void *)ext + (offset << 3);
    }
    return NULL;
}

int kprobe____xfrm_policy_check(struct pt_regs *ctx, void *sk, int dir,
                                struct sk_buff *skb) {
    struct data_t data = {};
    struct net_device *dev;
    struct iphdr *iph = (struct iphdr *)(skb->head + skb->network_header);
    u64 pid = bpf_get_current_pid_tgid();
    
    data.dir = dir;
    data.mark = skb->mark;
    
    data.proto = skb->protocol;
    data.proto = ntohs(data.proto);
    if (skb->protocol != htons(ETH_P_IP)) {
        data.saddr = 42;
        data.daddr = 43;
    } else {
        data.saddr = iph->saddr;
        data.saddr = ntohl(data.saddr);
        data.daddr = iph->daddr;
        data.daddr = ntohl(data.daddr);
    }

    struct tcphdr *tcp = (struct tcphdr *)(skb->head + skb->transport_header);
    data.sport = tcp->source;
    data.sport = ntohs(data.sport);
    data.dport = tcp->dest;
    data.dport = ntohs(data.dport);
    data.seq = tcp->seq;
    data.seq = ntohl(data.seq);
    data.ack_seq = tcp->ack_seq;
    data.ack_seq = ntohl(data.ack_seq);

    data.skb = skb;
    data.sp = get_sec_path(skb);
    if (data.sp) {
        data.sp_len = data.sp->len;
        if (data.sp_len > 0) {
            struct xfrm_state *state = data.sp->xvec[0];
            data.sa_mark_val = state->mark.v;
            data.sa_spi = state->id.spi;
            data.sa_reqid = state->props.reqid;
        }
    }

    dev = skb->dev;
    bpf_probe_read(&data.device, sizeof(data.device), &dev->name);
    
    data.stack_id = stack_traces.get_stackid(ctx, 0);
    if (args.update(&pid, &data))
        bpf_trace_printk("Update failed!\\n");
    return 0;
}

int kretprobe____xfrm_policy_check(struct pt_regs *ctx) {
    struct data_t *datap;
    int ret = PT_REGS_RC(ctx);
    u64 pid = bpf_get_current_pid_tgid();
    if (1) {
        datap = args.lookup(&pid);
        if (datap) {
            struct sec_path *sp = get_sec_path(datap->skb);
            if (datap->sp != sp) {
                datap->sp = (struct sec_path *)42;
            } else {
                datap->sp = sp;
            }
            datap->ret = ret;
            events.perf_submit(ctx, datap, sizeof(*datap));
            args.delete(&pid);
        } else {
            struct data_t data = {};
            data.ret = -1;
            events.perf_submit(ctx, &data, sizeof(data));
        }
    }
    return 0;
}
""")

IFNAMSIZ = 16
class Data(ct.Structure):
    _fields_ = [
        ("stack_id", ct.c_ulonglong),
        ("ret", ct.c_int),
        ("dir", ct.c_int),
        ("mark", ct.c_uint),
        ("proto", ct.c_uint16),
        ("saddr", ct.c_uint),
        ("daddr", ct.c_uint),
        ("sport", ct.c_uint16),
        ("dport", ct.c_uint16),
        ("seq", ct.c_uint),
        ("ack_seq", ct.c_uint),
        ("skb", ct.c_ulonglong),
        ("sp", ct.c_ulonglong),
        ("sp_len", ct.c_int),
        ("sa_mark_val", ct.c_uint),
        ("sa_spi", ct.c_uint),
        ("sa_reqid", ct.c_uint),
        ("device", ct.c_char * IFNAMSIZ),
    ]

stack_traces = b.get_table("stack_traces")
start_ts = time.time()

print("%-18s  %-6s %-6s %-6s %-12s %-6s" %
      ("TIME(s)", "RET", "DIR", "MARK", "DEVICE", "SP"))

def print_event(cpu, data, size):
    event = ct.cast(data, ct.POINTER(Data)).contents
    ts = time.time() - start_ts
    saddr = socket.inet_ntoa(struct.pack("!I", event.saddr))
    daddr = socket.inet_ntoa(struct.pack("!I", event.daddr))
    print("%-18.9f %-6d %-6d %-6x %-12s %-6x %s:%d -> %s:%d (sp_len: %d, mark: %x, spi: %d, reqid: %d)" %
          (ts, event.ret, event.dir, event.mark, event.device.decode(), event.sp, saddr, event.sport, daddr, event.dport, event.sp_len, event.sa_mark_val, event.sa_spi, event.sa_reqid))

    # for addr in stack_traces.walk(event.stack_id):
    #     sym = b.ksym(addr, show_offset=True)
    #     print("\t%s" % sym)

    # print()

b["events"].open_perf_buffer(print_event)
while 1:
    try:
        b.perf_buffer_poll()
    except KeyboardInterrupt:
        exit()