Skip to content

Commit caecb8c

Browse files
Merge pull request #59 from pythonbpf/vmlinux-handler
LGTM
2 parents 66caa3c + 1a0e21e commit caecb8c

File tree

10 files changed

+231
-13
lines changed

10 files changed

+231
-13
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 & 1 deletion
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,10 +58,13 @@ 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)
62-
print("DEBUG:", vmlinux_symtab)
6368
structs_sym_tab = structs_proc(tree, module, bpf_chunks)
6469
map_sym_tab = maps_proc(tree, module, bpf_chunks)
6570
func_proc(tree, module, bpf_chunks, map_sym_tab, structs_sym_tab)

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: 21 additions & 3 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

@@ -130,7 +142,12 @@ def get_operand_value(
130142
logger.info(f"var is {var}, base_type is {base_type}, depth is {depth}")
131143
val = deref_to_depth(func, builder, var, depth)
132144
return val
133-
raise ValueError(f"Undefined variable: {operand.id}")
145+
else:
146+
# Check if it's a vmlinux enum/constant
147+
vmlinux_result = VmlinuxHandlerRegistry.handle_name(operand.id)
148+
if vmlinux_result is not None:
149+
val, _ = vmlinux_result
150+
return val
134151
elif isinstance(operand, ast.Constant):
135152
if isinstance(operand.value, int):
136153
cst = ir.Constant(ir.IntType(64), int(operand.value))
@@ -332,6 +349,7 @@ def _handle_unary_op(
332349
neg_one = ir.Constant(ir.IntType(64), -1)
333350
result = builder.mul(operand, neg_one)
334351
return result, ir.IntType(64)
352+
return None
335353

336354

337355
# ============================================================================

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: 14 additions & 2 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
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
@@ -384,7 +390,13 @@ def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_t
384390
builder = ir.IRBuilder(block)
385391

386392
process_func_body(
387-
module, builder, func_node, func, ret_type, map_sym_tab, structs_sym_tab
393+
module,
394+
builder,
395+
func_node,
396+
func,
397+
ret_type,
398+
map_sym_tab,
399+
structs_sym_tab,
388400
)
389401
return func
390402

pythonbpf/helper/printk_formatter.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
from llvmlite import ir
55
from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
6+
from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
67

78
logger = logging.getLogger(__name__)
89

@@ -108,6 +109,16 @@ def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab):
108109
if local_sym_tab and name_node.id in local_sym_tab:
109110
_, var_type, tmp = local_sym_tab[name_node.id]
110111
_populate_fval(var_type, name_node, fmt_parts, exprs)
112+
else:
113+
# Try to resolve through vmlinux registry if not in local symbol table
114+
result = VmlinuxHandlerRegistry.handle_name(name_node.id)
115+
if result:
116+
val, var_type = result
117+
_populate_fval(var_type, name_node, fmt_parts, exprs)
118+
else:
119+
raise ValueError(
120+
f"Variable '{name_node.id}' not found in symbol table or vmlinux"
121+
)
111122

112123

113124
def _process_attr_in_fval(attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab):

pythonbpf/maps/maps_pass.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
from .maps_utils import MapProcessorRegistry
77
from .map_types import BPFMapType
88
from .map_debug_info import create_map_debug_info, create_ringbuf_debug_info
9+
from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
10+
911

1012
logger: Logger = logging.getLogger(__name__)
1113

@@ -51,7 +53,7 @@ def _parse_map_params(rval, expected_args=None):
5153
"""Parse map parameters from call arguments and keywords."""
5254

5355
params = {}
54-
56+
handler = VmlinuxHandlerRegistry.get_handler()
5557
# Parse positional arguments
5658
if expected_args:
5759
for i, arg_name in enumerate(expected_args):
@@ -65,7 +67,12 @@ def _parse_map_params(rval, expected_args=None):
6567
# Parse keyword arguments (override positional)
6668
for keyword in rval.keywords:
6769
if isinstance(keyword.value, ast.Name):
68-
params[keyword.arg] = keyword.value.id
70+
name = keyword.value.id
71+
if handler and handler.is_vmlinux_enum(name):
72+
result = handler.get_vmlinux_enum_value(name)
73+
params[keyword.arg] = result if result is not None else name
74+
else:
75+
params[keyword.arg] = name
6976
elif isinstance(keyword.value, ast.Constant):
7077
params[keyword.arg] = keyword.value.value
7178

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
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 get_vmlinux_enum_value(self, name):
58+
"""Handle vmlinux enum constants by returning LLVM IR constants"""
59+
if self.is_vmlinux_enum(name):
60+
value = self.vmlinux_symtab[name]["value"]
61+
logger.info(f"The value of vmlinux enum {name} = {value}")
62+
return value
63+
return None
64+
65+
def handle_vmlinux_struct(self, struct_name, module, builder):
66+
"""Handle vmlinux struct initializations"""
67+
if self.is_vmlinux_struct(struct_name):
68+
# TODO: Implement core-specific struct handling
69+
# This will be more complex and depends on the BTF information
70+
logger.info(f"Handling vmlinux struct {struct_name}")
71+
# Return struct type and allocated pointer
72+
# This is a stub, actual implementation will be more complex
73+
return None
74+
return None
75+
76+
def handle_vmlinux_struct_field(
77+
self, struct_var_name, field_name, module, builder, local_sym_tab
78+
):
79+
"""Handle access to vmlinux struct fields"""
80+
# Check if it's a variable of vmlinux struct type
81+
if struct_var_name in local_sym_tab:
82+
var_info = local_sym_tab[struct_var_name] # noqa: F841
83+
# Need to check if this variable is a vmlinux struct
84+
# This will depend on how you track vmlinux struct types in your symbol table
85+
logger.info(
86+
f"Attempting to access field {field_name} of possible vmlinux struct {struct_var_name}"
87+
)
88+
# Return pointer to field and field type
89+
return None
90+
return None
Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,26 @@
1-
from pythonbpf import bpf, section, bpfglobal, compile_to_ir
1+
import logging
2+
3+
from pythonbpf import bpf, section, bpfglobal, compile_to_ir, map
4+
from pythonbpf import compile # noqa: F401
25
from vmlinux import TASK_COMM_LEN # noqa: F401
36
from vmlinux import struct_trace_event_raw_sys_enter # noqa: F401
7+
from ctypes import c_uint64, c_int32, c_int64
8+
from pythonbpf.maps import HashMap
49

510
# from vmlinux import struct_uinput_device
611
# from vmlinux import struct_blk_integrity_iter
7-
from ctypes import c_int64
12+
13+
14+
@bpf
15+
@map
16+
def mymap() -> HashMap:
17+
return HashMap(key=c_int32, value=c_uint64, max_entries=TASK_COMM_LEN)
18+
19+
20+
@bpf
21+
@map
22+
def mymap2() -> HashMap:
23+
return HashMap(key=c_int32, value=c_uint64, max_entries=18)
824

925

1026
# Instructions to how to run this program
@@ -16,8 +32,9 @@
1632
@bpf
1733
@section("tracepoint/syscalls/sys_enter_execve")
1834
def hello_world(ctx: struct_trace_event_raw_sys_enter) -> c_int64:
19-
print("Hello, World!")
20-
return c_int64(0)
35+
a = 2 + TASK_COMM_LEN + TASK_COMM_LEN
36+
print(f"Hello, World{TASK_COMM_LEN} and {a}")
37+
return c_int64(TASK_COMM_LEN + 2)
2138

2239

2340
@bpf
@@ -26,4 +43,5 @@ def LICENSE() -> str:
2643
return "GPL"
2744

2845

29-
compile_to_ir("simple_struct_test.py", "simple_struct_test.ll")
46+
compile_to_ir("simple_struct_test.py", "simple_struct_test.ll", loglevel=logging.DEBUG)
47+
# compile()

0 commit comments

Comments
 (0)