Source code for ebpfcat.hashmap

# ebpfcat, A Python-based EBPF generator and EtherCAT master
# Copyright (C) 2021 Martin Teichmann <martin.teichmann@gmail.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

"""\
:mod:`!ebpfcat.hashmap` --- eBPF hash maps
==========================================

This module makes eBPF hash maps available, used for dictionary or global
variables.
"""

from collections.abc import MutableMapping
from contextlib import contextmanager
from struct import pack, unpack, unpack

from .ebpf import AssembleError, Expression, Opcode, Map, FuncId
from .bpf import (
    MapType, UpdateFlags, create_map, delete_elem, get_next_key, lookup_elem,
    lookup_and_delete_elem, update_elem)


__all__ = ["HashMap", "Dict"]


class HashGlobalVar(Expression):
    def __init__(self, ebpf, count, fmt):
        self.ebpf = ebpf
        self.count = count
        self.fmt = fmt
        self.signed = fmt.islower()
        self.fixed = fmt == "x"

    @contextmanager
    def get_address(self, dst, long, force=False):
        with self.ebpf.save_registers([i for i in range(6) if i != dst]), \
                self.ebpf.get_stack(4) as stack:
            self.ebpf.append(Opcode.ST, 10, 0, stack, self.count)
            self.ebpf.r1 = self.ebpf.get_fd(self.fd)
            self.ebpf.r2 = self.ebpf.r10 + stack
            self.ebpf.call(FuncId.map_lookup_elem)
            with self.ebpf.r0 == 0:
                self.ebpf.exit()
            if dst != 0 and force:
                self.ebpf.append(Opcode.MOV + Opcode.LONG + Opcode.REG, dst,
                                 0, 0, 0)
            else:
                dst = 0
        yield dst, self.fmt


class HashGlobalVarDesc:
    def __init__(self, count, fmt, default=0):
        self.count = count
        self.fmt = fmt
        self.default = default

    def __get__(self, instance, owner):
        if instance is None:
            return self
        if instance.loaded:
            fd = instance.__dict__[self.name].fd
            return lookup_elem(fd, pack("B", self.count), self.fmt)
        ret = instance.__dict__.get(self.name, None)
        if ret is None:
            ret = HashGlobalVar(instance, self.count, self.fmt)
            instance.__dict__[self.name] = ret
        return ret

    def __set_name__(self, owner, name):
        self.name = name

    def __set__(self, ebpf, value):
        if ebpf.loaded:
            fd = ebpf.__dict__[self.name].fd
            update_elem(fd, pack("B", self.count),
                        pack("q" if self.fmt.islower() else "Q", value))
            return
        with ebpf.save_registers([3]):
            with value.get_address(3, True, True):
                with ebpf.save_registers([0, 1, 2, 4, 5]), \
                        ebpf.get_stack(4) as stack:
                    ebpf.r1 = ebpf.get_fd(ebpf.__dict__[self.name].fd)
                    ebpf.append(Opcode.ST, 10, 0, stack, self.count)
                    ebpf.r2 = ebpf.r10 + stack
                    ebpf.r4 = 0
                    ebpf.call(FuncId.map_update_elem)


