feat:non struct field values can be cast

Author: varun-r-mallya

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()"
    ]
   },

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
 

pythonbpf/type_deducer.py

@@ -18,6 +18,8 @@
     "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)),
 }

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -121,7 +121,12 @@ def handle_vmlinux_struct_field(
 
                 # 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 @@ def handle_vmlinux_struct_field(
 
     @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 @@ def load_struct_field(
                     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)
 

tests/c-form/xdp_test.bpf.c

@@ -0,0 +1,72 @@
+// xdp_ip_map.c
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_endian.h>
+#include <linux/if_ether.h>
+#include <linux/ip.h>
+
+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";