Skip to content

Commit ebb872f

Browse files
r41k0ur41k0u
committed
Alocate space at the entry bb of each bpf chunk
1 parent aeb9a45 commit ebb872f

2 files changed

Lines changed: 75 additions & 15 deletions

File tree

‎examples/execve3.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,7 @@ def hello_again(ctx: c_void_p) -> c_int64:
3232
# last().delete(key)
3333
if True:
3434
print("we prevailed")
35-
# ts = ktime()
36-
ktime()
35+
ts = ktime()
3736
# last().update(key, ts)
3837
return c_int64(0)
3938

‎pythonbpf/functions_pass.py‎

Lines changed: 74 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -39,30 +39,32 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab):
3939
if isinstance(rval, ast.Constant):
4040
if isinstance(rval.value, int):
4141
# Assume c_int64 for now
42-
var = builder.alloca(ir.IntType(64), name=var_name)
43-
var.align = 8
44-
builder.store(ir.Constant(ir.IntType(64), rval.value), var)
45-
local_sym_tab[var_name] = var
42+
# var = builder.alloca(ir.IntType(64), name=var_name)
43+
# var.align = 8
44+
builder.store(ir.Constant(ir.IntType(64), rval.value),
45+
local_sym_tab[var_name])
46+
# local_sym_tab[var_name] = var
4647
print(f"Assigned constant {rval.value} to {var_name}")
4748
elif isinstance(rval, ast.Call):
4849
if isinstance(rval.func, ast.Name):
4950
call_type = rval.func.id
5051
print(f"Assignment call type: {call_type}")
5152
if call_type in num_types and len(rval.args) == 1 and isinstance(rval.args[0], ast.Constant) and isinstance(rval.args[0].value, int):
5253
ir_type = ctypes_to_ir(call_type)
53-
var = builder.alloca(ir_type, name=var_name)
54-
var.align = ir_type.width // 8
55-
builder.store(ir.Constant(ir_type, rval.args[0].value), var)
54+
# var = builder.alloca(ir_type, name=var_name)
55+
# var.align = ir_type.width // 8
56+
builder.store(ir.Constant(
57+
ir_type, rval.args[0].value), local_sym_tab[var_name])
5658
print(f"Assigned {call_type} constant "
5759
f"{rval.args[0].value} to {var_name}")
58-
local_sym_tab[var_name] = var
60+
# local_sym_tab[var_name] = var
5961
elif call_type in helper_func_list:
60-
var = builder.alloca(ir.IntType(64), name=var_name)
61-
var.align = 8
62+
# var = builder.alloca(ir.IntType(64), name=var_name)
63+
# var.align = 8
6264
val = handle_helper_call(
6365
rval, module, builder, None, local_sym_tab, map_sym_tab)
64-
builder.store(val, var)
65-
local_sym_tab[var_name] = var
66+
builder.store(val, local_sym_tab[var_name])
67+
# local_sym_tab[var_name] = var
6668
print(f"Assigned constant {rval.func.id} to {var_name}")
6769
else:
6870
print(f"Unsupported assignment call type: {call_type}")
@@ -73,8 +75,12 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab):
7375
if map_name in map_sym_tab:
7476
map_ptr = map_sym_tab[map_name]
7577
if method_name in helper_func_list:
76-
handle_helper_call(
78+
val = handle_helper_call(
7779
rval, module, builder, func, local_sym_tab, map_sym_tab)
80+
# var = builder.alloca(ir.IntType(64), name=var_name)
81+
# var.align = 8
82+
builder.store(val, local_sym_tab[var_name])
83+
# local_sym_tab[var_name] = var
7884
else:
7985
print("Unsupported assignment call structure")
8086
else:
@@ -188,6 +194,61 @@ def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab):
188194

189195
local_sym_tab = {}
190196

197+
# pre-allocate dynamic variables
198+
for stmt in func_node.body:
199+
if isinstance(stmt, ast.Assign):
200+
if len(stmt.targets) != 1:
201+
print("Unsupported multiassignment")
202+
continue
203+
target = stmt.targets[0]
204+
if not isinstance(target, ast.Name):
205+
print("Unsupported assignment target")
206+
continue
207+
var_name = target.id
208+
rval = stmt.value
209+
if isinstance(rval, ast.Call):
210+
if isinstance(rval.func, ast.Name):
211+
call_type = rval.func.id
212+
if call_type in ("c_int32", "c_int64", "c_uint32", "c_uint64"):
213+
ir_type = ctypes_to_ir(call_type)
214+
var = builder.alloca(ir_type, name=var_name)
215+
var.align = ir_type.width // 8
216+
print(
217+
f"Pre-allocated variable {var_name} of type {call_type}")
218+
elif call_type in helper_func_list:
219+
# Assume return type is int64 for now
220+
ir_type = ir.IntType(64)
221+
var = builder.alloca(ir_type, name=var_name)
222+
var.align = ir_type.width // 8
223+
print(
224+
f"Pre-allocated variable {var_name} for helper")
225+
elif isinstance(rval.func, ast.Attribute):
226+
ir_type = ir.PointerType(ir.IntType(64))
227+
var = builder.alloca(ir_type, name=var_name)
228+
# var.align = ir_type.width // 8
229+
print(
230+
f"Pre-allocated variable {var_name} for map")
231+
else:
232+
print("Unsupported assignment call function type")
233+
continue
234+
elif isinstance(rval, ast.Constant):
235+
if isinstance(rval.value, int):
236+
# Assume c_int64 for now
237+
ir_type = ir.IntType(64)
238+
var = builder.alloca(ir_type, name=var_name)
239+
var.align = ir_type.width // 8
240+
print(
241+
f"Pre-allocated variable {var_name} of type c_int64")
242+
else:
243+
print("Unsupported constant type")
244+
continue
245+
else:
246+
print("Unsupported assignment value type")
247+
continue
248+
local_sym_tab[var_name] = var
249+
250+
print(f"Local symbol table: {local_sym_tab.keys()}")
251+
191252
for stmt in func_node.body:
192253
did_return = process_stmt(func, module, builder, stmt, local_sym_tab,
193254
map_sym_tab, did_return, ret_type)

0 commit comments

Comments
 (0)