mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
Merge branch 'master' into vmlinux-handler
This commit is contained in:
varunrmallya
@ -1,7 +1,7 @@
|
||||
from .helper_registry import HelperHandlerRegistry
|
||||
from .helper_utils import reset_scratch_pool
|
||||
from .bpf_helper_handler import handle_helper_call
|
||||
from .helpers import ktime, pid, deref, XDP_DROP, XDP_PASS
|
||||
from .bpf_helper_handler import handle_helper_call, emit_probe_read_kernel_str_call
|
||||
from .helpers import ktime, pid, deref, comm, probe_read_str, XDP_DROP, XDP_PASS
|
||||
|
||||
|
||||
# Register the helper handler with expr module
|
||||
@ -59,9 +59,12 @@ __all__ = [
|
||||
"HelperHandlerRegistry",
|
||||
"reset_scratch_pool",
|
||||
"handle_helper_call",
|
||||
"emit_probe_read_kernel_str_call",
|
||||
"ktime",
|
||||
"pid",
|
||||
"deref",
|
||||
"comm",
|
||||
"probe_read_str",
|
||||
"XDP_DROP",
|
||||
"XDP_PASS",
|
||||
]
|
||||
|
||||
@ -7,6 +7,9 @@ from .helper_utils import (
|
||||
get_or_create_ptr_from_arg,
|
||||
get_flags_val,
|
||||
get_data_ptr_and_size,
|
||||
get_buffer_ptr_and_size,
|
||||
get_char_array_ptr_and_size,
|
||||
get_ptr_from_arg,
|
||||
)
|
||||
from .printk_formatter import simple_string_print, handle_fstring_print
|
||||
|
||||
@ -23,7 +26,9 @@ class BPFHelperID(Enum):
|
||||
BPF_KTIME_GET_NS = 5
|
||||
BPF_PRINTK = 6
|
||||
BPF_GET_CURRENT_PID_TGID = 14
|
||||
BPF_GET_CURRENT_COMM = 16
|
||||
BPF_PERF_EVENT_OUTPUT = 25
|
||||
BPF_PROBE_READ_KERNEL_STR = 115
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register("ktime")
|
||||
@ -234,6 +239,63 @@ def bpf_map_delete_elem_emitter(
|
||||
return result, None
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register("comm")
|
||||
def bpf_get_current_comm_emitter(
|
||||
call,
|
||||
map_ptr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab=None,
|
||||
struct_sym_tab=None,
|
||||
map_sym_tab=None,
|
||||
):
|
||||
"""
|
||||
Emit LLVM IR for bpf_get_current_comm helper function call.
|
||||
|
||||
Accepts: comm(dataobj.field) or comm(my_buffer)
|
||||
"""
|
||||
if not call.args or len(call.args) != 1:
|
||||
raise ValueError(
|
||||
f"comm expects exactly one argument (buffer), got {len(call.args)}"
|
||||
)
|
||||
|
||||
buf_arg = call.args[0]
|
||||
|
||||
# Extract buffer pointer and size
|
||||
buf_ptr, buf_size = get_buffer_ptr_and_size(
|
||||
buf_arg, builder, local_sym_tab, struct_sym_tab
|
||||
)
|
||||
|
||||
# Validate it's a char array
|
||||
if not isinstance(
|
||||
buf_ptr.type.pointee, ir.ArrayType
|
||||
) or buf_ptr.type.pointee.element != ir.IntType(8):
|
||||
raise ValueError(
|
||||
f"comm expects a char array buffer, got {buf_ptr.type.pointee}"
|
||||
)
|
||||
|
||||
# Cast to void* and call helper
|
||||
buf_void_ptr = builder.bitcast(buf_ptr, ir.PointerType())
|
||||
|
||||
fn_type = ir.FunctionType(
|
||||
ir.IntType(64),
|
||||
[ir.PointerType(), ir.IntType(32)],
|
||||
var_arg=False,
|
||||
)
|
||||
fn_ptr = builder.inttoptr(
|
||||
ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_COMM.value),
|
||||
ir.PointerType(fn_type),
|
||||
)
|
||||
|
||||
result = builder.call(
|
||||
fn_ptr, [buf_void_ptr, ir.Constant(ir.IntType(32), buf_size)], tail=False
|
||||
)
|
||||
|
||||
logger.info(f"Emitted bpf_get_current_comm with {buf_size} byte buffer")
|
||||
return result, None
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register("pid")
|
||||
def bpf_get_current_pid_tgid_emitter(
|
||||
call,
|
||||
@ -309,6 +371,68 @@ def bpf_perf_event_output_handler(
|
||||
return result, None
|
||||
|
||||
|
||||
def emit_probe_read_kernel_str_call(builder, dst_ptr, dst_size, src_ptr):
|
||||
"""Emit LLVM IR call to bpf_probe_read_kernel_str"""
|
||||
|
||||
fn_type = ir.FunctionType(
|
||||
ir.IntType(64),
|
||||
[ir.PointerType(), ir.IntType(32), ir.PointerType()],
|
||||
var_arg=False,
|
||||
)
|
||||
fn_ptr = builder.inttoptr(
|
||||
ir.Constant(ir.IntType(64), BPFHelperID.BPF_PROBE_READ_KERNEL_STR.value),
|
||||
ir.PointerType(fn_type),
|
||||
)
|
||||
|
||||
result = builder.call(
|
||||
fn_ptr,
|
||||
[
|
||||
builder.bitcast(dst_ptr, ir.PointerType()),
|
||||
ir.Constant(ir.IntType(32), dst_size),
|
||||
builder.bitcast(src_ptr, ir.PointerType()),
|
||||
],
|
||||
tail=False,
|
||||
)
|
||||
|
||||
logger.info(f"Emitted bpf_probe_read_kernel_str (size={dst_size})")
|
||||
return result
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register("probe_read_str")
|
||||
def bpf_probe_read_kernel_str_emitter(
|
||||
call,
|
||||
map_ptr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab=None,
|
||||
struct_sym_tab=None,
|
||||
map_sym_tab=None,
|
||||
):
|
||||
"""Emit LLVM IR for bpf_probe_read_kernel_str helper."""
|
||||
|
||||
if len(call.args) != 2:
|
||||
raise ValueError(
|
||||
f"probe_read_str expects 2 args (dst, src), got {len(call.args)}"
|
||||
)
|
||||
|
||||
# Get destination buffer (char array -> i8*)
|
||||
dst_ptr, dst_size = get_char_array_ptr_and_size(
|
||||
call.args[0], builder, local_sym_tab, struct_sym_tab
|
||||
)
|
||||
|
||||
# Get source pointer (evaluate expression)
|
||||
src_ptr, src_type = get_ptr_from_arg(
|
||||
call.args[1], func, module, builder, local_sym_tab, map_sym_tab, struct_sym_tab
|
||||
)
|
||||
|
||||
# Emit the helper call
|
||||
result = emit_probe_read_kernel_str_call(builder, dst_ptr, dst_size, src_ptr)
|
||||
|
||||
logger.info(f"Emitted bpf_probe_read_kernel_str (size={dst_size})")
|
||||
return result, ir.IntType(64)
|
||||
|
||||
|
||||
def handle_helper_call(
|
||||
call,
|
||||
module,
|
||||
|
||||
@ -4,6 +4,7 @@ import logging
|
||||
from llvmlite import ir
|
||||
from pythonbpf.expr import (
|
||||
get_operand_value,
|
||||
eval_expr,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -136,3 +137,140 @@ def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab):
|
||||
raise NotImplementedError(
|
||||
"Only simple object names are supported as data in perf event output."
|
||||
)
|
||||
|
||||
|
||||
def get_buffer_ptr_and_size(buf_arg, builder, local_sym_tab, struct_sym_tab):
|
||||
"""Extract buffer pointer and size from either a struct field or variable."""
|
||||
|
||||
# Case 1: Struct field (obj.field)
|
||||
if isinstance(buf_arg, ast.Attribute):
|
||||
if not isinstance(buf_arg.value, ast.Name):
|
||||
raise ValueError(
|
||||
"Only simple struct field access supported (e.g., obj.field)"
|
||||
)
|
||||
|
||||
struct_name = buf_arg.value.id
|
||||
field_name = buf_arg.attr
|
||||
|
||||
# Lookup struct
|
||||
if not local_sym_tab or struct_name not in local_sym_tab:
|
||||
raise ValueError(f"Struct '{struct_name}' not found")
|
||||
|
||||
struct_type = local_sym_tab[struct_name].metadata
|
||||
if not struct_sym_tab or struct_type not in struct_sym_tab:
|
||||
raise ValueError(f"Struct type '{struct_type}' not found")
|
||||
|
||||
struct_info = struct_sym_tab[struct_type]
|
||||
|
||||
# Get field pointer and type
|
||||
struct_ptr = local_sym_tab[struct_name].var
|
||||
field_ptr = struct_info.gep(builder, struct_ptr, field_name)
|
||||
field_type = struct_info.field_type(field_name)
|
||||
|
||||
if not isinstance(field_type, ir.ArrayType):
|
||||
raise ValueError(f"Field '{field_name}' must be an array type")
|
||||
|
||||
return field_ptr, field_type.count
|
||||
|
||||
# Case 2: Variable name
|
||||
elif isinstance(buf_arg, ast.Name):
|
||||
var_name = buf_arg.id
|
||||
|
||||
if not local_sym_tab or var_name not in local_sym_tab:
|
||||
raise ValueError(f"Variable '{var_name}' not found")
|
||||
|
||||
var_ptr = local_sym_tab[var_name].var
|
||||
var_type = local_sym_tab[var_name].ir_type
|
||||
|
||||
if not isinstance(var_type, ir.ArrayType):
|
||||
raise ValueError(f"Variable '{var_name}' must be an array type")
|
||||
|
||||
return var_ptr, var_type.count
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
"comm expects either a struct field (obj.field) or variable name"
|
||||
)
|
||||
|
||||
|
||||
def get_char_array_ptr_and_size(buf_arg, builder, local_sym_tab, struct_sym_tab):
|
||||
"""Get pointer to char array and its size."""
|
||||
|
||||
# Struct field: obj.field
|
||||
if isinstance(buf_arg, ast.Attribute) and isinstance(buf_arg.value, ast.Name):
|
||||
var_name = buf_arg.value.id
|
||||
field_name = buf_arg.attr
|
||||
|
||||
if not (local_sym_tab and var_name in local_sym_tab):
|
||||
raise ValueError(f"Variable '{var_name}' not found")
|
||||
|
||||
struct_type = local_sym_tab[var_name].metadata
|
||||
if not (struct_sym_tab and struct_type in struct_sym_tab):
|
||||
raise ValueError(f"Struct type '{struct_type}' not found")
|
||||
|
||||
struct_info = struct_sym_tab[struct_type]
|
||||
if field_name not in struct_info.fields:
|
||||
raise ValueError(f"Field '{field_name}' not found")
|
||||
|
||||
field_type = struct_info.field_type(field_name)
|
||||
if not _is_char_array(field_type):
|
||||
raise ValueError("Expected char array field")
|
||||
|
||||
struct_ptr = local_sym_tab[var_name].var
|
||||
field_ptr = struct_info.gep(builder, struct_ptr, field_name)
|
||||
|
||||
# GEP to first element: [N x i8]* -> i8*
|
||||
buf_ptr = builder.gep(
|
||||
field_ptr,
|
||||
[ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), 0)],
|
||||
inbounds=True,
|
||||
)
|
||||
return buf_ptr, field_type.count
|
||||
|
||||
elif isinstance(buf_arg, ast.Name):
|
||||
# NOTE: We shouldn't be doing this as we can't get size info
|
||||
var_name = buf_arg.id
|
||||
if not (local_sym_tab and var_name in local_sym_tab):
|
||||
raise ValueError(f"Variable '{var_name}' not found")
|
||||
|
||||
var_ptr = local_sym_tab[var_name].var
|
||||
var_type = local_sym_tab[var_name].ir_type
|
||||
|
||||
if not isinstance(var_type, ir.PointerType) or not isinstance(
|
||||
var_type.pointee, ir.IntType(8)
|
||||
):
|
||||
raise ValueError("Expected str ptr variable")
|
||||
|
||||
return var_ptr, 256 # Size unknown for str ptr, using 256 as default
|
||||
|
||||
else:
|
||||
raise ValueError("Expected struct field or variable name")
|
||||
|
||||
|
||||
def _is_char_array(ir_type):
|
||||
"""Check if IR type is [N x i8]."""
|
||||
return (
|
||||
isinstance(ir_type, ir.ArrayType)
|
||||
and isinstance(ir_type.element, ir.IntType)
|
||||
and ir_type.element.width == 8
|
||||
)
|
||||
|
||||
|
||||
def get_ptr_from_arg(
|
||||
arg, func, module, builder, local_sym_tab, map_sym_tab, struct_sym_tab
|
||||
):
|
||||
"""Evaluate argument and return pointer value"""
|
||||
|
||||
result = eval_expr(
|
||||
func, module, builder, arg, local_sym_tab, map_sym_tab, struct_sym_tab
|
||||
)
|
||||
|
||||
if not result:
|
||||
raise ValueError("Failed to evaluate argument")
|
||||
|
||||
val, val_type = result
|
||||
|
||||
if not isinstance(val_type, ir.PointerType):
|
||||
raise ValueError(f"Expected pointer type, got {val_type}")
|
||||
|
||||
return val, val_type
|
||||
|
||||
@ -2,19 +2,31 @@ import ctypes
|
||||
|
||||
|
||||
def ktime():
|
||||
"""get current ktime"""
|
||||
return ctypes.c_int64(0)
|
||||
|
||||
|
||||
def pid():
|
||||
"""get current process id"""
|
||||
return ctypes.c_int32(0)
|
||||
|
||||
|
||||
def deref(ptr):
|
||||
"dereference a pointer"
|
||||
"""dereference a pointer"""
|
||||
result = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_void_p)).contents.value
|
||||
return result if result is not None else 0
|
||||
|
||||
|
||||
def comm(buf):
|
||||
"""get current process command name"""
|
||||
return ctypes.c_int64(0)
|
||||
|
||||
|
||||
def probe_read_str(dst, src):
|
||||
"""Safely read a null-terminated string from kernel memory"""
|
||||
return ctypes.c_int64(0)
|
||||
|
||||
|
||||
XDP_ABORTED = ctypes.c_int64(0)
|
||||
XDP_DROP = ctypes.c_int64(1)
|
||||
XDP_PASS = ctypes.c_int64(2)
|
||||
|
||||
@ -184,6 +184,15 @@ def _populate_fval(ftype, node, fmt_parts, exprs):
|
||||
raise NotImplementedError(
|
||||
f"Unsupported pointer target type in f-string: {target}"
|
||||
)
|
||||
elif isinstance(ftype, ir.ArrayType):
|
||||
if isinstance(ftype.element, ir.IntType) and ftype.element.width == 8:
|
||||
# Char array
|
||||
fmt_parts.append("%s")
|
||||
exprs.append(node)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported array element type in f-string: {ftype.element}"
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported field type in f-string: {ftype}")
|
||||
|
||||
@ -208,44 +217,100 @@ def _create_format_string_global(fmt_str, func, module, builder):
|
||||
|
||||
def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_tab):
|
||||
"""Evaluate and prepare an expression to use as an arg for bpf_printk."""
|
||||
val, _ = eval_expr(
|
||||
func,
|
||||
module,
|
||||
builder,
|
||||
expr,
|
||||
local_sym_tab,
|
||||
None,
|
||||
struct_sym_tab,
|
||||
|
||||
# Special case: struct field char array needs pointer to first element
|
||||
char_array_ptr = _get_struct_char_array_ptr(
|
||||
expr, builder, local_sym_tab, struct_sym_tab
|
||||
)
|
||||
if char_array_ptr:
|
||||
return char_array_ptr
|
||||
|
||||
# Regular expression evaluation
|
||||
val, _ = eval_expr(func, module, builder, expr, local_sym_tab, None, struct_sym_tab)
|
||||
|
||||
if not val:
|
||||
logger.warning("Failed to evaluate expression for bpf_printk, defaulting to 0")
|
||||
return ir.Constant(ir.IntType(64), 0)
|
||||
|
||||
# Convert value to bpf_printk compatible type
|
||||
if isinstance(val.type, ir.PointerType):
|
||||
return _handle_pointer_arg(val, func, builder)
|
||||
elif isinstance(val.type, ir.IntType):
|
||||
return _handle_int_arg(val, builder)
|
||||
else:
|
||||
logger.warning(f"Unsupported type {val.type} in bpf_printk, defaulting to 0")
|
||||
return ir.Constant(ir.IntType(64), 0)
|
||||
|
||||
|
||||
def _get_struct_char_array_ptr(expr, builder, local_sym_tab, struct_sym_tab):
|
||||
"""Get pointer to first element of char array in struct field, or None."""
|
||||
if not (isinstance(expr, ast.Attribute) and isinstance(expr.value, ast.Name)):
|
||||
return None
|
||||
|
||||
var_name = expr.value.id
|
||||
field_name = expr.attr
|
||||
|
||||
# Check if it's a valid struct field
|
||||
if not (
|
||||
local_sym_tab
|
||||
and var_name in local_sym_tab
|
||||
and struct_sym_tab
|
||||
and local_sym_tab[var_name].metadata in struct_sym_tab
|
||||
):
|
||||
return None
|
||||
|
||||
struct_type = local_sym_tab[var_name].metadata
|
||||
struct_info = struct_sym_tab[struct_type]
|
||||
|
||||
if field_name not in struct_info.fields:
|
||||
return None
|
||||
|
||||
field_type = struct_info.field_type(field_name)
|
||||
|
||||
# Check if it's a char array
|
||||
is_char_array = (
|
||||
isinstance(field_type, ir.ArrayType)
|
||||
and isinstance(field_type.element, ir.IntType)
|
||||
and field_type.element.width == 8
|
||||
)
|
||||
|
||||
if val:
|
||||
if isinstance(val.type, ir.PointerType):
|
||||
target, depth = get_base_type_and_depth(val.type)
|
||||
if isinstance(target, ir.IntType):
|
||||
if target.width >= 32:
|
||||
val = deref_to_depth(func, builder, val, depth)
|
||||
val = builder.sext(val, ir.IntType(64))
|
||||
elif target.width == 8 and depth == 1:
|
||||
# NOTE: i8* is string, no need to deref
|
||||
pass
|
||||
if not is_char_array:
|
||||
return None
|
||||
|
||||
else:
|
||||
logger.warning(
|
||||
"Only int and ptr supported in bpf_printk args. Others default to 0."
|
||||
)
|
||||
val = ir.Constant(ir.IntType(64), 0)
|
||||
elif isinstance(val.type, ir.IntType):
|
||||
if val.type.width < 64:
|
||||
val = builder.sext(val, ir.IntType(64))
|
||||
else:
|
||||
logger.warning(
|
||||
"Only int and ptr supported in bpf_printk args. Others default to 0."
|
||||
)
|
||||
val = ir.Constant(ir.IntType(64), 0)
|
||||
return val
|
||||
else:
|
||||
logger.warning(
|
||||
"Failed to evaluate expression for bpf_printk argument. "
|
||||
"It will be converted to 0."
|
||||
)
|
||||
# Get field pointer and GEP to first element: [N x i8]* -> i8*
|
||||
struct_ptr = local_sym_tab[var_name].var
|
||||
field_ptr = struct_info.gep(builder, struct_ptr, field_name)
|
||||
|
||||
return builder.gep(
|
||||
field_ptr,
|
||||
[ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), 0)],
|
||||
inbounds=True,
|
||||
)
|
||||
|
||||
|
||||
def _handle_pointer_arg(val, func, builder):
|
||||
"""Convert pointer type for bpf_printk."""
|
||||
target, depth = get_base_type_and_depth(val.type)
|
||||
|
||||
if not isinstance(target, ir.IntType):
|
||||
logger.warning("Only int pointers supported in bpf_printk, defaulting to 0")
|
||||
return ir.Constant(ir.IntType(64), 0)
|
||||
|
||||
# i8* is string - use as-is
|
||||
if target.width == 8 and depth == 1:
|
||||
return val
|
||||
|
||||
# Integer pointers: dereference and sign-extend to i64
|
||||
if target.width >= 32:
|
||||
val = deref_to_depth(func, builder, val, depth)
|
||||
return builder.sext(val, ir.IntType(64))
|
||||
|
||||
logger.warning("Unsupported pointer width in bpf_printk, defaulting to 0")
|
||||
return ir.Constant(ir.IntType(64), 0)
|
||||
|
||||
|
||||
def _handle_int_arg(val, builder):
|
||||
"""Convert integer type for bpf_printk (sign-extend to i64)."""
|
||||
if val.type.width < 64:
|
||||
return builder.sext(val, ir.IntType(64))
|
||||
return val
|
||||
|
||||
Reference in New Issue
Block a user