[docs] class HashMap(Map): """global variables as a eBPF hash map""" count = 0 def __init__(self): self.vars = [] def globalVar(self, fmt="I", default=0): self.count += 1 ret = HashGlobalVarDesc(self.count, fmt, default) self.vars.append(ret) return ret def init(self, ebpf, fd): if fd is None: fd = create_map(MapType.HASH, 1, 8, self.count) for v in self.vars: getattr(ebpf, v.name).fd = fd def load(self, ebpf): for v in self.vars: setattr(ebpf, v.name, ebpf.__class__.__dict__[v.name].default)
class TheDict(MutableMapping): def __init__(self, ht, ebpf, fd): self.key = ht.Key() self.key.addr_offset = ht.key_offset self.key.data = None self.key.ebpf = ebpf self.value = ht.Value() self.value.addr_offset = ht.value_offset self.value.data = None self.value.ebpf = ebpf self.ebpf = ebpf self.fd = fd def __setitem__(self, key, value): assert isinstance(key, type(self.key)) assert isinstance(value, type(self.value)) update_elem(self.fd, key.data, value.data) def __getitem__(self, key): assert isinstance(key, type(self.key)) ret = type(self.value)() ret.data = lookup_elem(self.fd, key.data, self.value.stack) return ret __marker = object() def pop(self, key, default=__marker): assert isinstance(key, type(self.key)) ret = type(self.value)() try: ret.data = lookup_and_delete_elem(self.fd, key.data, self.value.stack) except KeyError: if default is self.__marker: raise return default return ret def __delitem__(self, key): assert isinstance(key, type(self.key)) delete_elem(self.fd, key.data) def __len__(self): """there is no way to actually tell how many elements are in a hash map, but we need a __len__ method for MutableMapping. We just raise a TypeError, so that list(table) still works. """ raise TypeError def __iter__(self): current = get_next_key(self.fd, self.key.stack) while True: ret = type(self.key)() ret.data = current yield ret try: current = get_next_key(self.fd, current) except StopIteration: return def update(self, flags=UpdateFlags.ANY): assert isinstance(flags, UpdateFlags) ebpf = self.ebpf with ebpf.save_registers([1, 2, 3, 4, 5]): ebpf.r1 = ebpf.get_fd(self.fd) ebpf.r2 = ebpf.r10 + self.key.addr_offset ebpf.r3 = ebpf.r10 + self.value.addr_offset ebpf.r4 = flags.value ebpf.call(FuncId.map_update_elem) @contextmanager def lookup(self): ebpf = self.ebpf with ebpf.save_registers([1, 2, 3, 4, 5]): ebpf.r1 = ebpf.get_fd(self.fd) ebpf.r2 = ebpf.r10 + self.key.addr_offset ebpf.call(FuncId.map_lookup_elem) with ebpf.r0 != 0 as Else: value = type(self.value)() value.addr_offset = 0 value.base_register = 0 value.data = None value.ebpf = ebpf yield value, Else
[docs] class Dict(Map): """A dictionary, implemented using a bpf hash map This is a lookup table similar to a Python :class:`dict`. Both key and value are subclasses of :class:`~ebpfcat.ebpf.Structure`:: class Key(Structure): some_key = Member('I') class Value(Structure): some_value = Member('q') A dictionary can then be declared in an EBPF program:: class Program(EBPF): table = Dict(key=Key, value=Value) On the Python side, the usage is like a Python :class:`dict`:: e = Program() e.load() k = Key() k.some_key = 3 v = Value() v.some_value = 7 # update a value e.table[k] = v # look up a value: v = e.table[k] On the EBPF side, things are a bit more complicated. We always keep a key on the local stack for lookups, as well as a value for updates. Possible errors are returned in register 0, which one may immediately after:: def program(self): self.table.key.some_key = 3 self.table.update() with self.r0 != 0: # do some error handling, if needed Lookups need to be enclosed in a ``with`` statement, during which the looked up value will be valid. One may modify the looked up value, which will indeed change the value in the dictionary. In case of an error (e.g., the looked up key does not exist), we skip over the entire ``with`` clause. The lookup function also returns an ``Else`` handler for later error handling:: def program(self): self.table.key.some_key = 7 with self.table.lookup() as value, Else: value.some_value = 3 # this changes the value in the Dict with Else: # do some error handling :param size: the maximum number of elements in the Dict :param lru: if set to ``True``, the dictionary behaves as last-recently-used, so if it is full, it will just remove the oldest keys instead of giving an error. """ def __init__(self, key, value, size=31, lru=False): self.Key = key self.Value = value if lru: self.mapType = MapType.LRU_HASH else: self.mapType = MapType.HASH self.size = size def __set_name__(self, owner, name): owner.stack -= self.Key.stack owner.stack &= -8 self.key_offset = owner.stack owner.stack -= self.Value.stack owner.stack &= -8 self.value_offset = owner.stack self.name = name def init(self, ebpf, fd): if fd is None: fd = create_map(self.mapType, self.Key.stack, self.Value.stack, self.size) setattr(ebpf, self.name, TheDict(self, ebpf, fd))