Testing a XDP program with BPF_PROG_TEST_RUN through Python

drop_ipv4_tcp_80.c

#include <uapi/linux/if_ether.h>
#include <uapi/linux/in.h>
#include <uapi/linux/ip.h>
#include <uapi/linux/tcp.h>

int drop_ipv4_tcp_80(struct xdp_md *ctx) {
  void *data = (void *)(long)ctx->data;
  void *data_end = (void *)(long)ctx->data_end;

  struct ethhdr *eth;
  struct iphdr *iph;
  struct tcphdr *th;

  eth = data;
  if ((void *)(eth + 1) > data_end)
    return XDP_DROP;

  if (eth->h_proto != htons(ETH_P_IP))
    return XDP_PASS;

  iph = (struct iphdr *)(eth + 1);
  if ((void *)(iph + 1) > data_end)
    return XDP_DROP;

  if (iph->ihl != 5 || iph->frag_off & htons(0x2000 | 0x1FFF))
    return XDP_PASS;

  if (iph->protocol != IPPROTO_TCP) {
    return XDP_PASS;
  }

  th = (struct tcphdr *)(iph + 1);
  if ((void *)(th + 1) > data_end)
    return XDP_DROP;

  if (th->dest == htons(80))
    return XDP_DROP;

  return XDP_PASS;
}

test_xdp_prog.py

import unittest

from bcc import BPF, libbcc
import ctypes
from scapy.all import *


class PacketDropTestCase(unittest.TestCase):
    bpf = None
    func = None

    DATA_OUT_LEN = 1514

    def _run_test(self, data, data_out_expect, retval_expect, repeat=1):
        size = len(data)
        data = ctypes.create_string_buffer(raw(data), size)
        data_out = ctypes.create_string_buffer(self.DATA_OUT_LEN)
        size_out = ctypes.c_uint32()
        retval = ctypes.c_uint32()
        duration = ctypes.c_uint32()

        ret = libbcc.lib.bpf_prog_test_run(self.func.fd, repeat,
                                           ctypes.byref(data), size,
                                           ctypes.byref(data_out),
                                           ctypes.byref(size_out),
                                           ctypes.byref(retval),
                                           ctypes.byref(duration))
        self.assertEqual(ret, 0)

        self.assertEqual(retval.value, retval_expect)
        if data_out_expect:
            self.assertEqual(data_out[:size_out.value], raw(data_out_expect))

    def setUp(self):
        self.bpf = BPF(src_file=b"drop_ipv4_tcp_80.c")
        self.func = self.bpf.load_func(b"drop_ipv4_tcp_80", BPF.XDP)

    def test_ipv4_tcp_80(self):
        packet_in = Ether() / IP() / TCP(dport=80)
        self._run_test(packet_in, None, BPF.XDP_DROP)

    def test_ipv4_udp_80(self):
        packet_in = Ether() / IP() / UDP(dport=80)
        self._run_test(packet_in, packet_in, BPF.XDP_PASS)

    def test_ipv4_tcp_443(self):
        packet_in = Ether() / IP() / TCP(dport=443)
        self._run_test(packet_in, packet_in, BPF.XDP_PASS)

    def test_ipv6_tcp_80(self):
        packet_in = Ether() / IPv6() / TCP(dport=443)
        self._run_test(packet_in, packet_in, BPF.XDP_PASS)


if __name__ == '__main__':
    unittest.main()