diff --git a/ebpfcat/ebpf.py b/ebpfcat/ebpf.py index c21bd3e651f777cb0070913e150f7555fe40f8db..da2efbf14fe100938e83f9dd685c7c768470eb9c 100644 --- a/ebpfcat/ebpf.py +++ b/ebpfcat/ebpf.py @@ -23,6 +23,7 @@ from struct import pack, unpack, calcsize from enum import Enum from . import bpf +from .util import sub Instruction = namedtuple("Instruction", ["opcode", "dst", "src", "off", "imm"]) @@ -1264,7 +1265,7 @@ class EBPF: def assemble(self): """return the assembled program""" - self.program() + sub(EBPF, self).program() return b"".join( pack("<BBHI", i.opcode.value, i.dst | i.src << 4, i.off % 0x10000, i.imm % 0x100000000) diff --git a/ebpfcat/ebpf.rst b/ebpfcat/ebpf.rst index 570e7c869768deced196013571e067ae9fbd8e57..e7711836753388dbd04135d7fa7bda86e4d3bad7 100644 --- a/ebpfcat/ebpf.rst +++ b/ebpfcat/ebpf.rst @@ -126,6 +126,21 @@ packets:: self.count += 1 self.exit(XDPExitCode.PASS) +as a simplification, if the class attribute ``minimumPacketSize`` is set, +the ``program`` is called within a ``with`` statement like above, and all +the packet variables appear as variables of the object. The class +attribute ``defaultExitCode`` then gives the exit code in case the packet +is too small (by default ``XDPExitCode.PASS``). So the above example becomes:: + + class Program(XDP): + minimumPacketSize = 16 + userspace = HashMap() + count = userspace.globalVar() + + def program(self): + with self.pH[12] == 8: + self.count += 1 + Maps ---- diff --git a/ebpfcat/ebpf_test.py b/ebpfcat/ebpf_test.py index c61bb959b4ec94b3d99f0f7bca470aeeaa9bf53b..4dea9657d40bb7d6c36d44a53c47195a2fcbb207 100644 --- a/ebpfcat/ebpf_test.py +++ b/ebpfcat/ebpf_test.py @@ -899,6 +899,27 @@ class Tests(TestCase): Instruction(opcode=O.JMP, dst=0, src=0, off=1, imm=0), Instruction(opcode=O.MOV+O.LONG, dst=3, src=0, off=0, imm=77)]) + def test_xdp_minsize(self): + class P(XDP): + minimumPacketSize = 100 + + def program(self): + self.r3 = self.pH[22] + + p = P(license="GPL") + p.assemble() + self.assertEqual(p.opcodes, [ + Instruction(opcode=O.W+O.LD, dst=9, src=1, off=0, imm=0), + Instruction(opcode=O.W+O.LD, dst=0, src=1, off=4, imm=0), + Instruction(opcode=O.W+O.LD, dst=2, src=1, off=0, imm=0), + Instruction(opcode=O.LONG+O.ADD, dst=2, src=0, off=0, imm=100), + Instruction(opcode=O.JLE+O.REG, dst=0, src=2, off=1, imm=0), + Instruction(opcode=O.REG+O.LD, dst=3, src=9, off=22, imm=0), + Instruction(opcode=O.LONG+O.MOV, dst=0, src=0, off=0, imm=2), + Instruction(opcode=O.EXIT, dst=0, src=0, off=0, imm=0), + ]) + + class KernelTests(TestCase): def test_hashmap(self): class Global(EBPF): diff --git a/ebpfcat/util.py b/ebpfcat/util.py new file mode 100644 index 0000000000000000000000000000000000000000..840e9ec52d6ffaff6fbc60df6a3ce7bc4ac98816 --- /dev/null +++ b/ebpfcat/util.py @@ -0,0 +1,43 @@ +from itertools import chain + + +class sub: + def __init__(self, cls, base, default=False): + self.cls = cls + self.base = base + + def __getattr__(self, name): + mro = self.base.__class__.__mro__[::-1] + i = mro.index(self.cls) + for cls in chain(mro[i + 1 :], mro[:i + 1]): + func = cls.__dict__.get(name) + if func is not None: + return func.__get__(self.base, cls) + raise AttributeError(f"'sub' object has no attribute '{name}'") + + +if __name__ == "__main__": + class A: + def g(self): + print("A.f") + + class B(A): + def f(self): + print("B.f") + + class C(A): + def f(self): + print("C.f") + + class D(C, B): + def f(self): + print("D.f") + + + b = D() + print(D.__mro__) + sub(A, b).f() + sub(B, b).f() + sub(C, b).f() + sub(D, b).f() + diff --git a/ebpfcat/xdp.py b/ebpfcat/xdp.py index dee0e1e7ee1b3187f1c06405b71b20c5e85e360e..db17388715107a27e14d5c1c68a2abeb13c47e0f 100644 --- a/ebpfcat/xdp.py +++ b/ebpfcat/xdp.py @@ -26,6 +26,7 @@ from struct import pack, unpack from .ebpf import EBPF from .bpf import ProgType +from .util import sub class XDPExitCode(Enum): @@ -153,11 +154,26 @@ class PacketSize: class XDP(EBPF): """the base class for XDP programs""" + minimumPacketSize = None + defaultExitCode = XDPExitCode.PASS + def __init__(self, **kwargs): super().__init__(prog_type=ProgType.XDP, **kwargs) self.packetSize = PacketSize(self) + def program(self): + if self.minimumPacketSize is None: + sub(XDP, self).program() + else: + with self.packetSize > self.minimumPacketSize as packet: + self.pB = packet.pB + self.pH = packet.pH + self.pI = packet.pI + self.pQ = packet.pQ + sub(XDP, self).program() + self.exit(self.defaultExitCode) + async def _netlink(self, ifindex, fd, flags): future = Future() transport, proto = await get_event_loop().create_datagram_endpoint(