From 4905649700da0473afccba17ae6213aa6c732e51 Mon Sep 17 00:00:00 2001 From: varun-r-mallya Date: Wed, 26 Nov 2025 14:18:40 +0530 Subject: [PATCH] feat:non struct field values can be cast --- BCC-Examples/disksnoop.ipynb | 8 ++- pythonbpf/expr/expr_pass.py | 16 +++-- pythonbpf/type_deducer.py | 2 + .../vmlinux_parser/vmlinux_exports_handler.py | 15 +++- tests/c-form/xdp_test.bpf.c | 72 +++++++++++++++++++ 5 files changed, 103 insertions(+), 10 deletions(-) create mode 100644 tests/c-form/xdp_test.bpf.c diff --git a/BCC-Examples/disksnoop.ipynb b/BCC-Examples/disksnoop.ipynb index 1f81d01..2265a07 100644 --- a/BCC-Examples/disksnoop.ipynb +++ b/BCC-Examples/disksnoop.ipynb @@ -3025,19 +3025,20 @@ ], "source": [ "from vmlinux import struct_request, struct_pt_regs\n", - "from pythonbpf import bpf, section, bpfglobal, compile_to_ir, compile, map, BPF\n", + "from pythonbpf import bpf, section, bpfglobal, map, BPF\n", "from pythonbpf.helper import ktime\n", "from pythonbpf.maps import HashMap\n", - "import logging\n", "from ctypes import c_int64, c_uint64, c_int32\n", "\n", "REQ_WRITE = 1\n", "\n", + "\n", "@bpf\n", "@map\n", "def start() -> HashMap:\n", " return HashMap(key=c_uint64, value=c_uint64, max_entries=10240)\n", "\n", + "\n", "@bpf\n", "@section(\"kprobe/blk_mq_end_request\")\n", "def trace_completion(ctx: struct_pt_regs) -> c_int64:\n", @@ -3063,6 +3064,7 @@ "\n", " return c_int64(0)\n", "\n", + "\n", "@bpf\n", "@section(\"kprobe/blk_mq_start_request\")\n", "def trace_start(ctx1: struct_pt_regs) -> c_int32:\n", @@ -3071,11 +3073,13 @@ " start.update(req, ts)\n", " return c_int32(0)\n", "\n", + "\n", "@bpf\n", "@bpfglobal\n", "def LICENSE() -> str:\n", " return \"GPL\"\n", "\n", + "\n", "b = BPF()" ] }, diff --git a/pythonbpf/expr/expr_pass.py b/pythonbpf/expr/expr_pass.py index 489de9f..abafb86 100644 --- a/pythonbpf/expr/expr_pass.py +++ b/pythonbpf/expr/expr_pass.py @@ -666,13 +666,17 @@ def _handle_vmlinux_cast( # Cast the integer/value to a pointer to the struct # If arg_val is an integer type, we need to inttoptr it ptr_type = ir.PointerType() - # TODO: add a integer check here later - if ctypes_to_ir(arg_type.type.__name__): - # Cast integer to pointer - casted_ptr = builder.inttoptr(arg_val, ptr_type) + # TODO: add a field value type check here + print(arg_type) + if isinstance(arg_type, Field): + if ctypes_to_ir(arg_type.type.__name__): + # Cast integer to pointer + casted_ptr = builder.inttoptr(arg_val, ptr_type) + else: + logger.error(f"Unsupported type for vmlinux cast: {arg_type}") + return None else: - logger.error(f"Unsupported type for vmlinux cast: {arg_type}") - return None + casted_ptr = builder.inttoptr(arg_val, ptr_type) return casted_ptr, vmlinux_struct_type diff --git a/pythonbpf/type_deducer.py b/pythonbpf/type_deducer.py index fd589ae..734b80b 100644 --- a/pythonbpf/type_deducer.py +++ b/pythonbpf/type_deducer.py @@ -18,6 +18,8 @@ mapping = { "c_longlong": ir.IntType(64), "c_uint": ir.IntType(32), "c_int": ir.IntType(32), + "c_ushort": ir.IntType(16), + "c_short": ir.IntType(16), # Not so sure about this one "str": ir.PointerType(ir.IntType(8)), } diff --git a/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py b/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py index c26cac9..e5bc153 100644 --- a/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py +++ b/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py @@ -121,7 +121,12 @@ class VmlinuxHandler: # Use bpf_probe_read_kernel for non-context struct field access field_value = self.load_struct_field( - builder, struct_ptr, globvar_ir, field_data, struct_name + builder, + struct_ptr, + globvar_ir, + field_data, + struct_name, + local_sym_tab, ) # Return field value and field type return field_value, field_data @@ -130,7 +135,12 @@ class VmlinuxHandler: @staticmethod def load_struct_field( - builder, struct_ptr_int, offset_global, field_data, struct_name=None + builder, + struct_ptr_int, + offset_global, + field_data, + struct_name=None, + local_sym_tab=None, ): """ Generate LLVM IR to load a field from a regular (non-context) struct using bpf_probe_read_kernel. @@ -204,6 +214,7 @@ class VmlinuxHandler: logger.warning("Complex vmlinux field type, using default 64 bits") # Allocate local storage for the field value + # TODO: CRITICAL BUG. alloca cannot be used anywhere other than the basic block local_storage = builder.alloca(ir.IntType(int_width)) local_storage_i8_ptr = builder.bitcast(local_storage, i8_ptr_type) diff --git a/tests/c-form/xdp_test.bpf.c b/tests/c-form/xdp_test.bpf.c new file mode 100644 index 0000000..152b0ae --- /dev/null +++ b/tests/c-form/xdp_test.bpf.c @@ -0,0 +1,72 @@ +// xdp_ip_map.c +#include +#include +#include +#include +#include + +struct ip_key { + __u8 family; // 4 = IPv4 + __u8 pad[3]; // padding for alignment + __u8 addr[16]; // IPv4 uses first 4 bytes +}; + +// key → packet count +struct { + __uint(type, BPF_MAP_TYPE_HASH); + __uint(max_entries, 16384); + __type(key, struct ip_key); + __type(value, __u64); +} ip_count_map SEC(".maps"); + +SEC("xdp") +int xdp_ip_map(struct xdp_md *ctx) +{ + void *data_end = (void *)(long)ctx->data_end; + void *data = (void *)(long)ctx->data; + struct ethhdr *eth = data; + + if (eth + 1 > (struct ethhdr *)data_end) + return XDP_PASS; + + __u16 h_proto = eth->h_proto; + void *nh = data + sizeof(*eth); + + // VLAN handling: single tag + if (h_proto == bpf_htons(ETH_P_8021Q) || + h_proto == bpf_htons(ETH_P_8021AD)) { + + if (nh + 4 > data_end) + return XDP_PASS; + + h_proto = *(__u16 *)(nh + 2); + nh += 4; + } + + struct ip_key key = {}; + + // IPv4 + if (h_proto == bpf_htons(ETH_P_IP)) { + struct iphdr *iph = nh; + if (iph + 1 > (struct iphdr *)data_end) + return XDP_PASS; + + key.family = 4; + // Copy 4 bytes of IPv4 address + __builtin_memcpy(key.addr, &iph->saddr, 4); + + __u64 *val = bpf_map_lookup_elem(&ip_count_map, &key); + if (val) + (*val)++; + else { + __u64 init = 1; + bpf_map_update_elem(&ip_count_map, &key, &init, BPF_ANY); + } + + return XDP_PASS; + } + + return XDP_PASS; +} + +char _license[] SEC("license") = "GPL";