varun-r-mallya/python-bpf
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/Python-BPF.git synced 2025-12-31 21:06:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix ruff errors

Browse Source
This commit is contained in:
varun-r-mallya
2025-10-02 01:23:55 +05:30
parent 6831f11179
commit 59b3d6514b
5 changed files with 268 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

230
pythonbpf/helper/bpf_helper_handler.py
View File

@ -1,10 +1,14 @@
import ast
from llvmlite import ir
from enum import Enum
from .helper_utils import (HelperHandlerRegistry,
get_or_create_ptr_from_arg, get_flags_val,
handle_fstring_print, simple_string_print,
get_data_ptr_and_size)
from .helper_utils import (
HelperHandlerRegistry,
get_or_create_ptr_from_arg,
get_flags_val,
handle_fstring_print,
simple_string_print,
get_data_ptr_and_size,
)
class BPFHelperID(Enum):
@ -18,9 +22,16 @@ class BPFHelperID(Enum):
@HelperHandlerRegistry.register("ktime")
def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func,
local_sym_tab=None, struct_sym_tab=None,
local_var_metadata=None):
def bpf_ktime_get_ns_emitter(
call,
map_ptr,
module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_ktime_get_ns helper function call.
"""
@ -34,27 +45,34 @@ def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("lookup")
def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func,
local_sym_tab=None, struct_sym_tab=None,
local_var_metadata=None):
def bpf_map_lookup_elem_emitter(
call,
map_ptr,
module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_map_lookup_elem helper function call.
"""
if not call.args or len(call.args) != 1:
raise ValueError("Map lookup expects exactly one argument (key), got "
f"{len(call.args)}")
raise ValueError(
"Map lookup expects exactly one argument (key), got " f"{len(call.args)}"
)
key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
fn_type = ir.FunctionType(
ir.PointerType(), # Return type: void*
[ir.PointerType(), ir.PointerType()], # Args: (void*, void*)
var_arg=False
var_arg=False,
)
fn_ptr_type = ir.PointerType(fn_type)
fn_addr = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
@ -63,33 +81,46 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("print")
def bpf_printk_emitter(call, map_ptr, module, builder, func,
local_sym_tab=None, struct_sym_tab=None,
local_var_metadata=None):
def bpf_printk_emitter(
call,
map_ptr,
module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""Emit LLVM IR for bpf_printk helper function call."""
if not hasattr(func, "_fmt_counter"):
func._fmt_counter = 0
if not call.args:
raise ValueError(
"bpf_printk expects at least one argument (format string)")
raise ValueError("bpf_printk expects at least one argument (format string)")
args = []
if isinstance(call.args[0], ast.JoinedStr):
args = handle_fstring_print(call.args[0], module, builder, func,
local_sym_tab, struct_sym_tab,
local_var_metadata)
elif (isinstance(call.args[0], ast.Constant) and
isinstance(call.args[0].value, str)):
args = handle_fstring_print(
call.args[0],
module,
builder,
func,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
elif isinstance(call.args[0], ast.Constant) and isinstance(call.args[0].value, str):
# TODO: We are only supporting single arguments for now.
# In case of multiple args, the first one will be taken.
args = simple_string_print(call.args[0].value, module, builder, func)
else:
raise NotImplementedError(
"Only simple strings or f-strings are supported in bpf_printk.")
"Only simple strings or f-strings are supported in bpf_printk."
)
fn_type = ir.FunctionType(
ir.IntType(64), [ir.PointerType(), ir.IntType(32)], var_arg=True)
ir.IntType(64), [ir.PointerType(), ir.IntType(32)], var_arg=True
)
fn_ptr_type = ir.PointerType(fn_type)
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PRINTK.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
@ -99,18 +130,25 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("update")
def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func,
local_sym_tab=None, struct_sym_tab=None,
local_var_metadata=None):
def bpf_map_update_elem_emitter(
call,
map_ptr,
module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_map_update_elem helper function call.
Expected call signature: map.update(key, value, flags=0)
"""
if (not call.args or
len(call.args) < 2 or
len(call.args) > 3):
raise ValueError("Map update expects 2 or 3 args (key, value, flags), "
f"got {len(call.args)}")
if not call.args or len(call.args) < 2 or len(call.args) > 3:
raise ValueError(
"Map update expects 2 or 3 args (key, value, flags), "
f"got {len(call.args)}"
)
key_arg = call.args[0]
value_arg = call.args[1]
@ -124,12 +162,11 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func,
fn_type = ir.FunctionType(
ir.IntType(64),
[ir.PointerType(), ir.PointerType(), ir.PointerType(), ir.IntType(64)],
var_arg=False
var_arg=False,
)
fn_ptr_type = ir.PointerType(fn_type)
fn_addr = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
if isinstance(flags_val, int):
@ -138,22 +175,31 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func,
flags_const = flags_val
result = builder.call(
fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False)
fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False
)
return result, None
@HelperHandlerRegistry.register("delete")
def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func,
local_sym_tab=None, struct_sym_tab=None,
local_var_metadata=None):
def bpf_map_delete_elem_emitter(
call,
map_ptr,
module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_map_delete_elem helper function call.
Expected call signature: map.delete(key)
"""
if not call.args or len(call.args) != 1:
raise ValueError("Map delete expects exactly one argument (key), got "
f"{len(call.args)}")
raise ValueError(
"Map delete expects exactly one argument (key), got " f"{len(call.args)}"
)
key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
@ -161,12 +207,11 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func,
fn_type = ir.FunctionType(
ir.IntType(64), # Return type: int64 (status code)
[ir.PointerType(), ir.PointerType()], # Args: (void*, void*)
var_arg=False
var_arg=False,
)
fn_ptr_type = ir.PointerType(fn_type)
fn_addr = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
@ -175,15 +220,21 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("pid")
def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func,
local_sym_tab=None, struct_sym_tab=None,
local_var_metadata=None):
def bpf_get_current_pid_tgid_emitter(
call,
map_ptr,
module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
"""
# func is an arg to just have a uniform signature with other emitters
helper_id = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value)
helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value)
fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
fn_ptr_type = ir.PointerType(fn_type)
fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
@ -196,18 +247,26 @@ def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("output")
def bpf_perf_event_output_handler(call, map_ptr, module, builder, func,
local_sym_tab=None, struct_sym_tab=None,
local_var_metadata=None):
def bpf_perf_event_output_handler(
call,
map_ptr,
module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
if len(call.args) != 1:
raise ValueError("Perf event output expects exactly one argument, "
f"got {len(call.args)}")
raise ValueError(
"Perf event output expects exactly one argument, " f"got {len(call.args)}"
)
data_arg = call.args[0]
ctx_ptr = func.args[0] # First argument to the function is ctx
data_ptr, size_val = get_data_ptr_and_size(data_arg, local_sym_tab,
struct_sym_tab,
local_var_metadata)
data_ptr, size_val = get_data_ptr_and_size(
data_arg, local_sym_tab, struct_sym_tab, local_var_metadata
)
# BPF_F_CURRENT_CPU is -1 in 32 bit
flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
@ -216,36 +275,56 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func,
data_void_ptr = builder.bitcast(data_ptr, ir.PointerType())
fn_type = ir.FunctionType(
ir.IntType(64),
[ir.PointerType(ir.IntType(8)), ir.PointerType(), ir.IntType(64),
ir.PointerType(), ir.IntType(64)],
var_arg=False
[
ir.PointerType(ir.IntType(8)),
ir.PointerType(),
ir.IntType(64),
ir.PointerType(),
ir.IntType(64),
],
var_arg=False,
)
fn_ptr_type = ir.PointerType(fn_type)
# helper id
fn_addr = ir.Constant(ir.IntType(64),
BPFHelperID.BPF_PERF_EVENT_OUTPUT.value)
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
result = builder.call(
fn_ptr,
[ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val],
tail=False)
fn_ptr, [ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], tail=False
)
return result, None
def handle_helper_call(call, module, builder, func,
local_sym_tab=None, map_sym_tab=None,
struct_sym_tab=None, local_var_metadata=None):
def handle_helper_call(
call,
module,
builder,
func,
local_sym_tab=None,
map_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""Process a BPF helper function call and emit the appropriate LLVM IR."""
# Helper function to get map pointer and invoke handler
def invoke_helper(method_name, map_ptr=None):
handler = HelperHandlerRegistry.get_handler(method_name)
if not handler:
raise NotImplementedError(
f"Helper function '{method_name}' is not implemented.")
return handler(call, map_ptr, module, builder, func,
local_sym_tab, struct_sym_tab, local_var_metadata)
f"Helper function '{method_name}' is not implemented."
)
return handler(
call,
map_ptr,
module,
builder,
func,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
# Handle direct function calls (e.g., print(), ktime())
if isinstance(call.func, ast.Name):
@ -265,7 +344,8 @@ def handle_helper_call(call, module, builder, func,
map_name = value.id
else:
raise NotImplementedError(
f"Unsupported map access pattern: {ast.dump(value)}")
f"Unsupported map access pattern: {ast.dump(value)}"
)
# Verify map exists and get pointer
if not map_sym_tab or map_name not in map_sym_tab: