mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
Simplify maps_pass
This commit is contained in:
Pragyansh Chaturvedi
@ -1,10 +1,11 @@
|
||||
import ast
|
||||
import logging
|
||||
from logging import Logger
|
||||
from llvmlite import ir
|
||||
from enum import Enum
|
||||
|
||||
from .maps_utils import MapProcessorRegistry
|
||||
from pythonbpf.debuginfo import DebugInfoGenerator
|
||||
import logging
|
||||
from .map_types import BPFMapType
|
||||
from .map_debug_info import create_map_debug_info, create_ringbuf_debug_info
|
||||
|
||||
logger: Logger = logging.getLogger(__name__)
|
||||
|
||||
@ -26,44 +27,6 @@ def is_map(func_node):
|
||||
)
|
||||
|
||||
|
||||
class BPFMapType(Enum):
|
||||
UNSPEC = 0
|
||||
HASH = 1
|
||||
ARRAY = 2
|
||||
PROG_ARRAY = 3
|
||||
PERF_EVENT_ARRAY = 4
|
||||
PERCPU_HASH = 5
|
||||
PERCPU_ARRAY = 6
|
||||
STACK_TRACE = 7
|
||||
CGROUP_ARRAY = 8
|
||||
LRU_HASH = 9
|
||||
LRU_PERCPU_HASH = 10
|
||||
LPM_TRIE = 11
|
||||
ARRAY_OF_MAPS = 12
|
||||
HASH_OF_MAPS = 13
|
||||
DEVMAP = 14
|
||||
SOCKMAP = 15
|
||||
CPUMAP = 16
|
||||
XSKMAP = 17
|
||||
SOCKHASH = 18
|
||||
CGROUP_STORAGE_DEPRECATED = 19
|
||||
CGROUP_STORAGE = 19
|
||||
REUSEPORT_SOCKARRAY = 20
|
||||
PERCPU_CGROUP_STORAGE_DEPRECATED = 21
|
||||
PERCPU_CGROUP_STORAGE = 21
|
||||
QUEUE = 22
|
||||
STACK = 23
|
||||
SK_STORAGE = 24
|
||||
DEVMAP_HASH = 25
|
||||
STRUCT_OPS = 26
|
||||
RINGBUF = 27
|
||||
INODE_STORAGE = 28
|
||||
TASK_STORAGE = 29
|
||||
BLOOM_FILTER = 30
|
||||
USER_RINGBUF = 31
|
||||
CGRP_STORAGE = 32
|
||||
|
||||
|
||||
def create_bpf_map(module, map_name, map_params):
|
||||
"""Create a BPF map in the module with given parameters and debug info"""
|
||||
|
||||
@ -84,114 +47,37 @@ def create_bpf_map(module, map_name, map_params):
|
||||
return map_global
|
||||
|
||||
|
||||
def create_map_debug_info(module, map_global, map_name, map_params):
|
||||
"""Generate debug info metadata for BPF maps HASH and PERF_EVENT_ARRAY"""
|
||||
generator = DebugInfoGenerator(module)
|
||||
def _parse_map_params(rval, expected_args=None):
|
||||
"""Parse map parameters from call arguments and keywords."""
|
||||
|
||||
uint_type = generator.get_uint32_type()
|
||||
ulong_type = generator.get_uint64_type()
|
||||
array_type = generator.create_array_type(
|
||||
uint_type, map_params.get("type", BPFMapType.UNSPEC).value
|
||||
)
|
||||
type_ptr = generator.create_pointer_type(array_type, 64)
|
||||
key_ptr = generator.create_pointer_type(
|
||||
array_type if "key_size" in map_params else ulong_type, 64
|
||||
)
|
||||
value_ptr = generator.create_pointer_type(
|
||||
array_type if "value_size" in map_params else ulong_type, 64
|
||||
)
|
||||
params = {}
|
||||
|
||||
elements_arr = []
|
||||
# Parse positional arguments
|
||||
if expected_args:
|
||||
for i, arg_name in enumerate(expected_args):
|
||||
if i < len(rval.args):
|
||||
arg = rval.args[i]
|
||||
if isinstance(arg, ast.Name):
|
||||
params[arg_name] = arg.id
|
||||
elif isinstance(arg, ast.Constant):
|
||||
params[arg_name] = arg.value
|
||||
|
||||
# Create struct members
|
||||
# scope field does not appear for some reason
|
||||
cnt = 0
|
||||
for elem in map_params:
|
||||
if elem == "max_entries":
|
||||
continue
|
||||
if elem == "type":
|
||||
ptr = type_ptr
|
||||
elif "key" in elem:
|
||||
ptr = key_ptr
|
||||
else:
|
||||
ptr = value_ptr
|
||||
# TODO: the best way to do this is not 64, but get the size each time. this will not work for structs.
|
||||
member = generator.create_struct_member(elem, ptr, cnt * 64)
|
||||
elements_arr.append(member)
|
||||
cnt += 1
|
||||
# Parse keyword arguments (override positional)
|
||||
for keyword in rval.keywords:
|
||||
if isinstance(keyword.value, ast.Name):
|
||||
params[keyword.arg] = keyword.value.id
|
||||
elif isinstance(keyword.value, ast.Constant):
|
||||
params[keyword.arg] = keyword.value.value
|
||||
|
||||
if "max_entries" in map_params:
|
||||
max_entries_array = generator.create_array_type(
|
||||
uint_type, map_params["max_entries"]
|
||||
)
|
||||
max_entries_ptr = generator.create_pointer_type(max_entries_array, 64)
|
||||
max_entries_member = generator.create_struct_member(
|
||||
"max_entries", max_entries_ptr, cnt * 64
|
||||
)
|
||||
elements_arr.append(max_entries_member)
|
||||
|
||||
# Create the struct type
|
||||
struct_type = generator.create_struct_type(
|
||||
elements_arr, 64 * len(elements_arr), is_distinct=True
|
||||
)
|
||||
|
||||
# Create global variable debug info
|
||||
global_var = generator.create_global_var_debug_info(
|
||||
map_name, struct_type, is_local=False
|
||||
)
|
||||
|
||||
# Attach debug info to the global variable
|
||||
map_global.set_metadata("dbg", global_var)
|
||||
|
||||
return global_var
|
||||
|
||||
|
||||
def create_ringbuf_debug_info(module, map_global, map_name, map_params):
|
||||
"""Generate debug information metadata for BPF RINGBUF map"""
|
||||
generator = DebugInfoGenerator(module)
|
||||
|
||||
int_type = generator.get_int32_type()
|
||||
|
||||
type_array = generator.create_array_type(
|
||||
int_type, map_params.get("type", BPFMapType.RINGBUF).value
|
||||
)
|
||||
type_ptr = generator.create_pointer_type(type_array, 64)
|
||||
type_member = generator.create_struct_member("type", type_ptr, 0)
|
||||
|
||||
max_entries_array = generator.create_array_type(int_type, map_params["max_entries"])
|
||||
max_entries_ptr = generator.create_pointer_type(max_entries_array, 64)
|
||||
max_entries_member = generator.create_struct_member(
|
||||
"max_entries", max_entries_ptr, 64
|
||||
)
|
||||
|
||||
elements_arr = [type_member, max_entries_member]
|
||||
|
||||
struct_type = generator.create_struct_type(elements_arr, 128, is_distinct=True)
|
||||
|
||||
global_var = generator.create_global_var_debug_info(
|
||||
map_name, struct_type, is_local=False
|
||||
)
|
||||
map_global.set_metadata("dbg", global_var)
|
||||
return global_var
|
||||
return params
|
||||
|
||||
|
||||
@MapProcessorRegistry.register("RingBuf")
|
||||
def process_ringbuf_map(map_name, rval, module):
|
||||
"""Process a BPF_RINGBUF map declaration"""
|
||||
logger.info(f"Processing Ringbuf: {map_name}")
|
||||
map_params = {"type": BPFMapType.RINGBUF}
|
||||
|
||||
# Parse max_entries if present
|
||||
if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Constant):
|
||||
const_val = rval.args[0].value
|
||||
if isinstance(const_val, int):
|
||||
map_params["max_entries"] = const_val
|
||||
|
||||
for keyword in rval.keywords:
|
||||
if keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant):
|
||||
const_val = keyword.value.value
|
||||
if isinstance(const_val, int):
|
||||
map_params["max_entries"] = const_val
|
||||
map_params = _parse_map_params(rval, expected_args=["max_entries"])
|
||||
map_params["type"] = BPFMapType.RINGBUF
|
||||
|
||||
logger.info(f"Ringbuf map parameters: {map_params}")
|
||||
|
||||
@ -204,27 +90,8 @@ def process_ringbuf_map(map_name, rval, module):
|
||||
def process_hash_map(map_name, rval, module):
|
||||
"""Process a BPF_HASH map declaration"""
|
||||
logger.info(f"Processing HashMap: {map_name}")
|
||||
map_params = {"type": BPFMapType.HASH}
|
||||
|
||||
# Assuming order: key_type, value_type, max_entries
|
||||
if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name):
|
||||
map_params["key"] = rval.args[0].id
|
||||
if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name):
|
||||
map_params["value"] = rval.args[1].id
|
||||
if len(rval.args) >= 3 and isinstance(rval.args[2], ast.Constant):
|
||||
const_val = rval.args[2].value
|
||||
if isinstance(const_val, (int, str)): # safe check
|
||||
map_params["max_entries"] = const_val
|
||||
|
||||
for keyword in rval.keywords:
|
||||
if keyword.arg == "key" and isinstance(keyword.value, ast.Name):
|
||||
map_params["key"] = keyword.value.id
|
||||
elif keyword.arg == "value" and isinstance(keyword.value, ast.Name):
|
||||
map_params["value"] = keyword.value.id
|
||||
elif keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant):
|
||||
const_val = keyword.value.value
|
||||
if isinstance(const_val, (int, str)):
|
||||
map_params["max_entries"] = const_val
|
||||
map_params = _parse_map_params(rval, expected_args=["key", "value", "max_entries"])
|
||||
map_params["type"] = BPFMapType.HASH
|
||||
|
||||
logger.info(f"Map parameters: {map_params}")
|
||||
map_global = create_bpf_map(module, map_name, map_params)
|
||||
@ -237,18 +104,8 @@ def process_hash_map(map_name, rval, module):
|
||||
def process_perf_event_map(map_name, rval, module):
|
||||
"""Process a BPF_PERF_EVENT_ARRAY map declaration"""
|
||||
logger.info(f"Processing PerfEventArray: {map_name}")
|
||||
map_params = {"type": BPFMapType.PERF_EVENT_ARRAY}
|
||||
|
||||
if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name):
|
||||
map_params["key_size"] = rval.args[0].id
|
||||
if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name):
|
||||
map_params["value_size"] = rval.args[1].id
|
||||
|
||||
for keyword in rval.keywords:
|
||||
if keyword.arg == "key_size" and isinstance(keyword.value, ast.Name):
|
||||
map_params["key_size"] = keyword.value.id
|
||||
elif keyword.arg == "value_size" and isinstance(keyword.value, ast.Name):
|
||||
map_params["value_size"] = keyword.value.id
|
||||
map_params = _parse_map_params(rval, expected_args=["key_size", "value_size"])
|
||||
map_params["type"] = BPFMapType.PERF_EVENT_ARRAY
|
||||
|
||||
logger.info(f"Map parameters: {map_params}")
|
||||
map_global = create_bpf_map(module, map_name, map_params)
|
||||
|
||||
Reference in New Issue
Block a user