improving arraymap code

ebpfcat/arraymap.py

-from struct import pack, unpack, calcsize
+from struct import pack_into, unpack_from, calcsize
 
-from .ebpf import FuncId, Map, Memory, Opcode, SubProgram
+from .ebpf import FuncId, Map, Memory, MemoryDesc, Opcode
 from .bpf import create_map, lookup_elem, MapType, update_elem
 
 
-class ArrayGlobalVarDesc:
-    def __init__(self, map, size):
+class ArrayGlobalVarDesc(MemoryDesc):
+    base_register = 0
+
+    def __init__(self, map, fmt):
         self.map = map
-        self.fmt = size
+        self.fmt = fmt
 
-    def __get__(self, ebpf, owner):
-        if ebpf is None:
-            return self
-        position = ebpf.__dict__[self.name]
-        if isinstance(ebpf, SubProgram):
-            ebpf = ebpf.ebpf
-        if ebpf.loaded:
-            data = ebpf.__dict__[self.map.name].data[
-                    position : position+calcsize(self.fmt)]
-            return unpack(self.fmt, data)[0]
-        return Memory(ebpf, Memory.fmt_to_opcode[self.fmt],
-                      ebpf.r0 + position, self.fmt.islower())
+    def fmt_addr(self, ebpf):
+        return self.fmt, ebpf.__dict__[self.name]
 
     def __set_name__(self, owner, name):
         self.name = name
 
-    def __set__(self, ebpf, value):
-        position = ebpf.__dict__[self.name]
-        if isinstance(ebpf, SubProgram):
-            ebpf = ebpf.ebpf
-        if ebpf.loaded:
-            ebpf.__dict__[self.map.name].data[
-               position : position+calcsize(self.fmt)] = pack(self.fmt, value)
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+        fmt, addr = self.fmt_addr(instance)
+        if instance.ebpf.loaded:
+            data = instance.ebpf.__dict__[self.map.name].data
+            return unpack_from(fmt, data, addr)[0]
+        else:
+            return super().__get__(instance, owner)
+
+    def __set__(self, instance, value):
+        fmt, addr = self.fmt_addr(instance)
+        if instance.ebpf.loaded:
+            pack_into(fmt, instance.ebpf.__dict__[self.map.name].data,
+                      addr, value)
         else:
-            getattr(ebpf, f"m{self.fmt}")[ebpf.r0 + position] = value
+            super().__set__(instance, value)
 
 
 class ArrayMapAccess:
@@ -49,8 +49,8 @@ class ArrayMapAccess:
 
 
 class ArrayMap(Map):
-    def globalVar(self, signed=False, size=4):
-        return ArrayGlobalVarDesc(self, size, signed)
+    def globalVar(self, fmt="I"):
+        return ArrayGlobalVarDesc(self, fmt)
 
     def add_program(self, owner, prog):
         position = getattr(owner, self.name)

ebpfcat/ebpf.py

@@ -664,24 +664,18 @@ class MemoryDesc:
     def __get__(self, instance, owner):
         if instance is None:
             return self
-        elif isinstance(instance, SubProgram):
-            ebpf = instance.ebpf
-        else:
-            ebpf = instance
         fmt, addr = self.fmt_addr(instance)
-        return Memory(ebpf, Memory.fmt_to_opcode[fmt],
-                      ebpf.r[self.base_register] + addr, fmt.islower())
+        return Memory(instance.ebpf, Memory.fmt_to_opcode[fmt],
+                      instance.ebpf.r[self.base_register] + addr,
+                      fmt.islower())
 
     def __set__(self, instance, value):
-        if isinstance(instance, SubProgram):
-            ebpf = instance.ebpf
-        else:
-            ebpf = instance
+        ebpf = instance.ebpf
         fmt, addr = self.fmt_addr(instance)
         bits = Memory.fmt_to_opcode[fmt]
         if isinstance(value, int):
             ebpf.append(Opcode.ST + bits, self.base_register, 0,
-                        self.addr(instance), value)
+                        addr, value)
             return
         elif isinstance(value, IAdd):
             value = value.value
@@ -987,6 +981,10 @@ class EBPF:
         yield self.stack
         self.stack = oldstack
 
+    @property
+    def ebpf(self):
+        return self
+
     tmp = TemporaryDesc(None, "r")
     stmp = TemporaryDesc(None, "sr")
     wtmp = TemporaryDesc(None, "w")