xdp h2o connection tracing example

xdp.py

#!/usr/bin/python
#
# simple XDP eBPF program to flag connections using an eBPF map
# this example adds connections to ::1 ::1 to a an eBPF map
# depends on `bpftool` for now, as bcc is missing map pinning
#
# the map is spinned by the XDP program, and this python handler
# will call `bpftool` to pin the map to a path to allow system-wide
# access. the path will be removed when the program quits.
#
# example usage
#  1. start h2o listening on ::1 (optionally adding logging in h2o_trace_check_map)
#  2. attach usdt probe for connection tracing (see cmd below for sample)
#  3. start map monitoring (optional, below)
#  4. start this script, attaching to lo interface
#  5. curl h2o on ::1
#  6. observe map key/values and optionally, logging
#
# useful monitoring commands :
#  enable h2o conn tracing:
#     % sudo ./trace.py -T -p `pgrep -o h2o` u:/home/joe/code/h2o/h2o:h2o_conn_tracing
#  list maps and dump:
#     % watch sudo bpftool map show
#     % watch sudo bpftool map dump pinned /sys/fs/bpf/h2o_map

from bcc import BPF
import time
import sys
import ctypes as ct
import json
import subprocess

if len(sys.argv) != 2:
    print("Usage: {0} <in ifdev>".format(sys.argv[0]))
    print("e.g.: {0} eth0\n".format(sys.argv[0]))
    exit(1)

MAP_NAME = "h2o_map"
PATH_MAP = " /sys/fs/bpf/" + MAP_NAME

def run(cmd):
    return subprocess.check_output(cmd, shell=True).decode("utf-8")

# TODO: integrate this into BCC
def pin_map():
    for m in json.loads(run('sudo bpftool -j map list')):
        if 'name' in m and m['name'] == MAP_NAME:
            print('pinning ' + str(m['id']) + ' to path' + PATH_MAP)
            run('bpftool -j map pin id ' + str(m['id']) + PATH_MAP)


# load BPF program
b = BPF(text = """
#define KBUILD_MODNAME "foo"
#include <linux/in.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/unistd.h>
#include <linux/bpf.h>

struct keyType {
    struct {
        uint8_t ip[16];
        uint16_t port;
    } source;
    struct {
        uint8_t ip[16];
        uint16_t port;
    } destination;
    uint8_t family;
    uint8_t protocol;
};

#define htons(x) ((__be16)___constant_swab16((x)))
#define htonl(x) ((__be32)___constant_swab32((x)))

#define is_v4 (eth->h_proto == htons(ETH_P_IP))
#define is_v6 (eth->h_proto == htons(ETH_P_IPV6))
#define iph_sze (is_v4 ? sizeof(struct iphdr) : sizeof(struct ipv6hdr))
#define is_udp (prot == IPPROTO_UDP)
#define is_tcp (prot == IPPROTO_TCP)

#define check_eth() ({ if (data + eth_sze > data_end) return XDP_PASS; })
#define check_hdr() ({ if (data + eth_sze + iph_sze > data_end) return XDP_PASS; })
#define check_udp_hdr() ({ if (udph + 1 > (struct udphdr *)data_end) return XDP_PASS; })
#define check_tcp_hdr() ({ if (tcph + 1 > (struct tcphdr *)data_end) return XDP_PASS; })

BPF_TABLE("lru_percpu_hash", struct keyType, char, h2o_map, 1000);

int xdp_block(struct xdp_md *ctx) {
    void *data = (void *)(long)ctx->data;
    void *data_end = (void *)(long)ctx->data_end;
    struct ethhdr *eth = data;
    uint64_t eth_sze = sizeof(*eth);
    check_eth();

    struct keyType key;
    memset(&key, 0, sizeof(key));
    char val = 1;
    char record = 0;
    unsigned char prot = 0;

    // extract family - protocol
    key.family = is_v6 ? 6 : 4;
    // extract ip (only tcp/udp)
    if (is_v4) {
        struct iphdr *iph = data + eth_sze;
        check_hdr();
        prot = iph->protocol;
        memcpy(&key.destination.ip[0], &iph->saddr, sizeof(iph->saddr));
        memcpy(&key.source.ip[0], &iph->daddr, sizeof(iph->daddr));
    } else if (is_v6) {
        struct ipv6hdr *ip6h = data + eth_sze;
        check_hdr();
        prot = ip6h->nexthdr;
        memcpy(&key.destination.ip[0], ip6h->saddr.s6_addr, sizeof(ip6h->saddr.s6_addr));
        memcpy(&key.source.ip[0], ip6h->daddr.s6_addr, sizeof(ip6h->daddr.s6_addr));

        // only record ::1 ::1 for now
        if (ip6h->saddr.s6_addr[15] == ip6h->daddr.s6_addr[15] == 1) record = 1;
    }

    // extract port (only tcp/udp)
    if (is_udp) {
        struct udphdr *udph = data + eth_sze + iph_sze;
        check_udp_hdr();
        key.destination.port = udph->source;
        key.source.port = udph->dest;
        key.protocol = SOCK_DGRAM;
    } else if (is_tcp) {
        struct tcphdr *tcph = data + eth_sze + iph_sze;
        check_tcp_hdr();
        key.destination.port = tcph->source;
        key.source.port = tcph->dest;
        key.protocol = SOCK_STREAM;
    }

    if (record) h2o_map.insert(&key, &val);

    return XDP_PASS;
}
""", cflags=["-w"], debug=0)

in_if = sys.argv[1]
in_fn = b.load_func("xdp_block", BPF.XDP)
b.attach_xdp(in_if, in_fn, 0)

time.sleep(1)
pin_map()

print("XDP program running, hit CTRL+C to stop")
while 1:
    try:
        time.sleep(1)
    except KeyboardInterrupt:
        break

print("remove xdp program from device...\r\ncleanup map and wait for h2o to disconnect...")
b.remove_xdp(in_if, 0)
run("sudo rm -rf" + PATH_MAP)

fds = "1"
while fds != "0":
    fds = run('sudo ls -la /proc/`pgrep -o h2o`/fd  | grep "bpf-map" | wc -l').rstrip()
    print("{} fds left...".format(fds))
    time.sleep(1)

print("cleanup done!")