diff --git a/ebpfcat/ebpf.py b/ebpfcat/ebpf.py index 6fd6e5959064ccf639f23a3ac81c18e2e4df1fe3..a615bcfe783908e2a55beef1f4cca4bb99e0311e 100644 --- a/ebpfcat/ebpf.py +++ b/ebpfcat/ebpf.py @@ -902,6 +902,21 @@ class ktime(Expression): yield dst, True, False +class prandom(Expression): + """a function that returns the current ktime in ns""" + def __init__(self, ebpf): + self.ebpf = ebpf + + @contextmanager + def calculate(self, dst, long, signed, force=False): + with self.ebpf.get_free_register(dst) as dst: + with self.ebpf.save_registers([i for i in range(6) if i != dst]): + self.ebpf.call(FuncId.get_prandom_u32) + if dst != 0: + self.ebpf.r[dst] = self.ebpf.r0 + yield dst, True, False + + class RegisterDesc: def __init__(self, no, array): self.no = no diff --git a/ebpfcat/ebpfcat.py b/ebpfcat/ebpfcat.py index a8b1908c9392b485ab03f1489eeb7e3278feb037..df21ca4609c670cfc29007128d0d3ab2789da390 100644 --- a/ebpfcat/ebpfcat.py +++ b/ebpfcat/ebpfcat.py @@ -21,7 +21,7 @@ from struct import pack, unpack, calcsize, pack_into, unpack_from from time import time from .arraymap import ArrayMap, ArrayGlobalVarDesc from .ethercat import ECCmd, EtherCat, Packet, Terminal -from .ebpf import FuncId, MemoryDesc, SubProgram, ktime +from .ebpf import FuncId, MemoryDesc, SubProgram, prandom from .xdp import XDP, XDPExitCode from .bpf import ( ProgType, MapType, create_map, update_elem, prog_test_run, lookup_elem) @@ -277,11 +277,9 @@ class EtherXDP(XDP): rate = 0 def program(self): - with self.tmp: - self.ebpf.tmp = ktime(self.ebpf) - self.ebpf.tmp = self.ebpf.tmp * 0xcf019d85 + 1 - with self.ebpf.tmp & 0xffff < self.rate: - self.ebpf.exit(XDPExitCode.DROP) + with prandom(self.ebpf) & 0xffff < self.rate: + self.dropcounter += 1 + self.ebpf.exit(XDPExitCode.DROP) with self.packetSize > 24 as p, p.pH[12] == 0xA488, p.pB[16] == 0: self.r3 = p.pI[18] with self.counters.get_address(None, False, False) as (dst, _), \