Skip to content

Commit 6ccbab4

Browse files
r41k0ur41k0u
committed
Complete printk refactor
1 parent 7b01f1d commit 6ccbab4

2 files changed

Lines changed: 35 additions & 158 deletions

File tree

‎pythonbpf/helper/bpf_helper_handler.py‎

Lines changed: 23 additions & 152 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from llvmlite import ir
33
from pythonbpf.expr_pass import eval_expr
44
from enum import Enum
5-
from .helper_utils import HelperHandlerRegistry, get_or_create_ptr_from_arg, get_flags_val
5+
from .helper_utils import HelperHandlerRegistry, get_or_create_ptr_from_arg, get_flags_val, _handle_fstring_print, _simple_string_print
66

77

88
class BPFHelperID(Enum):
@@ -64,163 +64,34 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func,
6464
def bpf_printk_emitter(call, map_ptr, module, builder, func,
6565
local_sym_tab=None, struct_sym_tab=None,
6666
local_var_metadata=None):
67+
"""Emit LLVM IR for bpf_printk helper function call."""
6768
if not hasattr(func, "_fmt_counter"):
6869
func._fmt_counter = 0
6970

7071
if not call.args:
71-
raise ValueError("print expects at least one argument")
72+
raise ValueError(
73+
"bpf_printk expects at least one argument (format string)")
7274

75+
args = []
7376
if isinstance(call.args[0], ast.JoinedStr):
74-
fmt_parts = []
75-
exprs = []
76-
77-
for value in call.args[0].values:
78-
print("Value in f-string:", ast.dump(value))
79-
if isinstance(value, ast.Constant):
80-
if isinstance(value.value, str):
81-
fmt_parts.append(value.value)
82-
elif isinstance(value.value, int):
83-
fmt_parts.append("%lld")
84-
exprs.append(ir.Constant(ir.IntType(64), value.value))
85-
else:
86-
raise NotImplementedError(
87-
"Only string and integer constants are supported in f-string.")
88-
elif isinstance(value, ast.FormattedValue):
89-
print("Formatted value:", ast.dump(value))
90-
# TODO: Dirty handling here, only checks for int or str
91-
if isinstance(value.value, ast.Name):
92-
if local_sym_tab and value.value.id in local_sym_tab:
93-
var_ptr, var_type = local_sym_tab[value.value.id]
94-
if isinstance(var_type, ir.IntType):
95-
fmt_parts.append("%lld")
96-
exprs.append(value.value)
97-
elif var_type == ir.PointerType(ir.IntType(8)):
98-
# Case with string
99-
fmt_parts.append("%s")
100-
exprs.append(value.value)
101-
else:
102-
raise NotImplementedError(
103-
"Only integer and pointer types are supported in formatted values.")
104-
else:
105-
raise ValueError(
106-
f"Variable {value.value.id} not found in local symbol table.")
107-
elif isinstance(value.value, ast.Attribute):
108-
# object field access from struct
109-
if (isinstance(value.value.value, ast.Name) and
110-
local_sym_tab and
111-
value.value.value.id in local_sym_tab):
112-
var_name = value.value.value.id
113-
field_name = value.value.attr
114-
if local_var_metadata and var_name in local_var_metadata:
115-
var_type = local_var_metadata[var_name]
116-
if var_type in struct_sym_tab:
117-
struct_info = struct_sym_tab[var_type]
118-
if field_name in struct_info.fields:
119-
field_type = struct_info.field_type(
120-
field_name)
121-
if isinstance(field_type, ir.IntType):
122-
fmt_parts.append("%lld")
123-
exprs.append(value.value)
124-
elif field_type == ir.PointerType(ir.IntType(8)):
125-
fmt_parts.append("%s")
126-
exprs.append(value.value)
127-
else:
128-
raise NotImplementedError(
129-
"Only integer and pointer types are supported in formatted values.")
130-
else:
131-
raise ValueError(
132-
f"Field {field_name} not found in struct {var_type}.")
133-
else:
134-
raise ValueError(
135-
f"Struct type {var_type} for variable {var_name} not found in struct symbol table.")
136-
else:
137-
raise ValueError(
138-
f"Metadata for variable {var_name} not found in local variable metadata.")
139-
else:
140-
raise ValueError(
141-
f"Variable {value.value.value.id} not found in local symbol table.")
142-
else:
143-
raise NotImplementedError(
144-
"Only simple variable names are supported in formatted values.")
145-
else:
146-
raise NotImplementedError(
147-
"Unsupported value type in f-string.")
148-
149-
fmt_str = "".join(fmt_parts) + "\n" + "\0"
150-
fmt_name = f"{func.name}____fmt{func._fmt_counter}"
151-
func._fmt_counter += 1
152-
153-
fmt_gvar = ir.GlobalVariable(
154-
module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name)
155-
fmt_gvar.global_constant = True
156-
fmt_gvar.initializer = ir.Constant( # type: ignore
157-
ir.ArrayType(ir.IntType(8), len(fmt_str)),
158-
bytearray(fmt_str.encode("utf8"))
159-
)
160-
fmt_gvar.linkage = "internal"
161-
fmt_gvar.align = 1 # type: ignore
162-
163-
fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType())
164-
165-
args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]
166-
167-
# Only 3 args supported in bpf_printk
168-
if len(exprs) > 3:
169-
print(
170-
"Warning: bpf_printk supports up to 3 arguments, extra arguments will be ignored.")
171-
172-
for expr in exprs[:3]:
173-
print(f"{ast.dump(expr)}")
174-
val, _ = eval_expr(func, module, builder,
175-
expr, local_sym_tab, None, struct_sym_tab, local_var_metadata)
176-
if val:
177-
if isinstance(val.type, ir.PointerType):
178-
val = builder.ptrtoint(val, ir.IntType(64))
179-
elif isinstance(val.type, ir.IntType):
180-
if val.type.width < 64:
181-
val = builder.sext(val, ir.IntType(64))
182-
else:
183-
print(
184-
"Warning: Only integer and pointer types are supported in bpf_printk arguments. Others will be converted to 0.")
185-
val = ir.Constant(ir.IntType(64), 0)
186-
args.append(val)
187-
else:
188-
print(
189-
"Warning: Failed to evaluate expression for bpf_printk argument. It will be converted to 0.")
190-
args.append(ir.Constant(ir.IntType(64), 0))
191-
fn_type = ir.FunctionType(ir.IntType(
192-
64), [ir.PointerType(), ir.IntType(32)], var_arg=True)
193-
fn_ptr_type = ir.PointerType(fn_type)
194-
fn_addr = ir.Constant(ir.IntType(64), 6)
195-
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
196-
return builder.call(fn_ptr, args, tail=True)
197-
198-
for arg in call.args:
199-
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
200-
fmt_str = arg.value + "\n" + "\0"
201-
fmt_name = f"{func.name}____fmt{func._fmt_counter}"
202-
func._fmt_counter += 1
203-
204-
fmt_gvar = ir.GlobalVariable(
205-
module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name)
206-
fmt_gvar.global_constant = True
207-
fmt_gvar.initializer = ir.Constant( # type: ignore
208-
ir.ArrayType(ir.IntType(8), len(fmt_str)),
209-
bytearray(fmt_str.encode("utf8"))
210-
)
211-
fmt_gvar.linkage = "internal"
212-
fmt_gvar.align = 1 # type: ignore
213-
214-
fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType())
215-
216-
fn_type = ir.FunctionType(ir.IntType(
217-
64), [ir.PointerType(), ir.IntType(32)], var_arg=True)
218-
fn_ptr_type = ir.PointerType(fn_type)
219-
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PRINTK.value)
220-
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
221-
222-
builder.call(fn_ptr, [fmt_ptr, ir.Constant(
223-
ir.IntType(32), len(fmt_str))], tail=True)
77+
args = _handle_fstring_print(call.args[0], module, builder, func,
78+
local_sym_tab, struct_sym_tab,
79+
local_var_metadata)
80+
elif isinstance(call.args[0], ast.Constant) and isinstance(call.args[0].value, str):
81+
# TODO: We are onbly supporting single arguments for now.
82+
# In case of multiple args, the first one will be taken.
83+
args = _simple_string_print(call.args[0], module, builder, func)
84+
else:
85+
raise NotImplementedError(
86+
"Only simple string literals or f-strings are supported in bpf_printk.")
87+
88+
fn_type = ir.FunctionType(
89+
ir.IntType(64), [ir.PointerType(), ir.IntType(32)], var_arg=True)
90+
fn_ptr_type = ir.PointerType(fn_type)
91+
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PRINTK.value)
92+
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
93+
94+
builder.call(fn_ptr, args, tail=True)
22495
return None
22596

