an example tracer for h2o based off EBPF

picotracer.py

#!/usr/bin/env python
from __future__ import print_function
from bcc import BPF, USDT
import ctypes as ct
import sys
import signal
import sys
import os
import socket
import time

if len(sys.argv) == 1:
    print("""USAGE: picotracer [H2O_PID] [PROBE_NAME] [MISC_FILTER]
    e.g. sudo picotracer.py 1234
         sudo picotracer.py RxHeader
         sudo picotracer.py NewConnH1 10.0.2.15""")
    exit(1)

bpf_text = """
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>
BPF_PERF_OUTPUT(h2otrace);
struct xchange_t {
    char flag[30];
    long connids;
    long connidu;
    int fd;
    char payload[150];
 };
"""

func_txt = """
int do_{}(struct pt_regs *ctx) {{
    uint64_t addr;
    struct xchange_t data = {{}};
    bpf_usdt_readarg(1, ctx, &data.connids);
    bpf_usdt_readarg(2, ctx, &data.connidu);
    bpf_usdt_readarg(3, ctx, &data.fd);
    bpf_usdt_readarg(4, ctx, &addr);
    bpf_probe_read(&data.payload[0], 150, (void *)addr);
    sprintf(data.flag, "{}  ");
    h2otrace.perf_submit(ctx, &data, sizeof(data));
    return 0;
}};
"""

def toU64(n):
    return n.to_bytes(8, byteorder='little', signed=False)

# conn socket to h2o
TCP_IP = '127.0.0.1'
TCP_PORT = 4321

# probes available -- order needs to be identical as in tracing.h
PROBES = ["SSLNew", "NewConnH1", "NewConnH2", "NewReq", "NewRes", "TxHeader", "RxHeader", "Proxy", "ProxyNewReq", "ProxyNewRes", "ProxyTxHdr", "ProxyRxHdr"]
H2O_MSG = 0x0000000000000000
H2O_MSG_CLOSE = 1<<63

pid = int(sys.argv[1])
probe_wanted = str(sys.argv[2]) if len(sys.argv) > 2 else False
probe_filter = str(sys.argv[3]) if len(sys.argv) > 3 else False

# ebpf <> python exchange structure
class Xchange(ct.Structure):
    _fields_ = [("flag", ct.c_char * 30),
                ("connids", ct.c_long),
                ("connidu", ct.c_long),
                ("fd", ct.c_int),
                ("payload", ct.c_char * 150)]

# enable USDT probe from given PID
u = USDT(pid=pid)

for i in range(len(PROBES)):
    # attach probe if needed
    probe = PROBES[i]
    if probe_wanted and probe_wanted != probe: continue
    # add function text to ebpf program. push bit to H2O_MSG
    bpf_text += func_txt.format(probe, probe)
    H2O_MSG |= 1<<i
    u.enable_probe(probe=probe, fn_name="do_"+probe)

# setup socket connection to h2o
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect((TCP_IP, TCP_PORT))

# initialize BPF & send message to h2o
b = BPF(text=bpf_text, usdt_contexts=[u])
s.send(toU64(H2O_MSG))
print("\r\n>> connstart \t\t fd    probe \t\t output")

# reconstruct payload, filter if needed, and print
def print_event(cpu, data, size):
    d = ct.cast(data, ct.POINTER(Xchange)).contents
    payload = d.payload.decode("utf-8")
    if probe_filter and probe_filter not in payload: return
    print(">> %ld%ld\t %-5d %s  \t %s" % (d.connids, d.connidu, d.fd, d.flag.decode("utf-8"), payload))

# start and print
b["h2otrace"].open_perf_buffer(print_event)
while 1:
    try:
        b.perf_buffer_poll()
    except KeyboardInterrupt:
        s.send(toU64(H2O_MSG_CLOSE | H2O_MSG))
        time.sleep(1)
        exit()