diff --git a/BCC-Examples/hello_perf_output.py b/BCC-Examples/hello_perf_output.py index 3b41d79..03da1ed 100644 --- a/BCC-Examples/hello_perf_output.py +++ b/BCC-Examples/hello_perf_output.py @@ -1,5 +1,5 @@ from pythonbpf import bpf, map, struct, section, bpfglobal, compile -from pythonbpf.helper import ktime, pid +from pythonbpf.helper import ktime, pid, comm from pythonbpf.maps import PerfEventArray from ctypes import c_void_p, c_int64, c_uint64 @@ -25,7 +25,7 @@ def hello(ctx: c_void_p) -> c_int64: dataobj = data_t() strobj = "hellohellohello" dataobj.pid, dataobj.ts = pid(), ktime() - # get_curr_comm(dataobj.comm) + comm(dataobj.comm) print(f"clone called at {dataobj.ts} by pid {dataobj.pid}, comm {strobj}") events.output(dataobj) return 0 # type: ignore [return-value] diff --git a/pythonbpf/helper/__init__.py b/pythonbpf/helper/__init__.py index 9f301b2..26f792b 100644 --- a/pythonbpf/helper/__init__.py +++ b/pythonbpf/helper/__init__.py @@ -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 .helpers import ktime, pid, deref, comm, XDP_DROP, XDP_PASS # Register the helper handler with expr module @@ -62,6 +62,7 @@ __all__ = [ "ktime", "pid", "deref", + "comm", "XDP_DROP", "XDP_PASS", ] diff --git a/pythonbpf/helper/bpf_helper_handler.py b/pythonbpf/helper/bpf_helper_handler.py index e82dfc6..1c96cb4 100644 --- a/pythonbpf/helper/bpf_helper_handler.py +++ b/pythonbpf/helper/bpf_helper_handler.py @@ -7,6 +7,7 @@ from .helper_utils import ( get_or_create_ptr_from_arg, get_flags_val, get_data_ptr_and_size, + get_buffer_ptr_and_size, ) from .printk_formatter import simple_string_print, handle_fstring_print @@ -248,8 +249,48 @@ def bpf_get_current_comm_emitter( ): """ Emit LLVM IR for bpf_get_current_comm helper function call. + + Accepts: comm(dataobj.field) or comm(my_buffer) """ - pass # Not implemented yet + 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") diff --git a/pythonbpf/helper/helper_utils.py b/pythonbpf/helper/helper_utils.py index 4b04464..cf89c30 100644 --- a/pythonbpf/helper/helper_utils.py +++ b/pythonbpf/helper/helper_utils.py @@ -136,3 +136,57 @@ 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" + )