22697

‎pythonbpf/helper/helper_utils.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,15 @@ def get_flags_val(arg, builder, local_sym_tab):
7373
"Only simple variable names or integer constants are supported as flags in map helpers.")
7474

7575

76+
def _simple_string_print(string_value, module, builder, func):
77+
"""Emit code for a simple string print statement."""
78+
fmt_str = string_value + "\n\0"
79+
fmt_ptr = _create_format_string_global(fmt_str, func, module, builder)
80+
81+
args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]
82+
return args
83+
84+
7685
def _handle_fstring_print(joined_str, module, builder, func,
7786
local_sym_tab=None, struct_sym_tab=None,
7887
local_var_metadata=None):
@@ -93,10 +102,8 @@ def _handle_fstring_print(joined_str, module, builder, func,
93102
raise NotImplementedError(
94103
f"Unsupported f-string value type: {type(value)}")
95104

96-
fmt_str = "".join(fmt_parts) + "\n\0"
97-
fmt_ptr = _create_format_string_global(fmt_str, func, module, builder)
98-
99-
args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]
105+
fmt_str = "".join(fmt_parts)
106+
args = _simple_string_print(fmt_str, module, builder, func)
100107

101108
# NOTE: Process expressions (limited to 3 due to BPF constraints)
102109
if len(exprs) > 3:
@@ -109,8 +116,7 @@ def _handle_fstring_print(joined_str, module, builder, func,
109116
local_var_metadata)
110117
args.append(arg_value)
111118

112-
# Call the BPF_PRINTK helper
113-
return _call_bpf_printk_helper(args, builder)
119+
return args
114120

115121

116122
def _process_constant_in_fstring(cst, fmt_parts, exprs):

0 commit comments

Comments
 (0)