From a4cfc2b7aafd0f4c77a4da400b2c448c1d16834a Mon Sep 17 00:00:00 2001 From: varun-r-mallya Date: Wed, 15 Oct 2025 17:49:20 +0530 Subject: [PATCH] add assignments table and offset handler --- pythonbpf/vmlinux_parser/class_handler.py | 10 ++-- pythonbpf/vmlinux_parser/dependency_node.py | 49 ++++++++++++++++--- pythonbpf/vmlinux_parser/import_detector.py | 8 ++- .../vmlinux_parser/ir_gen/ir_generation.py | 5 +- tests/failing_tests/xdp_pass.py | 3 +- 5 files changed, 62 insertions(+), 13 deletions(-) diff --git a/pythonbpf/vmlinux_parser/class_handler.py b/pythonbpf/vmlinux_parser/class_handler.py index ce08530..0702939 100644 --- a/pythonbpf/vmlinux_parser/class_handler.py +++ b/pythonbpf/vmlinux_parser/class_handler.py @@ -71,7 +71,9 @@ def process_vmlinux_post_ast( if len(field_elem) == 2: field_name, field_type = field_elem elif len(field_elem) == 3: - raise NotImplementedError("Bitfields are not supported in the current version") + raise NotImplementedError( + "Bitfields are not supported in the current version" + ) field_name, field_type, bitfield_size = field_elem field_table[field_name] = [field_type, bitfield_size] elif hasattr(class_obj, "__annotations__"): @@ -145,7 +147,8 @@ def process_vmlinux_post_ast( process_vmlinux_post_ast( containing_type, llvm_handler, handler, processing_stack ) - new_dep_node.set_field_ready(elem_name, True) + size_of_containing_type = (handler[containing_type.__name__]).__sizeof__() + new_dep_node.set_field_ready(elem_name, True, size_of_containing_type) elif containing_type.__module__ == ctypes.__name__: logger.debug(f"Processing ctype internal{containing_type}") new_dep_node.set_field_ready(elem_name, True) @@ -162,7 +165,8 @@ def process_vmlinux_post_ast( process_vmlinux_post_ast( elem_type, llvm_handler, handler, processing_stack ) - new_dep_node.set_field_ready(elem_name, True) + size_of_containing_type = (handler[elem_type.__name__]).__sizeof__() + new_dep_node.set_field_ready(elem_name, True, size_of_containing_type) else: raise ValueError( f"{elem_name} with type {elem_type} from module {module_name} not supported in recursive resolver" diff --git a/pythonbpf/vmlinux_parser/dependency_node.py b/pythonbpf/vmlinux_parser/dependency_node.py index 8a512cd..a6d4013 100644 --- a/pythonbpf/vmlinux_parser/dependency_node.py +++ b/pythonbpf/vmlinux_parser/dependency_node.py @@ -1,5 +1,6 @@ from dataclasses import dataclass, field from typing import Dict, Any, Optional +import ctypes # TODO: FIX THE FUCKING TYPE NAME CONVENTION. @@ -140,11 +141,14 @@ class DependencyNode: type_size=type_size, ctype_complex_type=ctype_complex_type, bitfield_size=bitfield_size, - offset=offset + offset=offset, ) # Invalidate readiness cache self._ready_cache = None + def __sizeof__(self): + return self.current_offset + def get_field(self, name: str) -> Field: """Get a field by name.""" return self.fields[name] @@ -211,20 +215,53 @@ class DependencyNode: # Invalidate readiness cache self._ready_cache = None - def set_field_ready(self, name: str, is_ready: bool = False) -> None: + def set_field_ready(self, name: str, is_ready: bool = False, size_of_containing_type: Optional[int] = None) -> None: """Mark a field as ready or not ready.""" if name not in self.fields: raise KeyError(f"Field '{name}' does not exist in node '{self.name}'") self.fields[name].set_ready(is_ready) self.fields[name].set_offset(self.current_offset) - self.current_offset += self._calculate_size(name) - + self.current_offset += self._calculate_size(name, size_of_containing_type) # Invalidate readiness cache self._ready_cache = None - def _calculate_size(self, name: str) -> int: - pass + def _calculate_size(self, name: str, size_of_containing_type: Optional[int] = None) -> int: + processing_field = self.fields[name] + # size_of_field will be in bytes + if processing_field.type.__module__ == ctypes.__name__: + size_of_field = ctypes.sizeof(processing_field.type) + return size_of_field + elif processing_field.type.__module__ == "vmlinux": + size_of_field: int = 0 + if processing_field.ctype_complex_type is not None: + if issubclass(processing_field.ctype_complex_type, ctypes.Array): + if processing_field.containing_type.__module__ == ctypes.__name__: + size_of_field = ( + ctypes.sizeof(processing_field.containing_type) + * processing_field.type_size + ) + return size_of_field + elif processing_field.containing_type.__module__ == "vmlinux": + size_of_field = ( + size_of_containing_type + * processing_field.type_size + ) + return size_of_field + elif issubclass(processing_field.ctype_complex_type, ctypes._Pointer): + return ctypes.sizeof(ctypes.pointer()) + else: + raise NotImplementedError( + "This subclass of ctype not supported yet" + ) + else: + # search up pre-created stuff and get size + return size_of_containing_type + + else: + raise ModuleNotFoundError("Module is not supported for the operation") + raise RuntimeError("control should not reach here") + @property def is_ready(self) -> bool: """Check if the node is ready (all fields are ready).""" diff --git a/pythonbpf/vmlinux_parser/import_detector.py b/pythonbpf/vmlinux_parser/import_detector.py index f5789ce..e314a35 100644 --- a/pythonbpf/vmlinux_parser/import_detector.py +++ b/pythonbpf/vmlinux_parser/import_detector.py @@ -129,7 +129,13 @@ def vmlinux_proc(tree: ast.AST, module): ) IRGenerator(module, handler) + return assignments def process_vmlinux_assign(node, module, assignments: Dict[str, type]): - raise NotImplementedError("Assignment handling has not been implemented yet") + # Check if this is a simple assignment with a constant value + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + target_name = node.targets[0].id + if isinstance(node.value, ast.Constant): + assignments[target_name] = node.value.value + logger.info(f"Added assignment: {target_name} = {node.value.value}") diff --git a/pythonbpf/vmlinux_parser/ir_gen/ir_generation.py b/pythonbpf/vmlinux_parser/ir_gen/ir_generation.py index 1a2be62..d500cf0 100644 --- a/pythonbpf/vmlinux_parser/ir_gen/ir_generation.py +++ b/pythonbpf/vmlinux_parser/ir_gen/ir_generation.py @@ -30,6 +30,7 @@ class IRGenerator: # this part cannot yet resolve circular dependencies. Gets stuck on an infinite loop during that. self.generated.append(struct.name) - - def struct_name_generator(self, ): + def struct_name_generator( + self, + ) -> None: pass diff --git a/tests/failing_tests/xdp_pass.py b/tests/failing_tests/xdp_pass.py index f44910d..da438c8 100644 --- a/tests/failing_tests/xdp_pass.py +++ b/tests/failing_tests/xdp_pass.py @@ -1,8 +1,9 @@ from pythonbpf import bpf, map, section, bpfglobal, compile_to_ir from pythonbpf.maps import HashMap from pythonbpf.helper import XDP_PASS +from vmlinux import TASK_COMM_LEN # noqa: F401 +from vmlinux import struct_trace_event_raw_sys_enter # noqa: F401 # from vmlinux import struct_request -from vmlinux import struct_trace_event_raw_sys_enter from vmlinux import struct_xdp_md # from vmlinux import struct_trace_event_raw_sys_enter # noqa: F401 # from vmlinux import struct_ring_buffer_per_cpu # noqa: F401