Skip to content

Commit 5955db8

Browse files
add vmlinux expressions to eval expr
1 parent e499c29 commit 5955db8

File tree

8 files changed

+180
-10
lines changed

8 files changed

+180
-10
lines changed

pythonbpf/allocation_pass.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from dataclasses import dataclass
66
from typing import Any
77
from pythonbpf.helper import HelperHandlerRegistry
8+
from .expr import VmlinuxHandlerRegistry
89
from pythonbpf.type_deducer import ctypes_to_ir
910

1011
logger = logging.getLogger(__name__)
@@ -49,6 +50,15 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
4950
logger.debug(f"Variable {var_name} already allocated, skipping")
5051
return
5152

53+
# When allocating a variable, check if it's a vmlinux struct type
54+
if isinstance(stmt.value, ast.Name) and VmlinuxHandlerRegistry.is_vmlinux_struct(
55+
stmt.value.id
56+
):
57+
# Handle vmlinux struct allocation
58+
# This requires more implementation
59+
print(stmt.value)
60+
pass
61+
5262
# Determine type and allocate based on rval
5363
if isinstance(rval, ast.Call):
5464
_allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab)

pythonbpf/codegen.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
from .maps import maps_proc
66
from .structs import structs_proc
77
from .vmlinux_parser import vmlinux_proc
8+
from pythonbpf.vmlinux_parser.vmlinux_exports_handler import VmlinuxHandler
9+
from .expr import VmlinuxHandlerRegistry
810
from .globals_pass import (
911
globals_list_creation,
1012
globals_processing,
@@ -56,6 +58,10 @@ def processor(source_code, filename, module):
5658
logger.info(f"Found BPF function/struct: {func_node.name}")
5759

5860
vmlinux_symtab = vmlinux_proc(tree, module)
61+
if vmlinux_symtab:
62+
handler = VmlinuxHandler.initialize(vmlinux_symtab)
63+
VmlinuxHandlerRegistry.set_handler(handler)
64+
5965
populate_global_symbol_table(tree, module)
6066
license_processing(tree, module)
6167
globals_processing(tree, module)

pythonbpf/expr/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from .type_normalization import convert_to_bool, get_base_type_and_depth
33
from .ir_ops import deref_to_depth
44
from .call_registry import CallHandlerRegistry
5+
from .vmlinux_registry import VmlinuxHandlerRegistry
56

67
__all__ = [
78
"eval_expr",
@@ -11,4 +12,5 @@
1112
"deref_to_depth",
1213
"get_operand_value",
1314
"CallHandlerRegistry",
15+
"VmlinuxHandlerRegistry",
1416
]

pythonbpf/expr/expr_pass.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
get_base_type_and_depth,
1313
deref_to_depth,
1414
)
15+
from .vmlinux_registry import VmlinuxHandlerRegistry
1516

1617
logger: Logger = logging.getLogger(__name__)
1718

@@ -27,8 +28,12 @@ def _handle_name_expr(expr: ast.Name, local_sym_tab: Dict, builder: ir.IRBuilder
2728
val = builder.load(var)
2829
return val, local_sym_tab[expr.id].ir_type
2930
else:
30-
logger.info(f"Undefined variable {expr.id}")
31-
return None
31+
# Check if it's a vmlinux enum/constant
32+
vmlinux_result = VmlinuxHandlerRegistry.handle_name(expr.id)
33+
if vmlinux_result is not None:
34+
return vmlinux_result
35+
36+
raise SyntaxError(f"Undefined variable {expr.id}")
3237

3338

3439
def _handle_constant_expr(module, builder, expr: ast.Constant):
@@ -74,6 +79,13 @@ def _handle_attribute_expr(
7479
val = builder.load(gep)
7580
field_type = metadata.field_type(attr_name)
7681
return val, field_type
82+
83+
# Try vmlinux handler as fallback
84+
vmlinux_result = VmlinuxHandlerRegistry.handle_attribute(
85+
expr, local_sym_tab, None, builder
86+
)
87+
if vmlinux_result is not None:
88+
return vmlinux_result
7789
return None
7890

7991

pythonbpf/expr/vmlinux_registry.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
import ast
2+
3+
4+
class VmlinuxHandlerRegistry:
5+
"""Registry for vmlinux handler operations"""
6+
7+
_handler = None
8+
9+
@classmethod
10+
def set_handler(cls, handler):
11+
"""Set the vmlinux handler"""
12+
cls._handler = handler
13+
14+
@classmethod
15+
def get_handler(cls):
16+
"""Get the vmlinux handler"""
17+
return cls._handler
18+
19+
@classmethod
20+
def handle_name(cls, name):
21+
"""Try to handle a name as vmlinux enum/constant"""
22+
if cls._handler is None:
23+
return None
24+
return cls._handler.handle_vmlinux_enum(name)
25+
26+
@classmethod
27+
def handle_attribute(cls, expr, local_sym_tab, module, builder):
28+
"""Try to handle an attribute access as vmlinux struct field"""
29+
if cls._handler is None:
30+
return None
31+
32+
if isinstance(expr.value, ast.Name):
33+
var_name = expr.value.id
34+
field_name = expr.attr
35+
return cls._handler.handle_vmlinux_struct_field(
36+
var_name, field_name, module, builder, local_sym_tab
37+
)
38+
return None
39+
40+
@classmethod
41+
def is_vmlinux_struct(cls, name):
42+
"""Check if a name refers to a vmlinux struct"""
43+
if cls._handler is None:
44+
return False
45+
return cls._handler.is_vmlinux_struct(name)

pythonbpf/functions/functions_pass.py

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,13 @@ def process_stmt(
311311

312312

313313
def process_func_body(
314-
module, builder, func_node, func, ret_type, map_sym_tab, structs_sym_tab, vmlinux_symtab
314+
module,
315+
builder,
316+
func_node,
317+
func,
318+
ret_type,
319+
map_sym_tab,
320+
structs_sym_tab,
315321
):
316322
"""Process the body of a bpf function"""
317323
# TODO: A lot. We just have print -> bpf_trace_printk for now
@@ -350,7 +356,9 @@ def process_func_body(
350356
builder.ret(ir.Constant(ir.IntType(64), 0))
351357

352358

353-
def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_tab, vmlinux_symtab):
359+
def process_bpf_chunk(
360+
func_node, module, return_type, map_sym_tab, structs_sym_tab
361+
):
354362
"""Process a single BPF chunk (function) and emit corresponding LLVM IR."""
355363

356364
func_name = func_node.name
@@ -384,7 +392,13 @@ def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_t
384392
builder = ir.IRBuilder(block)
385393

386394
process_func_body(
387-
module, builder, func_node, func, ret_type, map_sym_tab, structs_sym_tab, vmlinux_symtab
395+
module,
396+
builder,
397+
func_node,
398+
func,
399+
ret_type,
400+
map_sym_tab,
401+
structs_sym_tab,
388402
)
389403
return func
390404

@@ -394,7 +408,7 @@ def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_t
394408
# ============================================================================
395409

396410

397-
def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab, vmlinux_symtab):
411+
def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab):
398412
for func_node in chunks:
399413
if is_global_function(func_node):
400414
continue
@@ -406,8 +420,7 @@ def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab, vmlinux_symtab
406420
module,
407421
ctypes_to_ir(infer_return_type(func_node)),
408422
map_sym_tab,
409-
structs_sym_tab,
410-
vmlinux_symtab
423+
structs_sym_tab
411424
)
412425

413426

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
import logging
2+
from llvmlite import ir
3+
4+
from pythonbpf.vmlinux_parser.assignment_info import AssignmentType
5+
6+
logger = logging.getLogger(__name__)
7+
8+
9+
class VmlinuxHandler:
10+
"""Handler for vmlinux-related operations"""
11+
12+
_instance = None
13+
14+
@classmethod
15+
def get_instance(cls):
16+
"""Get the singleton instance"""
17+
if cls._instance is None:
18+
logger.warning("VmlinuxHandler used before initialization")
19+
return None
20+
return cls._instance
21+
22+
@classmethod
23+
def initialize(cls, vmlinux_symtab):
24+
"""Initialize the handler with vmlinux symbol table"""
25+
cls._instance = cls(vmlinux_symtab)
26+
return cls._instance
27+
28+
def __init__(self, vmlinux_symtab):
29+
"""Initialize with vmlinux symbol table"""
30+
self.vmlinux_symtab = vmlinux_symtab
31+
logger.info(
32+
f"VmlinuxHandler initialized with {len(vmlinux_symtab) if vmlinux_symtab else 0} symbols"
33+
)
34+
35+
def is_vmlinux_enum(self, name):
36+
"""Check if name is a vmlinux enum constant"""
37+
return (
38+
name in self.vmlinux_symtab
39+
and self.vmlinux_symtab[name]["value_type"] == AssignmentType.CONSTANT
40+
)
41+
42+
def is_vmlinux_struct(self, name):
43+
"""Check if name is a vmlinux struct"""
44+
return (
45+
name in self.vmlinux_symtab
46+
and self.vmlinux_symtab[name]["value_type"] == AssignmentType.STRUCT
47+
)
48+
49+
def handle_vmlinux_enum(self, name):
50+
"""Handle vmlinux enum constants by returning LLVM IR constants"""
51+
if self.is_vmlinux_enum(name):
52+
value = self.vmlinux_symtab[name]["value"]
53+
logger.info(f"Resolving vmlinux enum {name} = {value}")
54+
return ir.Constant(ir.IntType(64), value), ir.IntType(64)
55+
return None
56+
57+
def handle_vmlinux_struct(self, struct_name, module, builder):
58+
"""Handle vmlinux struct initializations"""
59+
if self.is_vmlinux_struct(struct_name):
60+
# TODO: Implement core-specific struct handling
61+
# This will be more complex and depends on the BTF information
62+
logger.info(f"Handling vmlinux struct {struct_name}")
63+
# Return struct type and allocated pointer
64+
# This is a stub, actual implementation will be more complex
65+
return None
66+
return None
67+
68+
def handle_vmlinux_struct_field(
69+
self, struct_var_name, field_name, module, builder, local_sym_tab
70+
):
71+
"""Handle access to vmlinux struct fields"""
72+
# Check if it's a variable of vmlinux struct type
73+
if struct_var_name in local_sym_tab:
74+
var_info = local_sym_tab[struct_var_name]
75+
# Need to check if this variable is a vmlinux struct
76+
# This will depend on how you track vmlinux struct types in your symbol table
77+
logger.info(
78+
f"Attempting to access field {field_name} of possible vmlinux struct {struct_var_name}"
79+
)
80+
# Return pointer to field and field type
81+
return None
82+
return None

tests/passing_tests/vmlinux/simple_struct_test.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616
@bpf
1717
@section("tracepoint/syscalls/sys_enter_execve")
1818
def hello_world(ctx: struct_trace_event_raw_sys_enter) -> c_int64:
19-
print("Hello, World!")
20-
return c_int64(0)
19+
print("Hello, World")
20+
return c_int64(TASK_COMM_LEN)
2121

2222

2323
@bpf

0 commit comments

Comments
 (0)