mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
Compare commits
138 Commits
refactor_a
...
v0.1.5
| Author | SHA1 | Date | |
|---|---|---|---|
| 7ae629e8f7 | |||
| dd734ea2aa | |||
| 71d005b6b1 | |||
| 5d9a29ee8e | |||
| 041e538b53 | |||
| 5413cc793b | |||
| f21837aefe | |||
| 0f5c1fa752 | |||
| de02731ea1 | |||
| c22d85ceb8 | |||
| 2b3c81affa | |||
| 8372111616 | |||
| eb4ee64ee5 | |||
| ce7b170fea | |||
| 9a60dd87e3 | |||
| c499fe7421 | |||
| 8239097fbb | |||
| a4cfc2b7aa | |||
| 69b73003ca | |||
| 11e8e72188 | |||
| d3f0e3b2ef | |||
| 09ba749b46 | |||
| a03d3e5d4c | |||
| e1f9ac6ba0 | |||
| 27ab3aaf1e | |||
| b34f7dd68a | |||
| 69d8669e44 | |||
| d4f5a9c36e | |||
| b2a57edf11 | |||
| 20ec307288 | |||
| 0b4c6264a8 | |||
| 6345fcdeff | |||
| 6b41f1fb84 | |||
| 74d8014ade | |||
| 5d0a888542 | |||
| 0042280ff1 | |||
| 7a67041ea3 | |||
| 45e6ce5e5c | |||
| c5f0a2806f | |||
| b0ea93a786 | |||
| fc058c4341 | |||
| 158cc42e1e | |||
| 2a1eabc10d | |||
| e5741562f6 | |||
| 93634a4769 | |||
| 9b8462f1ed | |||
| 785182787c | |||
| 80396c78a6 | |||
| 31645f0316 | |||
| 8774277000 | |||
| 8743ea17f3 | |||
| f8844104a6 | |||
| 3343bedd11 | |||
| 75d3ad4fe2 | |||
| abbf17748d | |||
| 7c559840f0 | |||
| 06773c895f | |||
| 1e3d775865 | |||
| 168e26268e | |||
| 2cf7b28793 | |||
| d24d59c2ba | |||
| f190a33e21 | |||
| eb636ef731 | |||
| 2ae3aade60 | |||
| f227fe9310 | |||
| 7940d02bc7 | |||
| 2483ef2840 | |||
| 68e9693f9a | |||
| e4575a6b1e | |||
| 3ec3ab30fe | |||
| 7fb3ecff48 | |||
| ec59dad025 | |||
| 28b7b1620c | |||
| 9f8e240a38 | |||
| e6c05ab494 | |||
| 8aa9cf7119 | |||
| 9683e3799f | |||
| 200d293750 | |||
| ed196caebf | |||
| a049796b81 | |||
| 384fc9dd40 | |||
| 5f2df57e64 | |||
| 130d8a9edc | |||
| 40ae3d825a | |||
| 484624104e | |||
| e7c4bdb150 | |||
| 7210366e7d | |||
| 435bf27176 | |||
| 1ba27ac7cf | |||
| e4ddec3a02 | |||
| bc7b5c97d1 | |||
| fa720f8e6b | |||
| eff0f66d95 | |||
| b43c252224 | |||
| aae7aa981d | |||
| 6f9a7301af | |||
| 48923d03d4 | |||
| 019a83cf11 | |||
| 140d9e6e35 | |||
| a351b0f1b5 | |||
| 3cb73ff0c3 | |||
| 3b08c2bede | |||
| 86378d6cc4 | |||
| 00d1c583af | |||
| cfc246c80d | |||
| f3c80f9e5f | |||
| 0d3a5748dd | |||
| 079431754c | |||
| 46f5eca33d | |||
| 7081e939fb | |||
| 1e29460d6f | |||
| e180a89644 | |||
| 34a267e982 | |||
| c81aad7c67 | |||
| 2e677c2c7b | |||
| 4ea7b22b44 | |||
| b8b937bfca | |||
| 6cc29c4fa1 | |||
| 5451ba646d | |||
| 7720437ca5 | |||
| eb0a7a917d | |||
| 6f65903552 | |||
| 97e74d09be | |||
| 9c7560ed2e | |||
| 2979ceedcf | |||
| 745f59278f | |||
| 49c59b32ca | |||
| ff78140a7d | |||
| 82ff71b753 | |||
| f46e7cd846 | |||
| 9d73eb67c4 | |||
| 21ce041353 | |||
| 6402cf7be5 | |||
| 9a96e1247b | |||
| 989134f4be | |||
| b95fbd0ed0 | |||
| d84ce0c6fa | |||
| 8d07a4cd05 |
3
.gitignore
vendored
3
.gitignore
vendored
@ -7,3 +7,6 @@ __pycache__/
|
||||
*.ll
|
||||
*.o
|
||||
.ipynb_checkpoints/
|
||||
vmlinux.py
|
||||
~*
|
||||
vmlinux.h
|
||||
|
||||
@ -12,7 +12,7 @@
|
||||
#
|
||||
# See https://github.com/pre-commit/pre-commit
|
||||
|
||||
exclude: 'vmlinux.*\.py$'
|
||||
exclude: 'vmlinux.py'
|
||||
|
||||
ci:
|
||||
autoupdate_commit_msg: "chore: update pre-commit hooks"
|
||||
@ -41,7 +41,7 @@ repos:
|
||||
- id: ruff
|
||||
args: ["--fix", "--show-fixes"]
|
||||
- id: ruff-format
|
||||
exclude: ^(docs)|^(tests)|^(examples)
|
||||
# exclude: ^(docs)|^(tests)|^(examples)
|
||||
|
||||
# Checking static types
|
||||
- repo: https://github.com/pre-commit/mirrors-mypy
|
||||
|
||||
@ -83,14 +83,14 @@ def hist() -> HashMap:
|
||||
def hello(ctx: c_void_p) -> c_int64:
|
||||
process_id = pid()
|
||||
one = 1
|
||||
prev = hist().lookup(process_id)
|
||||
prev = hist.lookup(process_id)
|
||||
if prev:
|
||||
previous_value = prev + 1
|
||||
print(f"count: {previous_value} with {process_id}")
|
||||
hist().update(process_id, previous_value)
|
||||
hist.update(process_id, previous_value)
|
||||
return c_int64(0)
|
||||
else:
|
||||
hist().update(process_id, one)
|
||||
hist.update(process_id, one)
|
||||
return c_int64(0)
|
||||
|
||||
|
||||
|
||||
13
TODO.md
13
TODO.md
@ -1,13 +0,0 @@
|
||||
## Short term
|
||||
|
||||
- Implement enough functionality to port the BCC tutorial examples in PythonBPF
|
||||
- Add all maps
|
||||
- XDP support in pylibbpf
|
||||
- ringbuf support
|
||||
- Add oneline IfExpr conditionals (wishlist)
|
||||
|
||||
## Long term
|
||||
|
||||
- Refactor the codebase to be better than a hackathon project
|
||||
- Port to C++ and use actual LLVM?
|
||||
- Fix struct_kioctx issue in the vmlinux transpiler
|
||||
@ -308,6 +308,7 @@
|
||||
"def hist() -> HashMap:\n",
|
||||
" return HashMap(key=c_int32, value=c_uint64, max_entries=4096)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"@bpf\n",
|
||||
"@section(\"tracepoint/syscalls/sys_enter_clone\")\n",
|
||||
"def hello(ctx: c_void_p) -> c_int64:\n",
|
||||
@ -329,6 +330,7 @@
|
||||
"def LICENSE() -> str:\n",
|
||||
" return \"GPL\"\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"b = BPF()"
|
||||
]
|
||||
},
|
||||
@ -357,7 +359,6 @@
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"\n",
|
||||
"b.load_and_attach()\n",
|
||||
"hist = BpfMap(b, hist)\n",
|
||||
"print(\"Recording\")\n",
|
||||
|
||||
@ -8,12 +8,14 @@ def hello_world(ctx: c_void_p) -> c_int64:
|
||||
print("Hello, World!")
|
||||
return c_int64(0)
|
||||
|
||||
|
||||
@bpf
|
||||
@section("kprobe/do_unlinkat")
|
||||
def hello_world2(ctx: c_void_p) -> c_int64:
|
||||
print("Hello, World!")
|
||||
return c_int64(0)
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def LICENSE() -> str:
|
||||
|
||||
@ -27,7 +27,7 @@ def hello(ctx: c_void_p) -> c_int32:
|
||||
dataobj.pid = pid()
|
||||
dataobj.ts = ktime()
|
||||
# dataobj.comm = strobj
|
||||
print(f"clone called at {dataobj.ts} by pid" f"{dataobj.pid}, comm {strobj}")
|
||||
print(f"clone called at {dataobj.ts} by pid{dataobj.pid}, comm {strobj}")
|
||||
events.output(dataobj)
|
||||
return c_int32(0)
|
||||
|
||||
|
||||
248446
examples/vmlinux.py
248446
examples/vmlinux.py
File diff suppressed because it is too large
Load Diff
@ -1,8 +1,8 @@
|
||||
from pythonbpf import bpf, map, section, bpfglobal, compile
|
||||
from pythonbpf import bpf, map, section, bpfglobal, compile, compile_to_ir
|
||||
from pythonbpf.helper import XDP_PASS
|
||||
from pythonbpf.maps import HashMap
|
||||
from ctypes import c_int64, c_void_p
|
||||
|
||||
from ctypes import c_void_p, c_int64
|
||||
|
||||
# Instructions to how to run this program
|
||||
# 1. Install PythonBPF: pip install pythonbpf
|
||||
@ -41,4 +41,5 @@ def LICENSE() -> str:
|
||||
return "GPL"
|
||||
|
||||
|
||||
compile_to_ir("xdp_pass.py", "xdp_pass.ll")
|
||||
compile()
|
||||
|
||||
@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "pythonbpf"
|
||||
version = "0.1.4"
|
||||
version = "0.1.5"
|
||||
description = "Reduced Python frontend for eBPF"
|
||||
authors = [
|
||||
{ name = "r41k0u", email="pragyanshchaturvedi18@gmail.com" },
|
||||
|
||||
@ -22,19 +22,6 @@ class LocalSymbol:
|
||||
yield self.metadata
|
||||
|
||||
|
||||
def _is_helper_call(call_node):
|
||||
"""Check if a call node is a BPF helper function call."""
|
||||
if isinstance(call_node.func, ast.Name):
|
||||
# Exclude print from requiring temps (handles f-strings differently)
|
||||
func_name = call_node.func.id
|
||||
return HelperHandlerRegistry.has_handler(func_name) and func_name != "print"
|
||||
|
||||
elif isinstance(call_node.func, ast.Attribute):
|
||||
return HelperHandlerRegistry.has_handler(call_node.func.attr)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
|
||||
"""Handle memory allocation for assignment statements."""
|
||||
|
||||
|
||||
@ -1,110 +0,0 @@
|
||||
import ast
|
||||
from llvmlite import ir
|
||||
from logging import Logger
|
||||
import logging
|
||||
|
||||
from pythonbpf.expr import get_base_type_and_depth, deref_to_depth, eval_expr
|
||||
|
||||
logger: Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_operand_value(
|
||||
func, module, operand, builder, local_sym_tab, map_sym_tab, structs_sym_tab=None
|
||||
):
|
||||
"""Extract the value from an operand, handling variables and constants."""
|
||||
logger.info(f"Getting operand value for: {ast.dump(operand)}")
|
||||
if isinstance(operand, ast.Name):
|
||||
if operand.id in local_sym_tab:
|
||||
var = local_sym_tab[operand.id].var
|
||||
var_type = var.type
|
||||
base_type, depth = get_base_type_and_depth(var_type)
|
||||
logger.info(f"var is {var}, base_type is {base_type}, depth is {depth}")
|
||||
val = deref_to_depth(func, builder, var, depth)
|
||||
return val
|
||||
raise ValueError(f"Undefined variable: {operand.id}")
|
||||
elif isinstance(operand, ast.Constant):
|
||||
if isinstance(operand.value, int):
|
||||
cst = ir.Constant(ir.IntType(64), int(operand.value))
|
||||
return cst
|
||||
raise TypeError(f"Unsupported constant type: {type(operand.value)}")
|
||||
elif isinstance(operand, ast.BinOp):
|
||||
res = handle_binary_op_impl(
|
||||
func, module, operand, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
return res
|
||||
else:
|
||||
res = eval_expr(
|
||||
func, module, builder, operand, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
if res is None:
|
||||
raise ValueError(f"Failed to evaluate call expression: {operand}")
|
||||
val, _ = res
|
||||
logger.info(f"Evaluated expr to {val} of type {val.type}")
|
||||
base_type, depth = get_base_type_and_depth(val.type)
|
||||
if depth > 0:
|
||||
val = deref_to_depth(func, builder, val, depth)
|
||||
return val
|
||||
raise TypeError(f"Unsupported operand type: {type(operand)}")
|
||||
|
||||
|
||||
def handle_binary_op_impl(
|
||||
func, module, rval, builder, local_sym_tab, map_sym_tab, structs_sym_tab=None
|
||||
):
|
||||
op = rval.op
|
||||
left = get_operand_value(
|
||||
func, module, rval.left, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
right = get_operand_value(
|
||||
func, module, rval.right, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
logger.info(f"left is {left}, right is {right}, op is {op}")
|
||||
|
||||
# NOTE: Before doing the operation, if the operands are integers
|
||||
# we always extend them to i64. The assignment to LHS will take
|
||||
# care of truncation if needed.
|
||||
if isinstance(left.type, ir.IntType) and left.type.width < 64:
|
||||
left = builder.sext(left, ir.IntType(64))
|
||||
if isinstance(right.type, ir.IntType) and right.type.width < 64:
|
||||
right = builder.sext(right, ir.IntType(64))
|
||||
|
||||
# Map AST operation nodes to LLVM IR builder methods
|
||||
op_map = {
|
||||
ast.Add: builder.add,
|
||||
ast.Sub: builder.sub,
|
||||
ast.Mult: builder.mul,
|
||||
ast.Div: builder.sdiv,
|
||||
ast.Mod: builder.srem,
|
||||
ast.LShift: builder.shl,
|
||||
ast.RShift: builder.lshr,
|
||||
ast.BitOr: builder.or_,
|
||||
ast.BitXor: builder.xor,
|
||||
ast.BitAnd: builder.and_,
|
||||
ast.FloorDiv: builder.udiv,
|
||||
}
|
||||
|
||||
if type(op) in op_map:
|
||||
result = op_map[type(op)](left, right)
|
||||
return result
|
||||
else:
|
||||
raise SyntaxError("Unsupported binary operation")
|
||||
|
||||
|
||||
def handle_binary_op(
|
||||
func,
|
||||
module,
|
||||
rval,
|
||||
builder,
|
||||
var_name,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
structs_sym_tab=None,
|
||||
):
|
||||
result = handle_binary_op_impl(
|
||||
func, module, rval, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
if var_name and var_name in local_sym_tab:
|
||||
logger.info(
|
||||
f"Storing result {result} into variable {local_sym_tab[var_name].var}"
|
||||
)
|
||||
builder.store(result, local_sym_tab[var_name].var)
|
||||
return result, result.type
|
||||
@ -4,6 +4,7 @@ from .license_pass import license_processing
|
||||
from .functions import func_proc
|
||||
from .maps import maps_proc
|
||||
from .structs import structs_proc
|
||||
from .vmlinux_parser import vmlinux_proc
|
||||
from .globals_pass import (
|
||||
globals_list_creation,
|
||||
globals_processing,
|
||||
@ -18,10 +19,20 @@ from pylibbpf import BpfProgram
|
||||
import tempfile
|
||||
from logging import Logger
|
||||
import logging
|
||||
import re
|
||||
|
||||
logger: Logger = logging.getLogger(__name__)
|
||||
|
||||
VERSION = "v0.1.4"
|
||||
VERSION = "v0.1.5"
|
||||
|
||||
|
||||
def finalize_module(original_str):
|
||||
"""After all IR generation is complete, we monkey patch btf_ama attribute"""
|
||||
|
||||
# Create a string with applied transformation of btf_ama attribute addition to BTF struct field accesses.
|
||||
pattern = r'(@"llvm\.[^"]+:[^"]*" = external global i64, !llvm\.preserve\.access\.index ![0-9]+)'
|
||||
replacement = r'\1 "btf_ama"'
|
||||
return re.sub(pattern, replacement, original_str)
|
||||
|
||||
|
||||
def find_bpf_chunks(tree):
|
||||
@ -44,6 +55,7 @@ def processor(source_code, filename, module):
|
||||
for func_node in bpf_chunks:
|
||||
logger.info(f"Found BPF function/struct: {func_node.name}")
|
||||
|
||||
vmlinux_proc(tree, module)
|
||||
populate_global_symbol_table(tree, module)
|
||||
license_processing(tree, module)
|
||||
globals_processing(tree, module)
|
||||
@ -119,15 +131,44 @@ def compile_to_ir(filename: str, output: str, loglevel=logging.INFO):
|
||||
|
||||
module.add_named_metadata("llvm.ident", [f"PythonBPF {VERSION}"])
|
||||
|
||||
module_string = finalize_module(str(module))
|
||||
|
||||
logger.info(f"IR written to {output}")
|
||||
with open(output, "w") as f:
|
||||
f.write(f'source_filename = "{filename}"\n')
|
||||
f.write(str(module))
|
||||
f.write(module_string)
|
||||
f.write("\n")
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _run_llc(ll_file, obj_file):
|
||||
"""Compile LLVM IR to BPF object file using llc."""
|
||||
|
||||
logger.info(f"Compiling IR to object: {ll_file} -> {obj_file}")
|
||||
result = subprocess.run(
|
||||
[
|
||||
"llc",
|
||||
"-march=bpf",
|
||||
"-filetype=obj",
|
||||
"-O2",
|
||||
str(ll_file),
|
||||
"-o",
|
||||
str(obj_file),
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
if result.returncode == 0:
|
||||
logger.info(f"Object file written to {obj_file}")
|
||||
return True
|
||||
else:
|
||||
logger.error(f"llc compilation failed: {result.stderr}")
|
||||
return False
|
||||
|
||||
|
||||
def compile(loglevel=logging.INFO) -> bool:
|
||||
# Look one level up the stack to the caller of this function
|
||||
caller_frame = inspect.stack()[1]
|
||||
@ -141,21 +182,7 @@ def compile(loglevel=logging.INFO) -> bool:
|
||||
compile_to_ir(str(caller_file), str(ll_file), loglevel=loglevel) and success
|
||||
)
|
||||
|
||||
success = bool(
|
||||
subprocess.run(
|
||||
[
|
||||
"llc",
|
||||
"-march=bpf",
|
||||
"-filetype=obj",
|
||||
"-O2",
|
||||
str(ll_file),
|
||||
"-o",
|
||||
str(o_file),
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
and success
|
||||
)
|
||||
success = _run_llc(ll_file, o_file) and success
|
||||
|
||||
logger.info(f"Object written to {o_file}")
|
||||
return success
|
||||
@ -175,17 +202,6 @@ def BPF(loglevel=logging.INFO) -> BpfProgram:
|
||||
f.flush()
|
||||
source = f.name
|
||||
compile_to_ir(source, str(inter.name), loglevel=loglevel)
|
||||
subprocess.run(
|
||||
[
|
||||
"llc",
|
||||
"-march=bpf",
|
||||
"-filetype=obj",
|
||||
"-O2",
|
||||
str(inter.name),
|
||||
"-o",
|
||||
str(obj_file.name),
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
_run_llc(str(inter.name), str(obj_file.name))
|
||||
|
||||
return BpfProgram(str(obj_file.name))
|
||||
|
||||
@ -1,5 +1,7 @@
|
||||
from .expr_pass import eval_expr, handle_expr
|
||||
from .type_normalization import convert_to_bool, get_base_type_and_depth, deref_to_depth
|
||||
from .expr_pass import eval_expr, handle_expr, get_operand_value
|
||||
from .type_normalization import convert_to_bool, get_base_type_and_depth
|
||||
from .ir_ops import deref_to_depth
|
||||
from .call_registry import CallHandlerRegistry
|
||||
|
||||
__all__ = [
|
||||
"eval_expr",
|
||||
@ -7,4 +9,6 @@ __all__ = [
|
||||
"convert_to_bool",
|
||||
"get_base_type_and_depth",
|
||||
"deref_to_depth",
|
||||
"get_operand_value",
|
||||
"CallHandlerRegistry",
|
||||
]
|
||||
|
||||
20
pythonbpf/expr/call_registry.py
Normal file
20
pythonbpf/expr/call_registry.py
Normal file
@ -0,0 +1,20 @@
|
||||
class CallHandlerRegistry:
|
||||
"""Registry for handling different types of calls (helpers, etc.)"""
|
||||
|
||||
_handler = None
|
||||
|
||||
@classmethod
|
||||
def set_handler(cls, handler):
|
||||
"""Set the handler for unknown calls"""
|
||||
cls._handler = handler
|
||||
|
||||
@classmethod
|
||||
def handle_call(
|
||||
cls, call, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
):
|
||||
"""Handle a call using the registered handler"""
|
||||
if cls._handler is None:
|
||||
return None
|
||||
return cls._handler(
|
||||
call, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
@ -5,10 +5,20 @@ import logging
|
||||
from typing import Dict
|
||||
|
||||
from pythonbpf.type_deducer import ctypes_to_ir, is_ctypes
|
||||
from .type_normalization import convert_to_bool, handle_comparator
|
||||
from .call_registry import CallHandlerRegistry
|
||||
from .type_normalization import (
|
||||
convert_to_bool,
|
||||
handle_comparator,
|
||||
get_base_type_and_depth,
|
||||
deref_to_depth,
|
||||
)
|
||||
|
||||
logger: Logger = logging.getLogger(__name__)
|
||||
|
||||
# ============================================================================
|
||||
# Leaf Handlers (No Recursive eval_expr calls)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _handle_name_expr(expr: ast.Name, local_sym_tab: Dict, builder: ir.IRBuilder):
|
||||
"""Handle ast.Name expressions."""
|
||||
@ -21,10 +31,24 @@ def _handle_name_expr(expr: ast.Name, local_sym_tab: Dict, builder: ir.IRBuilder
|
||||
return None
|
||||
|
||||
|
||||
def _handle_constant_expr(expr: ast.Constant):
|
||||
def _handle_constant_expr(module, builder, expr: ast.Constant):
|
||||
"""Handle ast.Constant expressions."""
|
||||
if isinstance(expr.value, int) or isinstance(expr.value, bool):
|
||||
return ir.Constant(ir.IntType(64), int(expr.value)), ir.IntType(64)
|
||||
elif isinstance(expr.value, str):
|
||||
str_name = f".str.{id(expr)}"
|
||||
str_bytes = expr.value.encode("utf-8") + b"\x00"
|
||||
str_type = ir.ArrayType(ir.IntType(8), len(str_bytes))
|
||||
str_constant = ir.Constant(str_type, bytearray(str_bytes))
|
||||
|
||||
# Create global variable
|
||||
global_str = ir.GlobalVariable(module, str_type, name=str_name)
|
||||
global_str.linkage = "internal"
|
||||
global_str.global_constant = True
|
||||
global_str.initializer = str_constant
|
||||
|
||||
str_ptr = builder.bitcast(global_str, ir.PointerType(ir.IntType(8)))
|
||||
return str_ptr, ir.PointerType(ir.IntType(8))
|
||||
else:
|
||||
logger.error(f"Unsupported constant type {ast.dump(expr)}")
|
||||
return None
|
||||
@ -88,6 +112,118 @@ def _handle_deref_call(expr: ast.Call, local_sym_tab: Dict, builder: ir.IRBuilde
|
||||
return val, local_sym_tab[arg.id].ir_type
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Binary Operations
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def get_operand_value(
|
||||
func, module, operand, builder, local_sym_tab, map_sym_tab, structs_sym_tab=None
|
||||
):
|
||||
"""Extract the value from an operand, handling variables and constants."""
|
||||
logger.info(f"Getting operand value for: {ast.dump(operand)}")
|
||||
if isinstance(operand, ast.Name):
|
||||
if operand.id in local_sym_tab:
|
||||
var = local_sym_tab[operand.id].var
|
||||
var_type = var.type
|
||||
base_type, depth = get_base_type_and_depth(var_type)
|
||||
logger.info(f"var is {var}, base_type is {base_type}, depth is {depth}")
|
||||
val = deref_to_depth(func, builder, var, depth)
|
||||
return val
|
||||
raise ValueError(f"Undefined variable: {operand.id}")
|
||||
elif isinstance(operand, ast.Constant):
|
||||
if isinstance(operand.value, int):
|
||||
cst = ir.Constant(ir.IntType(64), int(operand.value))
|
||||
return cst
|
||||
raise TypeError(f"Unsupported constant type: {type(operand.value)}")
|
||||
elif isinstance(operand, ast.BinOp):
|
||||
res = _handle_binary_op_impl(
|
||||
func, module, operand, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
return res
|
||||
else:
|
||||
res = eval_expr(
|
||||
func, module, builder, operand, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
if res is None:
|
||||
raise ValueError(f"Failed to evaluate call expression: {operand}")
|
||||
val, _ = res
|
||||
logger.info(f"Evaluated expr to {val} of type {val.type}")
|
||||
base_type, depth = get_base_type_and_depth(val.type)
|
||||
if depth > 0:
|
||||
val = deref_to_depth(func, builder, val, depth)
|
||||
return val
|
||||
raise TypeError(f"Unsupported operand type: {type(operand)}")
|
||||
|
||||
|
||||
def _handle_binary_op_impl(
|
||||
func, module, rval, builder, local_sym_tab, map_sym_tab, structs_sym_tab=None
|
||||
):
|
||||
op = rval.op
|
||||
left = get_operand_value(
|
||||
func, module, rval.left, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
right = get_operand_value(
|
||||
func, module, rval.right, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
logger.info(f"left is {left}, right is {right}, op is {op}")
|
||||
|
||||
# NOTE: Before doing the operation, if the operands are integers
|
||||
# we always extend them to i64. The assignment to LHS will take
|
||||
# care of truncation if needed.
|
||||
if isinstance(left.type, ir.IntType) and left.type.width < 64:
|
||||
left = builder.sext(left, ir.IntType(64))
|
||||
if isinstance(right.type, ir.IntType) and right.type.width < 64:
|
||||
right = builder.sext(right, ir.IntType(64))
|
||||
|
||||
# Map AST operation nodes to LLVM IR builder methods
|
||||
op_map = {
|
||||
ast.Add: builder.add,
|
||||
ast.Sub: builder.sub,
|
||||
ast.Mult: builder.mul,
|
||||
ast.Div: builder.sdiv,
|
||||
ast.Mod: builder.srem,
|
||||
ast.LShift: builder.shl,
|
||||
ast.RShift: builder.lshr,
|
||||
ast.BitOr: builder.or_,
|
||||
ast.BitXor: builder.xor,
|
||||
ast.BitAnd: builder.and_,
|
||||
ast.FloorDiv: builder.udiv,
|
||||
}
|
||||
|
||||
if type(op) in op_map:
|
||||
result = op_map[type(op)](left, right)
|
||||
return result
|
||||
else:
|
||||
raise SyntaxError("Unsupported binary operation")
|
||||
|
||||
|
||||
def _handle_binary_op(
|
||||
func,
|
||||
module,
|
||||
rval,
|
||||
builder,
|
||||
var_name,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
structs_sym_tab=None,
|
||||
):
|
||||
result = _handle_binary_op_impl(
|
||||
func, module, rval, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
if var_name and var_name in local_sym_tab:
|
||||
logger.info(
|
||||
f"Storing result {result} into variable {local_sym_tab[var_name].var}"
|
||||
)
|
||||
builder.store(result, local_sym_tab[var_name].var)
|
||||
return result, result.type
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Comparison and Unary Operations
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _handle_ctypes_call(
|
||||
func,
|
||||
module,
|
||||
@ -180,8 +316,6 @@ def _handle_unary_op(
|
||||
logger.error("Only 'not' and '-' unary operators are supported")
|
||||
return None
|
||||
|
||||
from pythonbpf.binary_ops import get_operand_value
|
||||
|
||||
operand = get_operand_value(
|
||||
func, module, expr.operand, builder, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
@ -200,6 +334,11 @@ def _handle_unary_op(
|
||||
return result, ir.IntType(64)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Boolean Operations
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _handle_and_op(func, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab):
|
||||
"""Handle `and` boolean operations."""
|
||||
|
||||
@ -330,6 +469,11 @@ def _handle_boolean_op(
|
||||
return None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Expression Dispatcher
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def eval_expr(
|
||||
func,
|
||||
module,
|
||||
@ -343,7 +487,7 @@ def eval_expr(
|
||||
if isinstance(expr, ast.Name):
|
||||
return _handle_name_expr(expr, local_sym_tab, builder)
|
||||
elif isinstance(expr, ast.Constant):
|
||||
return _handle_constant_expr(expr)
|
||||
return _handle_constant_expr(module, builder, expr)
|
||||
elif isinstance(expr, ast.Call):
|
||||
if isinstance(expr.func, ast.Name) and expr.func.id == "deref":
|
||||
return _handle_deref_call(expr, local_sym_tab, builder)
|
||||
@ -359,57 +503,18 @@ def eval_expr(
|
||||
structs_sym_tab,
|
||||
)
|
||||
|
||||
# delayed import to avoid circular dependency
|
||||
from pythonbpf.helper import HelperHandlerRegistry, handle_helper_call
|
||||
result = CallHandlerRegistry.handle_call(
|
||||
expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
if result is not None:
|
||||
return result
|
||||
|
||||
if isinstance(expr.func, ast.Name) and HelperHandlerRegistry.has_handler(
|
||||
expr.func.id
|
||||
):
|
||||
return handle_helper_call(
|
||||
expr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
elif isinstance(expr.func, ast.Attribute):
|
||||
logger.info(f"Handling method call: {ast.dump(expr.func)}")
|
||||
if isinstance(expr.func.value, ast.Call) and isinstance(
|
||||
expr.func.value.func, ast.Name
|
||||
):
|
||||
method_name = expr.func.attr
|
||||
if HelperHandlerRegistry.has_handler(method_name):
|
||||
return handle_helper_call(
|
||||
expr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
elif isinstance(expr.func.value, ast.Name):
|
||||
obj_name = expr.func.value.id
|
||||
method_name = expr.func.attr
|
||||
if obj_name in map_sym_tab:
|
||||
if HelperHandlerRegistry.has_handler(method_name):
|
||||
return handle_helper_call(
|
||||
expr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
logger.warning(f"Unknown call: {ast.dump(expr)}")
|
||||
return None
|
||||
elif isinstance(expr, ast.Attribute):
|
||||
return _handle_attribute_expr(expr, local_sym_tab, structs_sym_tab, builder)
|
||||
elif isinstance(expr, ast.BinOp):
|
||||
from pythonbpf.binary_ops import handle_binary_op
|
||||
|
||||
return handle_binary_op(
|
||||
return _handle_binary_op(
|
||||
func,
|
||||
module,
|
||||
expr,
|
||||
|
||||
50
pythonbpf/expr/ir_ops.py
Normal file
50
pythonbpf/expr/ir_ops.py
Normal file
@ -0,0 +1,50 @@
|
||||
import logging
|
||||
from llvmlite import ir
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def deref_to_depth(func, builder, val, target_depth):
|
||||
"""Dereference a pointer to a certain depth."""
|
||||
|
||||
cur_val = val
|
||||
cur_type = val.type
|
||||
|
||||
for depth in range(target_depth):
|
||||
if not isinstance(val.type, ir.PointerType):
|
||||
logger.error("Cannot dereference further, non-pointer type")
|
||||
return None
|
||||
|
||||
# dereference with null check
|
||||
pointee_type = cur_type.pointee
|
||||
null_check_block = builder.block
|
||||
not_null_block = func.append_basic_block(name=f"deref_not_null_{depth}")
|
||||
merge_block = func.append_basic_block(name=f"deref_merge_{depth}")
|
||||
|
||||
null_ptr = ir.Constant(cur_type, None)
|
||||
is_not_null = builder.icmp_signed("!=", cur_val, null_ptr)
|
||||
logger.debug(f"Inserted null check for pointer at depth {depth}")
|
||||
|
||||
builder.cbranch(is_not_null, not_null_block, merge_block)
|
||||
|
||||
builder.position_at_end(not_null_block)
|
||||
dereferenced_val = builder.load(cur_val)
|
||||
logger.debug(f"Dereferenced to depth {depth - 1}, type: {pointee_type}")
|
||||
builder.branch(merge_block)
|
||||
|
||||
builder.position_at_end(merge_block)
|
||||
phi = builder.phi(pointee_type, name=f"deref_result_{depth}")
|
||||
|
||||
zero_value = (
|
||||
ir.Constant(pointee_type, 0)
|
||||
if isinstance(pointee_type, ir.IntType)
|
||||
else ir.Constant(pointee_type, None)
|
||||
)
|
||||
phi.add_incoming(zero_value, null_check_block)
|
||||
|
||||
phi.add_incoming(dereferenced_val, not_null_block)
|
||||
|
||||
# Continue with phi result
|
||||
cur_val = phi
|
||||
cur_type = pointee_type
|
||||
return cur_val
|
||||
@ -1,6 +1,7 @@
|
||||
from llvmlite import ir
|
||||
import logging
|
||||
import ast
|
||||
from llvmlite import ir
|
||||
from .ir_ops import deref_to_depth
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -26,52 +27,6 @@ def get_base_type_and_depth(ir_type):
|
||||
return cur_type, depth
|
||||
|
||||
|
||||
def deref_to_depth(func, builder, val, target_depth):
|
||||
"""Dereference a pointer to a certain depth."""
|
||||
|
||||
cur_val = val
|
||||
cur_type = val.type
|
||||
|
||||
for depth in range(target_depth):
|
||||
if not isinstance(val.type, ir.PointerType):
|
||||
logger.error("Cannot dereference further, non-pointer type")
|
||||
return None
|
||||
|
||||
# dereference with null check
|
||||
pointee_type = cur_type.pointee
|
||||
null_check_block = builder.block
|
||||
not_null_block = func.append_basic_block(name=f"deref_not_null_{depth}")
|
||||
merge_block = func.append_basic_block(name=f"deref_merge_{depth}")
|
||||
|
||||
null_ptr = ir.Constant(cur_type, None)
|
||||
is_not_null = builder.icmp_signed("!=", cur_val, null_ptr)
|
||||
logger.debug(f"Inserted null check for pointer at depth {depth}")
|
||||
|
||||
builder.cbranch(is_not_null, not_null_block, merge_block)
|
||||
|
||||
builder.position_at_end(not_null_block)
|
||||
dereferenced_val = builder.load(cur_val)
|
||||
logger.debug(f"Dereferenced to depth {depth - 1}, type: {pointee_type}")
|
||||
builder.branch(merge_block)
|
||||
|
||||
builder.position_at_end(merge_block)
|
||||
phi = builder.phi(pointee_type, name=f"deref_result_{depth}")
|
||||
|
||||
zero_value = (
|
||||
ir.Constant(pointee_type, 0)
|
||||
if isinstance(pointee_type, ir.IntType)
|
||||
else ir.Constant(pointee_type, None)
|
||||
)
|
||||
phi.add_incoming(zero_value, null_check_block)
|
||||
|
||||
phi.add_incoming(dereferenced_val, not_null_block)
|
||||
|
||||
# Continue with phi result
|
||||
cur_val = phi
|
||||
cur_type = pointee_type
|
||||
return cur_val
|
||||
|
||||
|
||||
def _normalize_types(func, builder, lhs, rhs):
|
||||
"""Normalize types for comparison."""
|
||||
|
||||
|
||||
@ -1,22 +0,0 @@
|
||||
from typing import Dict
|
||||
|
||||
|
||||
class StatementHandlerRegistry:
|
||||
"""Registry for statement handlers."""
|
||||
|
||||
_handlers: Dict = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, stmt_type):
|
||||
"""Register a handler for a specific statement type."""
|
||||
|
||||
def decorator(handler):
|
||||
cls._handlers[stmt_type] = handler
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def __getitem__(cls, stmt_type):
|
||||
"""Get the handler for a specific statement type."""
|
||||
return cls._handlers.get(stmt_type, None)
|
||||
88
pythonbpf/functions/function_metadata.py
Normal file
88
pythonbpf/functions/function_metadata.py
Normal file
@ -0,0 +1,88 @@
|
||||
import ast
|
||||
|
||||
|
||||
def get_probe_string(func_node):
|
||||
"""Extract the probe string from the decorator of the function node"""
|
||||
# TODO: right now we have the whole string in the section decorator
|
||||
# But later we can implement typed tuples for tracepoints and kprobes
|
||||
# For helper functions, we return "helper"
|
||||
|
||||
for decorator in func_node.decorator_list:
|
||||
if isinstance(decorator, ast.Name) and decorator.id == "bpfglobal":
|
||||
return None
|
||||
if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Name):
|
||||
if decorator.func.id == "section" and len(decorator.args) == 1:
|
||||
arg = decorator.args[0]
|
||||
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
|
||||
return arg.value
|
||||
return "helper"
|
||||
|
||||
|
||||
def is_global_function(func_node):
|
||||
"""Check if the function is a global"""
|
||||
for decorator in func_node.decorator_list:
|
||||
if isinstance(decorator, ast.Name) and decorator.id in (
|
||||
"map",
|
||||
"bpfglobal",
|
||||
"struct",
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def infer_return_type(func_node: ast.FunctionDef):
|
||||
if not isinstance(func_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
raise TypeError("Expected ast.FunctionDef")
|
||||
if func_node.returns is not None:
|
||||
try:
|
||||
return ast.unparse(func_node.returns)
|
||||
except Exception:
|
||||
node = func_node.returns
|
||||
if isinstance(node, ast.Name):
|
||||
return node.id
|
||||
if isinstance(node, ast.Attribute):
|
||||
return getattr(node, "attr", type(node).__name__)
|
||||
try:
|
||||
return str(node)
|
||||
except Exception:
|
||||
return type(node).__name__
|
||||
found_type = None
|
||||
|
||||
def _expr_type(e):
|
||||
if e is None:
|
||||
return "None"
|
||||
if isinstance(e, ast.Constant):
|
||||
return type(e.value).__name__
|
||||
if isinstance(e, ast.Name):
|
||||
return e.id
|
||||
if isinstance(e, ast.Call):
|
||||
f = e.func
|
||||
if isinstance(f, ast.Name):
|
||||
return f.id
|
||||
if isinstance(f, ast.Attribute):
|
||||
try:
|
||||
return ast.unparse(f)
|
||||
except Exception:
|
||||
return getattr(f, "attr", type(f).__name__)
|
||||
try:
|
||||
return ast.unparse(f)
|
||||
except Exception:
|
||||
return type(f).__name__
|
||||
if isinstance(e, ast.Attribute):
|
||||
try:
|
||||
return ast.unparse(e)
|
||||
except Exception:
|
||||
return getattr(e, "attr", type(e).__name__)
|
||||
try:
|
||||
return ast.unparse(e)
|
||||
except Exception:
|
||||
return type(e).__name__
|
||||
|
||||
for walked_node in ast.walk(func_node):
|
||||
if isinstance(walked_node, ast.Return):
|
||||
t = _expr_type(walked_node.value)
|
||||
if found_type is None:
|
||||
found_type = t
|
||||
elif found_type != t:
|
||||
raise ValueError(f"Conflicting return types: {found_type} vs {t}")
|
||||
return found_type or "None"
|
||||
@ -14,27 +14,125 @@ from pythonbpf.assign_pass import (
|
||||
)
|
||||
from pythonbpf.allocation_pass import handle_assign_allocation, allocate_temp_pool
|
||||
|
||||
from .return_utils import _handle_none_return, _handle_xdp_return, _is_xdp_name
|
||||
from .return_utils import handle_none_return, handle_xdp_return, is_xdp_name
|
||||
from .function_metadata import get_probe_string, is_global_function, infer_return_type
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_probe_string(func_node):
|
||||
"""Extract the probe string from the decorator of the function node."""
|
||||
# TODO: right now we have the whole string in the section decorator
|
||||
# But later we can implement typed tuples for tracepoints and kprobes
|
||||
# For helper functions, we return "helper"
|
||||
# ============================================================================
|
||||
# SECTION 1: Memory Allocation
|
||||
# ============================================================================
|
||||
|
||||
for decorator in func_node.decorator_list:
|
||||
if isinstance(decorator, ast.Name) and decorator.id == "bpfglobal":
|
||||
return None
|
||||
if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Name):
|
||||
if decorator.func.id == "section" and len(decorator.args) == 1:
|
||||
arg = decorator.args[0]
|
||||
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
|
||||
return arg.value
|
||||
return "helper"
|
||||
|
||||
def count_temps_in_call(call_node, local_sym_tab):
|
||||
"""Count the number of temporary variables needed for a function call."""
|
||||
|
||||
count = 0
|
||||
is_helper = False
|
||||
|
||||
# NOTE: We exclude print calls for now
|
||||
if isinstance(call_node.func, ast.Name):
|
||||
if (
|
||||
HelperHandlerRegistry.has_handler(call_node.func.id)
|
||||
and call_node.func.id != "print"
|
||||
):
|
||||
is_helper = True
|
||||
elif isinstance(call_node.func, ast.Attribute):
|
||||
if HelperHandlerRegistry.has_handler(call_node.func.attr):
|
||||
is_helper = True
|
||||
|
||||
if not is_helper:
|
||||
return 0
|
||||
|
||||
for arg in call_node.args:
|
||||
# NOTE: Count all non-name arguments
|
||||
# For struct fields, if it is being passed as an argument,
|
||||
# The struct object should already exist in the local_sym_tab
|
||||
if not isinstance(arg, ast.Name) and not (
|
||||
isinstance(arg, ast.Attribute) and arg.value.id in local_sym_tab
|
||||
):
|
||||
count += 1
|
||||
|
||||
return count
|
||||
|
||||
|
||||
def handle_if_allocation(
|
||||
module, builder, stmt, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
|
||||
):
|
||||
"""Recursively handle allocations in if/else branches."""
|
||||
if stmt.body:
|
||||
allocate_mem(
|
||||
module,
|
||||
builder,
|
||||
stmt.body,
|
||||
func,
|
||||
ret_type,
|
||||
map_sym_tab,
|
||||
local_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
if stmt.orelse:
|
||||
allocate_mem(
|
||||
module,
|
||||
builder,
|
||||
stmt.orelse,
|
||||
func,
|
||||
ret_type,
|
||||
map_sym_tab,
|
||||
local_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
|
||||
|
||||
def allocate_mem(
|
||||
module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
|
||||
):
|
||||
max_temps_needed = 0
|
||||
|
||||
def update_max_temps_for_stmt(stmt):
|
||||
nonlocal max_temps_needed
|
||||
temps_needed = 0
|
||||
|
||||
if isinstance(stmt, ast.If):
|
||||
for s in stmt.body:
|
||||
update_max_temps_for_stmt(s)
|
||||
for s in stmt.orelse:
|
||||
update_max_temps_for_stmt(s)
|
||||
return
|
||||
|
||||
for node in ast.walk(stmt):
|
||||
if isinstance(node, ast.Call):
|
||||
temps_needed += count_temps_in_call(node, local_sym_tab)
|
||||
max_temps_needed = max(max_temps_needed, temps_needed)
|
||||
|
||||
for stmt in body:
|
||||
update_max_temps_for_stmt(stmt)
|
||||
|
||||
# Handle allocations
|
||||
if isinstance(stmt, ast.If):
|
||||
handle_if_allocation(
|
||||
module,
|
||||
builder,
|
||||
stmt,
|
||||
func,
|
||||
ret_type,
|
||||
map_sym_tab,
|
||||
local_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
elif isinstance(stmt, ast.Assign):
|
||||
handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab)
|
||||
|
||||
allocate_temp_pool(builder, max_temps_needed, local_sym_tab)
|
||||
|
||||
return local_sym_tab
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SECTION 2: Statement Handlers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def handle_assign(
|
||||
@ -146,9 +244,9 @@ def handle_if(
|
||||
def handle_return(builder, stmt, local_sym_tab, ret_type):
|
||||
logger.info(f"Handling return statement: {ast.dump(stmt)}")
|
||||
if stmt.value is None:
|
||||
return _handle_none_return(builder)
|
||||
elif isinstance(stmt.value, ast.Name) and _is_xdp_name(stmt.value.id):
|
||||
return _handle_xdp_return(stmt, builder, ret_type)
|
||||
return handle_none_return(builder)
|
||||
elif isinstance(stmt.value, ast.Name) and is_xdp_name(stmt.value.id):
|
||||
return handle_xdp_return(stmt, builder, ret_type)
|
||||
else:
|
||||
val = eval_expr(
|
||||
func=None,
|
||||
@ -207,108 +305,9 @@ def process_stmt(
|
||||
return did_return
|
||||
|
||||
|
||||
def handle_if_allocation(
|
||||
module, builder, stmt, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
|
||||
):
|
||||
"""Recursively handle allocations in if/else branches."""
|
||||
if stmt.body:
|
||||
allocate_mem(
|
||||
module,
|
||||
builder,
|
||||
stmt.body,
|
||||
func,
|
||||
ret_type,
|
||||
map_sym_tab,
|
||||
local_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
if stmt.orelse:
|
||||
allocate_mem(
|
||||
module,
|
||||
builder,
|
||||
stmt.orelse,
|
||||
func,
|
||||
ret_type,
|
||||
map_sym_tab,
|
||||
local_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
|
||||
|
||||
def count_temps_in_call(call_node, local_sym_tab):
|
||||
"""Count the number of temporary variables needed for a function call."""
|
||||
|
||||
count = 0
|
||||
is_helper = False
|
||||
|
||||
# NOTE: We exclude print calls for now
|
||||
if isinstance(call_node.func, ast.Name):
|
||||
if (
|
||||
HelperHandlerRegistry.has_handler(call_node.func.id)
|
||||
and call_node.func.id != "print"
|
||||
):
|
||||
is_helper = True
|
||||
elif isinstance(call_node.func, ast.Attribute):
|
||||
if HelperHandlerRegistry.has_handler(call_node.func.attr):
|
||||
is_helper = True
|
||||
|
||||
if not is_helper:
|
||||
return 0
|
||||
|
||||
for arg in call_node.args:
|
||||
# NOTE: Count all non-name arguments
|
||||
# For struct fields, if it is being passed as an argument,
|
||||
# The struct object should already exist in the local_sym_tab
|
||||
if not isinstance(arg, ast.Name) and not (
|
||||
isinstance(arg, ast.Attribute) and arg.value.id in local_sym_tab
|
||||
):
|
||||
count += 1
|
||||
|
||||
return count
|
||||
|
||||
|
||||
def allocate_mem(
|
||||
module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
|
||||
):
|
||||
max_temps_needed = 0
|
||||
|
||||
def update_max_temps_for_stmt(stmt):
|
||||
nonlocal max_temps_needed
|
||||
temps_needed = 0
|
||||
|
||||
if isinstance(stmt, ast.If):
|
||||
for s in stmt.body:
|
||||
update_max_temps_for_stmt(s)
|
||||
for s in stmt.orelse:
|
||||
update_max_temps_for_stmt(s)
|
||||
return
|
||||
|
||||
for node in ast.walk(stmt):
|
||||
if isinstance(node, ast.Call):
|
||||
temps_needed += count_temps_in_call(node, local_sym_tab)
|
||||
max_temps_needed = max(max_temps_needed, temps_needed)
|
||||
|
||||
for stmt in body:
|
||||
update_max_temps_for_stmt(stmt)
|
||||
|
||||
# Handle allocations
|
||||
if isinstance(stmt, ast.If):
|
||||
handle_if_allocation(
|
||||
module,
|
||||
builder,
|
||||
stmt,
|
||||
func,
|
||||
ret_type,
|
||||
map_sym_tab,
|
||||
local_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
elif isinstance(stmt, ast.Assign):
|
||||
handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab)
|
||||
|
||||
allocate_temp_pool(builder, max_temps_needed, local_sym_tab)
|
||||
|
||||
return local_sym_tab
|
||||
# ============================================================================
|
||||
# SECTION 3: Function Body Processing
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def process_func_body(
|
||||
@ -390,18 +389,14 @@ def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_t
|
||||
return func
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SECTION 4: Top-Level Function Processor
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab):
|
||||
for func_node in chunks:
|
||||
is_global = False
|
||||
for decorator in func_node.decorator_list:
|
||||
if isinstance(decorator, ast.Name) and decorator.id in (
|
||||
"map",
|
||||
"bpfglobal",
|
||||
"struct",
|
||||
):
|
||||
is_global = True
|
||||
break
|
||||
if is_global:
|
||||
if is_global_function(func_node):
|
||||
continue
|
||||
func_type = get_probe_string(func_node)
|
||||
logger.info(f"Found probe_string of {func_node.name}: {func_type}")
|
||||
@ -415,67 +410,7 @@ def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab):
|
||||
)
|
||||
|
||||
|
||||
def infer_return_type(func_node: ast.FunctionDef):
|
||||
if not isinstance(func_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
raise TypeError("Expected ast.FunctionDef")
|
||||
if func_node.returns is not None:
|
||||
try:
|
||||
return ast.unparse(func_node.returns)
|
||||
except Exception:
|
||||
node = func_node.returns
|
||||
if isinstance(node, ast.Name):
|
||||
return node.id
|
||||
if isinstance(node, ast.Attribute):
|
||||
return getattr(node, "attr", type(node).__name__)
|
||||
try:
|
||||
return str(node)
|
||||
except Exception:
|
||||
return type(node).__name__
|
||||
found_type = None
|
||||
|
||||
def _expr_type(e):
|
||||
if e is None:
|
||||
return "None"
|
||||
if isinstance(e, ast.Constant):
|
||||
return type(e.value).__name__
|
||||
if isinstance(e, ast.Name):
|
||||
return e.id
|
||||
if isinstance(e, ast.Call):
|
||||
f = e.func
|
||||
if isinstance(f, ast.Name):
|
||||
return f.id
|
||||
if isinstance(f, ast.Attribute):
|
||||
try:
|
||||
return ast.unparse(f)
|
||||
except Exception:
|
||||
return getattr(f, "attr", type(f).__name__)
|
||||
try:
|
||||
return ast.unparse(f)
|
||||
except Exception:
|
||||
return type(f).__name__
|
||||
if isinstance(e, ast.Attribute):
|
||||
try:
|
||||
return ast.unparse(e)
|
||||
except Exception:
|
||||
return getattr(e, "attr", type(e).__name__)
|
||||
try:
|
||||
return ast.unparse(e)
|
||||
except Exception:
|
||||
return type(e).__name__
|
||||
|
||||
for walked_node in ast.walk(func_node):
|
||||
if isinstance(walked_node, ast.Return):
|
||||
t = _expr_type(walked_node.value)
|
||||
if found_type is None:
|
||||
found_type = t
|
||||
elif found_type != t:
|
||||
raise ValueError(f"Conflicting return types: {found_type} vs {t}")
|
||||
return found_type or "None"
|
||||
|
||||
|
||||
# For string assignment to fixed-size arrays
|
||||
|
||||
|
||||
# TODO: WIP, for string assignment to fixed-size arrays
|
||||
def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_length):
|
||||
"""
|
||||
Copy a string (i8*) to a fixed-size array ([N x i8]*)
|
||||
|
||||
@ -14,19 +14,19 @@ XDP_ACTIONS = {
|
||||
}
|
||||
|
||||
|
||||
def _handle_none_return(builder) -> bool:
|
||||
def handle_none_return(builder) -> bool:
|
||||
"""Handle return or return None -> returns 0."""
|
||||
builder.ret(ir.Constant(ir.IntType(64), 0))
|
||||
logger.debug("Generated default return: 0")
|
||||
return True
|
||||
|
||||
|
||||
def _is_xdp_name(name: str) -> bool:
|
||||
def is_xdp_name(name: str) -> bool:
|
||||
"""Check if a name is an XDP action"""
|
||||
return name in XDP_ACTIONS
|
||||
|
||||
|
||||
def _handle_xdp_return(stmt: ast.Return, builder, ret_type) -> bool:
|
||||
def handle_xdp_return(stmt: ast.Return, builder, ret_type) -> bool:
|
||||
"""Handle XDP returns"""
|
||||
if not isinstance(stmt.value, ast.Name):
|
||||
return False
|
||||
@ -37,7 +37,6 @@ def _handle_xdp_return(stmt: ast.Return, builder, ret_type) -> bool:
|
||||
raise ValueError(
|
||||
f"Unknown XDP action: {action_name}. Available: {XDP_ACTIONS.keys()}"
|
||||
)
|
||||
return False
|
||||
|
||||
value = XDP_ACTIONS[action_name]
|
||||
builder.ret(ir.Constant(ret_type, value))
|
||||
|
||||
@ -1,7 +1,60 @@
|
||||
from .helper_utils import HelperHandlerRegistry, reset_scratch_pool
|
||||
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
|
||||
|
||||
|
||||
# Register the helper handler with expr module
|
||||
def _register_helper_handler():
|
||||
"""Register helper call handler with the expression evaluator"""
|
||||
from pythonbpf.expr.expr_pass import CallHandlerRegistry
|
||||
|
||||
def helper_call_handler(
|
||||
call, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
):
|
||||
"""Check if call is a helper and handle it"""
|
||||
import ast
|
||||
|
||||
# Check for direct helper calls (e.g., ktime(), print())
|
||||
if isinstance(call.func, ast.Name):
|
||||
if HelperHandlerRegistry.has_handler(call.func.id):
|
||||
return handle_helper_call(
|
||||
call,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
|
||||
# Check for method calls (e.g., map.lookup())
|
||||
elif isinstance(call.func, ast.Attribute):
|
||||
method_name = call.func.attr
|
||||
|
||||
# Handle: my_map.lookup(key)
|
||||
if isinstance(call.func.value, ast.Name):
|
||||
obj_name = call.func.value.id
|
||||
if map_sym_tab and obj_name in map_sym_tab:
|
||||
if HelperHandlerRegistry.has_handler(method_name):
|
||||
return handle_helper_call(
|
||||
call,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
structs_sym_tab,
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
CallHandlerRegistry.set_handler(helper_call_handler)
|
||||
|
||||
|
||||
# Register on module import
|
||||
_register_helper_handler()
|
||||
|
||||
__all__ = [
|
||||
"HelperHandlerRegistry",
|
||||
"reset_scratch_pool",
|
||||
|
||||
@ -1,14 +1,15 @@
|
||||
import ast
|
||||
from llvmlite import ir
|
||||
from enum import Enum
|
||||
|
||||
from .helper_registry import HelperHandlerRegistry
|
||||
from .helper_utils import (
|
||||
HelperHandlerRegistry,
|
||||
get_or_create_ptr_from_arg,
|
||||
get_flags_val,
|
||||
handle_fstring_print,
|
||||
simple_string_print,
|
||||
get_data_ptr_and_size,
|
||||
)
|
||||
from .printk_formatter import simple_string_print, handle_fstring_print
|
||||
|
||||
from logging import Logger
|
||||
import logging
|
||||
|
||||
@ -135,7 +136,7 @@ def bpf_printk_emitter(
|
||||
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
||||
|
||||
builder.call(fn_ptr, args, tail=True)
|
||||
return None
|
||||
return True
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register("update")
|
||||
|
||||
27
pythonbpf/helper/helper_registry.py
Normal file
27
pythonbpf/helper/helper_registry.py
Normal file
@ -0,0 +1,27 @@
|
||||
from typing import Callable
|
||||
|
||||
|
||||
class HelperHandlerRegistry:
|
||||
"""Registry for BPF helpers"""
|
||||
|
||||
_handlers: dict[str, Callable] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, helper_name):
|
||||
"""Decorator to register a handler function for a helper"""
|
||||
|
||||
def decorator(func):
|
||||
cls._handlers[helper_name] = func
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def get_handler(cls, helper_name):
|
||||
"""Get the handler function for a helper"""
|
||||
return cls._handlers.get(helper_name)
|
||||
|
||||
@classmethod
|
||||
def has_handler(cls, helper_name):
|
||||
"""Check if a handler function is registered for a helper"""
|
||||
return helper_name in cls._handlers
|
||||
@ -1,40 +1,14 @@
|
||||
import ast
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
|
||||
from llvmlite import ir
|
||||
from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
|
||||
from pythonbpf.binary_ops import get_operand_value
|
||||
from pythonbpf.expr import (
|
||||
get_operand_value,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HelperHandlerRegistry:
|
||||
"""Registry for BPF helpers"""
|
||||
|
||||
_handlers: dict[str, Callable] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, helper_name):
|
||||
"""Decorator to register a handler function for a helper"""
|
||||
|
||||
def decorator(func):
|
||||
cls._handlers[helper_name] = func
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def get_handler(cls, helper_name):
|
||||
"""Get the handler function for a helper"""
|
||||
return cls._handlers.get(helper_name)
|
||||
|
||||
@classmethod
|
||||
def has_handler(cls, helper_name):
|
||||
"""Check if a handler function is registered for a helper"""
|
||||
return helper_name in cls._handlers
|
||||
|
||||
|
||||
class ScratchPoolManager:
|
||||
"""Manage the temporary helper variables in local_sym_tab"""
|
||||
|
||||
@ -70,6 +44,11 @@ def reset_scratch_pool():
|
||||
_temp_pool_manager.reset()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Argument Preparation
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def get_var_ptr_from_name(var_name, local_sym_tab):
|
||||
"""Get a pointer to a variable from the symbol table."""
|
||||
if local_sym_tab and var_name in local_sym_tab:
|
||||
@ -134,234 +113,6 @@ def get_flags_val(arg, builder, local_sym_tab):
|
||||
)
|
||||
|
||||
|
||||
def simple_string_print(string_value, module, builder, func):
|
||||
"""Prepare arguments for bpf_printk from a simple string value"""
|
||||
fmt_str = string_value + "\n\0"
|
||||
fmt_ptr = _create_format_string_global(fmt_str, func, module, builder)
|
||||
|
||||
args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]
|
||||
return args
|
||||
|
||||
|
||||
def handle_fstring_print(
|
||||
joined_str,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab=None,
|
||||
struct_sym_tab=None,
|
||||
):
|
||||
"""Handle f-string formatting for bpf_printk emitter."""
|
||||
fmt_parts = []
|
||||
exprs = []
|
||||
|
||||
for value in joined_str.values:
|
||||
logger.debug(f"Processing f-string value: {ast.dump(value)}")
|
||||
|
||||
if isinstance(value, ast.Constant):
|
||||
_process_constant_in_fstring(value, fmt_parts, exprs)
|
||||
elif isinstance(value, ast.FormattedValue):
|
||||
_process_fval(
|
||||
value,
|
||||
fmt_parts,
|
||||
exprs,
|
||||
local_sym_tab,
|
||||
struct_sym_tab,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported f-string value type: {type(value)}")
|
||||
|
||||
fmt_str = "".join(fmt_parts)
|
||||
args = simple_string_print(fmt_str, module, builder, func)
|
||||
|
||||
# NOTE: Process expressions (limited to 3 due to BPF constraints)
|
||||
if len(exprs) > 3:
|
||||
logger.warning("bpf_printk supports up to 3 args, extra args will be ignored.")
|
||||
|
||||
for expr in exprs[:3]:
|
||||
arg_value = _prepare_expr_args(
|
||||
expr,
|
||||
func,
|
||||
module,
|
||||
builder,
|
||||
local_sym_tab,
|
||||
struct_sym_tab,
|
||||
)
|
||||
args.append(arg_value)
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def _process_constant_in_fstring(cst, fmt_parts, exprs):
|
||||
"""Process constant values in f-string."""
|
||||
if isinstance(cst.value, str):
|
||||
fmt_parts.append(cst.value)
|
||||
elif isinstance(cst.value, int):
|
||||
fmt_parts.append("%lld")
|
||||
exprs.append(ir.Constant(ir.IntType(64), cst.value))
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported constant type in f-string: {type(cst.value)}"
|
||||
)
|
||||
|
||||
|
||||
def _process_fval(fval, fmt_parts, exprs, local_sym_tab, struct_sym_tab):
|
||||
"""Process formatted values in f-string."""
|
||||
logger.debug(f"Processing formatted value: {ast.dump(fval)}")
|
||||
|
||||
if isinstance(fval.value, ast.Name):
|
||||
_process_name_in_fval(fval.value, fmt_parts, exprs, local_sym_tab)
|
||||
elif isinstance(fval.value, ast.Attribute):
|
||||
_process_attr_in_fval(
|
||||
fval.value,
|
||||
fmt_parts,
|
||||
exprs,
|
||||
local_sym_tab,
|
||||
struct_sym_tab,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported formatted value in f-string: {type(fval.value)}"
|
||||
)
|
||||
|
||||
|
||||
def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab):
|
||||
"""Process name nodes in formatted values."""
|
||||
if local_sym_tab and name_node.id in local_sym_tab:
|
||||
_, var_type, tmp = local_sym_tab[name_node.id]
|
||||
_populate_fval(var_type, name_node, fmt_parts, exprs)
|
||||
|
||||
|
||||
def _process_attr_in_fval(attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab):
|
||||
"""Process attribute nodes in formatted values."""
|
||||
if (
|
||||
isinstance(attr_node.value, ast.Name)
|
||||
and local_sym_tab
|
||||
and attr_node.value.id in local_sym_tab
|
||||
):
|
||||
var_name = attr_node.value.id
|
||||
field_name = attr_node.attr
|
||||
|
||||
var_type = local_sym_tab[var_name].metadata
|
||||
if var_type not in struct_sym_tab:
|
||||
raise ValueError(
|
||||
f"Struct '{var_type}' for '{var_name}' not in symbol table"
|
||||
)
|
||||
|
||||
struct_info = struct_sym_tab[var_type]
|
||||
if field_name not in struct_info.fields:
|
||||
raise ValueError(f"Field '{field_name}' not found in struct '{var_type}'")
|
||||
|
||||
field_type = struct_info.field_type(field_name)
|
||||
_populate_fval(field_type, attr_node, fmt_parts, exprs)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Only simple attribute on local vars is supported in f-strings."
|
||||
)
|
||||
|
||||
|
||||
def _populate_fval(ftype, node, fmt_parts, exprs):
|
||||
"""Populate format parts and expressions based on field type."""
|
||||
if isinstance(ftype, ir.IntType):
|
||||
# TODO: We print as signed integers only for now
|
||||
if ftype.width == 64:
|
||||
fmt_parts.append("%lld")
|
||||
exprs.append(node)
|
||||
elif ftype.width == 32:
|
||||
fmt_parts.append("%d")
|
||||
exprs.append(node)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported integer width in f-string: {ftype.width}"
|
||||
)
|
||||
elif isinstance(ftype, ir.PointerType):
|
||||
target, depth = get_base_type_and_depth(ftype)
|
||||
if isinstance(target, ir.IntType):
|
||||
if target.width == 64:
|
||||
fmt_parts.append("%lld")
|
||||
exprs.append(node)
|
||||
elif target.width == 32:
|
||||
fmt_parts.append("%d")
|
||||
exprs.append(node)
|
||||
elif target.width == 8 and depth == 1:
|
||||
# NOTE: Assume i8* is a string
|
||||
fmt_parts.append("%s")
|
||||
exprs.append(node)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported pointer target type in f-string: {target}"
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported pointer target type in f-string: {target}"
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported field type in f-string: {ftype}")
|
||||
|
||||
|
||||
def _create_format_string_global(fmt_str, func, module, builder):
|
||||
"""Create a global variable for the format string."""
|
||||
fmt_name = f"{func.name}____fmt{func._fmt_counter}"
|
||||
func._fmt_counter += 1
|
||||
|
||||
fmt_gvar = ir.GlobalVariable(
|
||||
module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name
|
||||
)
|
||||
fmt_gvar.global_constant = True
|
||||
fmt_gvar.initializer = ir.Constant(
|
||||
ir.ArrayType(ir.IntType(8), len(fmt_str)), bytearray(fmt_str.encode("utf8"))
|
||||
)
|
||||
fmt_gvar.linkage = "internal"
|
||||
fmt_gvar.align = 1
|
||||
|
||||
return builder.bitcast(fmt_gvar, ir.PointerType())
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
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."
|
||||
)
|
||||
return ir.Constant(ir.IntType(64), 0)
|
||||
|
||||
|
||||
def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab):
|
||||
"""Extract data pointer and size information for perf event output."""
|
||||
if isinstance(data_arg, ast.Name):
|
||||
|
||||
@ -15,5 +15,8 @@ def deref(ptr):
|
||||
return result if result is not None else 0
|
||||
|
||||
|
||||
XDP_ABORTED = ctypes.c_int64(0)
|
||||
XDP_DROP = ctypes.c_int64(1)
|
||||
XDP_PASS = ctypes.c_int64(2)
|
||||
XDP_TX = ctypes.c_int64(3)
|
||||
XDP_REDIRECT = ctypes.c_int64(4)
|
||||
|
||||
240
pythonbpf/helper/printk_formatter.py
Normal file
240
pythonbpf/helper/printk_formatter.py
Normal file
@ -0,0 +1,240 @@
|
||||
import ast
|
||||
import logging
|
||||
|
||||
from llvmlite import ir
|
||||
from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def simple_string_print(string_value, module, builder, func):
|
||||
"""Prepare arguments for bpf_printk from a simple string value"""
|
||||
fmt_str = string_value + "\n\0"
|
||||
fmt_ptr = _create_format_string_global(fmt_str, func, module, builder)
|
||||
|
||||
args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]
|
||||
return args
|
||||
|
||||
|
||||
def handle_fstring_print(
|
||||
joined_str,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab=None,
|
||||
struct_sym_tab=None,
|
||||
):
|
||||
"""Handle f-string formatting for bpf_printk emitter."""
|
||||
fmt_parts = []
|
||||
exprs = []
|
||||
|
||||
for value in joined_str.values:
|
||||
logger.debug(f"Processing f-string value: {ast.dump(value)}")
|
||||
|
||||
if isinstance(value, ast.Constant):
|
||||
_process_constant_in_fstring(value, fmt_parts, exprs)
|
||||
elif isinstance(value, ast.FormattedValue):
|
||||
_process_fval(
|
||||
value,
|
||||
fmt_parts,
|
||||
exprs,
|
||||
local_sym_tab,
|
||||
struct_sym_tab,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported f-string value type: {type(value)}")
|
||||
|
||||
fmt_str = "".join(fmt_parts)
|
||||
args = simple_string_print(fmt_str, module, builder, func)
|
||||
|
||||
# NOTE: Process expressions (limited to 3 due to BPF constraints)
|
||||
if len(exprs) > 3:
|
||||
logger.warning("bpf_printk supports up to 3 args, extra args will be ignored.")
|
||||
|
||||
for expr in exprs[:3]:
|
||||
arg_value = _prepare_expr_args(
|
||||
expr,
|
||||
func,
|
||||
module,
|
||||
builder,
|
||||
local_sym_tab,
|
||||
struct_sym_tab,
|
||||
)
|
||||
args.append(arg_value)
|
||||
|
||||
return args
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Internal Helpers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _process_constant_in_fstring(cst, fmt_parts, exprs):
|
||||
"""Process constant values in f-string."""
|
||||
if isinstance(cst.value, str):
|
||||
fmt_parts.append(cst.value)
|
||||
elif isinstance(cst.value, int):
|
||||
fmt_parts.append("%lld")
|
||||
exprs.append(ir.Constant(ir.IntType(64), cst.value))
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported constant type in f-string: {type(cst.value)}"
|
||||
)
|
||||
|
||||
|
||||
def _process_fval(fval, fmt_parts, exprs, local_sym_tab, struct_sym_tab):
|
||||
"""Process formatted values in f-string."""
|
||||
logger.debug(f"Processing formatted value: {ast.dump(fval)}")
|
||||
|
||||
if isinstance(fval.value, ast.Name):
|
||||
_process_name_in_fval(fval.value, fmt_parts, exprs, local_sym_tab)
|
||||
elif isinstance(fval.value, ast.Attribute):
|
||||
_process_attr_in_fval(
|
||||
fval.value,
|
||||
fmt_parts,
|
||||
exprs,
|
||||
local_sym_tab,
|
||||
struct_sym_tab,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported formatted value in f-string: {type(fval.value)}"
|
||||
)
|
||||
|
||||
|
||||
def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab):
|
||||
"""Process name nodes in formatted values."""
|
||||
if local_sym_tab and name_node.id in local_sym_tab:
|
||||
_, var_type, tmp = local_sym_tab[name_node.id]
|
||||
_populate_fval(var_type, name_node, fmt_parts, exprs)
|
||||
|
||||
|
||||
def _process_attr_in_fval(attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab):
|
||||
"""Process attribute nodes in formatted values."""
|
||||
if (
|
||||
isinstance(attr_node.value, ast.Name)
|
||||
and local_sym_tab
|
||||
and attr_node.value.id in local_sym_tab
|
||||
):
|
||||
var_name = attr_node.value.id
|
||||
field_name = attr_node.attr
|
||||
|
||||
var_type = local_sym_tab[var_name].metadata
|
||||
if var_type not in struct_sym_tab:
|
||||
raise ValueError(
|
||||
f"Struct '{var_type}' for '{var_name}' not in symbol table"
|
||||
)
|
||||
|
||||
struct_info = struct_sym_tab[var_type]
|
||||
if field_name not in struct_info.fields:
|
||||
raise ValueError(f"Field '{field_name}' not found in struct '{var_type}'")
|
||||
|
||||
field_type = struct_info.field_type(field_name)
|
||||
_populate_fval(field_type, attr_node, fmt_parts, exprs)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Only simple attribute on local vars is supported in f-strings."
|
||||
)
|
||||
|
||||
|
||||
def _populate_fval(ftype, node, fmt_parts, exprs):
|
||||
"""Populate format parts and expressions based on field type."""
|
||||
if isinstance(ftype, ir.IntType):
|
||||
# TODO: We print as signed integers only for now
|
||||
if ftype.width == 64:
|
||||
fmt_parts.append("%lld")
|
||||
exprs.append(node)
|
||||
elif ftype.width == 32:
|
||||
fmt_parts.append("%d")
|
||||
exprs.append(node)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported integer width in f-string: {ftype.width}"
|
||||
)
|
||||
elif isinstance(ftype, ir.PointerType):
|
||||
target, depth = get_base_type_and_depth(ftype)
|
||||
if isinstance(target, ir.IntType):
|
||||
if target.width == 64:
|
||||
fmt_parts.append("%lld")
|
||||
exprs.append(node)
|
||||
elif target.width == 32:
|
||||
fmt_parts.append("%d")
|
||||
exprs.append(node)
|
||||
elif target.width == 8 and depth == 1:
|
||||
# NOTE: Assume i8* is a string
|
||||
fmt_parts.append("%s")
|
||||
exprs.append(node)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported pointer target type in f-string: {target}"
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported pointer target type in f-string: {target}"
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported field type in f-string: {ftype}")
|
||||
|
||||
|
||||
def _create_format_string_global(fmt_str, func, module, builder):
|
||||
"""Create a global variable for the format string."""
|
||||
fmt_name = f"{func.name}____fmt{func._fmt_counter}"
|
||||
func._fmt_counter += 1
|
||||
|
||||
fmt_gvar = ir.GlobalVariable(
|
||||
module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name
|
||||
)
|
||||
fmt_gvar.global_constant = True
|
||||
fmt_gvar.initializer = ir.Constant(
|
||||
ir.ArrayType(ir.IntType(8), len(fmt_str)), bytearray(fmt_str.encode("utf8"))
|
||||
)
|
||||
fmt_gvar.linkage = "internal"
|
||||
fmt_gvar.align = 1
|
||||
|
||||
return builder.bitcast(fmt_gvar, ir.PointerType())
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
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."
|
||||
)
|
||||
return ir.Constant(ir.IntType(64), 0)
|
||||
93
pythonbpf/maps/map_debug_info.py
Normal file
93
pythonbpf/maps/map_debug_info.py
Normal file
@ -0,0 +1,93 @@
|
||||
from pythonbpf.debuginfo import DebugInfoGenerator
|
||||
from .map_types import BPFMapType
|
||||
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
elements_arr = []
|
||||
|
||||
# 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
|
||||
|
||||
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
|
||||
39
pythonbpf/maps/map_types.py
Normal file
39
pythonbpf/maps/map_types.py
Normal file
@ -0,0 +1,39 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
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
|
||||
@ -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)
|
||||
|
||||
3
pythonbpf/vmlinux_parser/__init__.py
Normal file
3
pythonbpf/vmlinux_parser/__init__.py
Normal file
@ -0,0 +1,3 @@
|
||||
from .import_detector import vmlinux_proc
|
||||
|
||||
__all__ = ["vmlinux_proc"]
|
||||
205
pythonbpf/vmlinux_parser/class_handler.py
Normal file
205
pythonbpf/vmlinux_parser/class_handler.py
Normal file
@ -0,0 +1,205 @@
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
import importlib
|
||||
from .dependency_handler import DependencyHandler
|
||||
from .dependency_node import DependencyNode
|
||||
import ctypes
|
||||
from typing import Optional, Any, Dict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_module_symbols(module_name: str):
|
||||
imported_module = importlib.import_module(module_name)
|
||||
return [name for name in dir(imported_module)], imported_module
|
||||
|
||||
|
||||
def process_vmlinux_class(node, llvm_module, handler: DependencyHandler):
|
||||
symbols_in_module, imported_module = get_module_symbols("vmlinux")
|
||||
if node.name in symbols_in_module:
|
||||
vmlinux_type = getattr(imported_module, node.name)
|
||||
process_vmlinux_post_ast(vmlinux_type, llvm_module, handler)
|
||||
else:
|
||||
raise ImportError(f"{node.name} not in vmlinux")
|
||||
|
||||
|
||||
def process_vmlinux_post_ast(
|
||||
elem_type_class, llvm_handler, handler: DependencyHandler, processing_stack=None
|
||||
):
|
||||
# Initialize processing stack on first call
|
||||
if processing_stack is None:
|
||||
processing_stack = set()
|
||||
symbols_in_module, imported_module = get_module_symbols("vmlinux")
|
||||
|
||||
current_symbol_name = elem_type_class.__name__
|
||||
logger.info(f"Begin {current_symbol_name} Processing")
|
||||
field_table: Dict[str, list] = {}
|
||||
is_complex_type = False
|
||||
containing_type: Optional[Any] = None
|
||||
ctype_complex_type: Optional[Any] = None
|
||||
type_length: Optional[int] = None
|
||||
module_name = getattr(elem_type_class, "__module__", None)
|
||||
|
||||
# Check if already processed
|
||||
if handler.has_node(current_symbol_name):
|
||||
logger.debug(f"Node {current_symbol_name} already processed and ready")
|
||||
return True
|
||||
|
||||
# XXX:Check it's use. It's probably not being used.
|
||||
if current_symbol_name in processing_stack:
|
||||
logger.debug(
|
||||
f"Dependency already in processing stack for {current_symbol_name}, skipping"
|
||||
)
|
||||
return True
|
||||
|
||||
processing_stack.add(current_symbol_name)
|
||||
|
||||
if module_name == "vmlinux":
|
||||
if hasattr(elem_type_class, "_type_"):
|
||||
pass
|
||||
else:
|
||||
new_dep_node = DependencyNode(name=current_symbol_name)
|
||||
|
||||
# elem_type_class is the actual vmlinux struct/class
|
||||
new_dep_node.set_ctype_struct(elem_type_class)
|
||||
|
||||
handler.add_node(new_dep_node)
|
||||
class_obj = getattr(imported_module, current_symbol_name)
|
||||
# Inspect the class fields
|
||||
if hasattr(class_obj, "_fields_"):
|
||||
for field_elem in class_obj._fields_:
|
||||
field_name: str = ""
|
||||
field_type: Optional[Any] = None
|
||||
bitfield_size: Optional[int] = None
|
||||
if len(field_elem) == 2:
|
||||
field_name, field_type = field_elem
|
||||
elif len(field_elem) == 3:
|
||||
field_name, field_type, bitfield_size = field_elem
|
||||
field_table[field_name] = [field_type, bitfield_size]
|
||||
elif hasattr(class_obj, "__annotations__"):
|
||||
for field_elem in class_obj.__annotations__.items():
|
||||
if len(field_elem) == 2:
|
||||
field_name, field_type = field_elem
|
||||
bitfield_size = None
|
||||
elif len(field_elem) == 3:
|
||||
field_name, field_type, bitfield_size = field_elem
|
||||
else:
|
||||
raise ValueError(
|
||||
"Number of fields in items() of class object unexpected"
|
||||
)
|
||||
field_table[field_name] = [field_type, bitfield_size]
|
||||
else:
|
||||
raise TypeError("Could not get required class and definition")
|
||||
|
||||
logger.debug(f"Extracted fields for {current_symbol_name}: {field_table}")
|
||||
for elem in field_table.items():
|
||||
elem_name, elem_temp_list = elem
|
||||
[elem_type, elem_bitfield_size] = elem_temp_list
|
||||
local_module_name = getattr(elem_type, "__module__", None)
|
||||
new_dep_node.add_field(elem_name, elem_type, ready=False)
|
||||
if local_module_name == ctypes.__name__:
|
||||
new_dep_node.set_field_bitfield_size(elem_name, elem_bitfield_size)
|
||||
new_dep_node.set_field_ready(elem_name, is_ready=True)
|
||||
logger.debug(
|
||||
f"Field {elem_name} is direct ctypes type: {elem_type}"
|
||||
)
|
||||
elif local_module_name == "vmlinux":
|
||||
new_dep_node.set_field_bitfield_size(elem_name, elem_bitfield_size)
|
||||
logger.debug(
|
||||
f"Processing vmlinux field: {elem_name}, type: {elem_type}"
|
||||
)
|
||||
if hasattr(elem_type, "_type_"):
|
||||
is_complex_type = True
|
||||
containing_type = elem_type._type_
|
||||
if hasattr(elem_type, "_length_") and is_complex_type:
|
||||
type_length = elem_type._length_
|
||||
|
||||
if containing_type.__module__ == "vmlinux":
|
||||
new_dep_node.add_dependent(
|
||||
elem_type._type_.__name__
|
||||
if hasattr(elem_type._type_, "__name__")
|
||||
else str(elem_type._type_)
|
||||
)
|
||||
elif containing_type.__module__ == ctypes.__name__:
|
||||
if isinstance(elem_type, type):
|
||||
if issubclass(elem_type, ctypes.Array):
|
||||
ctype_complex_type = ctypes.Array
|
||||
elif issubclass(elem_type, ctypes._Pointer):
|
||||
ctype_complex_type = ctypes._Pointer
|
||||
else:
|
||||
raise TypeError("Unsupported ctypes subclass")
|
||||
else:
|
||||
raise ImportError(
|
||||
f"Unsupported module of {containing_type}"
|
||||
)
|
||||
logger.debug(
|
||||
f"{containing_type} containing type of parent {elem_name} with {elem_type} and ctype {ctype_complex_type} and length {type_length}"
|
||||
)
|
||||
new_dep_node.set_field_containing_type(
|
||||
elem_name, containing_type
|
||||
)
|
||||
new_dep_node.set_field_type_size(elem_name, type_length)
|
||||
new_dep_node.set_field_ctype_complex_type(
|
||||
elem_name, ctype_complex_type
|
||||
)
|
||||
new_dep_node.set_field_type(elem_name, elem_type)
|
||||
if containing_type.__module__ == "vmlinux":
|
||||
containing_type_name = (
|
||||
containing_type.__name__
|
||||
if hasattr(containing_type, "__name__")
|
||||
else str(containing_type)
|
||||
)
|
||||
|
||||
# Check for self-reference or already processed
|
||||
if containing_type_name == current_symbol_name:
|
||||
# Self-referential pointer
|
||||
logger.debug(
|
||||
f"Self-referential pointer in {current_symbol_name}.{elem_name}"
|
||||
)
|
||||
new_dep_node.set_field_ready(elem_name, True)
|
||||
elif handler.has_node(containing_type_name):
|
||||
# Already processed
|
||||
logger.debug(
|
||||
f"Reusing already processed {containing_type_name}"
|
||||
)
|
||||
new_dep_node.set_field_ready(elem_name, True)
|
||||
else:
|
||||
# Process recursively - THIS WAS MISSING
|
||||
new_dep_node.add_dependent(containing_type_name)
|
||||
process_vmlinux_post_ast(
|
||||
containing_type,
|
||||
llvm_handler,
|
||||
handler,
|
||||
processing_stack,
|
||||
)
|
||||
new_dep_node.set_field_ready(elem_name, True)
|
||||
elif containing_type.__module__ == ctypes.__name__:
|
||||
logger.debug(f"Processing ctype internal{containing_type}")
|
||||
new_dep_node.set_field_ready(elem_name, True)
|
||||
else:
|
||||
raise TypeError(
|
||||
"Module not supported in recursive resolution"
|
||||
)
|
||||
else:
|
||||
new_dep_node.add_dependent(
|
||||
elem_type.__name__
|
||||
if hasattr(elem_type, "__name__")
|
||||
else str(elem_type)
|
||||
)
|
||||
process_vmlinux_post_ast(
|
||||
elem_type, llvm_handler, handler, processing_stack
|
||||
)
|
||||
new_dep_node.set_field_ready(elem_name, True)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"{elem_name} with type {elem_type} from module {module_name} not supported in recursive resolver"
|
||||
)
|
||||
|
||||
else:
|
||||
raise ImportError("UNSUPPORTED Module")
|
||||
|
||||
logging.info(
|
||||
f"{current_symbol_name} processed and handler readiness {handler.is_ready}"
|
||||
)
|
||||
return True
|
||||
173
pythonbpf/vmlinux_parser/dependency_handler.py
Normal file
173
pythonbpf/vmlinux_parser/dependency_handler.py
Normal file
@ -0,0 +1,173 @@
|
||||
from typing import Optional, Dict, List, Iterator
|
||||
from .dependency_node import DependencyNode
|
||||
|
||||
|
||||
class DependencyHandler:
|
||||
"""
|
||||
Manages a collection of DependencyNode objects with no duplicates.
|
||||
|
||||
Ensures that no two nodes with the same name can be added and provides
|
||||
methods to check readiness and retrieve specific nodes.
|
||||
|
||||
Example usage:
|
||||
# Create a handler
|
||||
handler = DependencyHandler()
|
||||
|
||||
# Create some dependency nodes
|
||||
node1 = DependencyNode(name="node1")
|
||||
node1.add_field("field1", str)
|
||||
node1.set_field_value("field1", "value1")
|
||||
|
||||
node2 = DependencyNode(name="node2")
|
||||
node2.add_field("field1", int)
|
||||
|
||||
# Add nodes to the handler
|
||||
handler.add_node(node1)
|
||||
handler.add_node(node2)
|
||||
|
||||
# Check if a specific node exists
|
||||
print(handler.has_node("node1")) # True
|
||||
|
||||
# Get a reference to a node and modify it
|
||||
node = handler.get_node("node2")
|
||||
node.set_field_value("field1", 42)
|
||||
|
||||
# Check if all nodes are ready
|
||||
print(handler.is_ready) # False (node2 is ready, but node1 isn't)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Using a dictionary with node names as keys ensures name uniqueness
|
||||
# and provides efficient lookups
|
||||
self._nodes: Dict[str, DependencyNode] = {}
|
||||
|
||||
def add_node(self, node: DependencyNode) -> bool:
|
||||
"""
|
||||
Add a dependency node to the handler.
|
||||
|
||||
Args:
|
||||
node: The DependencyNode to add
|
||||
|
||||
Returns:
|
||||
bool: True if the node was added, False if a node with the same name already exists
|
||||
|
||||
Raises:
|
||||
TypeError: If the provided object is not a DependencyNode
|
||||
"""
|
||||
if not isinstance(node, DependencyNode):
|
||||
raise TypeError(f"Expected DependencyNode, got {type(node).__name__}")
|
||||
|
||||
# Check if a node with this name already exists
|
||||
if node.name in self._nodes:
|
||||
return False
|
||||
|
||||
self._nodes[node.name] = node
|
||||
return True
|
||||
|
||||
@property
|
||||
def is_ready(self) -> bool:
|
||||
"""
|
||||
Check if all nodes are ready.
|
||||
|
||||
Returns:
|
||||
bool: True if all nodes are ready (or if there are no nodes), False otherwise
|
||||
"""
|
||||
if not self._nodes:
|
||||
return True
|
||||
|
||||
return all(node.is_ready for node in self._nodes.values())
|
||||
|
||||
def has_node(self, name: str) -> bool:
|
||||
"""
|
||||
Check if a node with the given name exists.
|
||||
|
||||
Args:
|
||||
name: The name to check
|
||||
|
||||
Returns:
|
||||
bool: True if a node with the given name exists, False otherwise
|
||||
"""
|
||||
return name in self._nodes
|
||||
|
||||
def get_node(self, name: str) -> Optional[DependencyNode]:
|
||||
"""
|
||||
Get a node by name for manipulation.
|
||||
|
||||
Args:
|
||||
name: The name of the node to retrieve
|
||||
|
||||
Returns:
|
||||
Optional[DependencyNode]: The node with the given name, or None if not found
|
||||
"""
|
||||
return self._nodes.get(name)
|
||||
|
||||
def remove_node(self, node_or_name) -> bool:
|
||||
"""
|
||||
Remove a node by name or reference.
|
||||
|
||||
Args:
|
||||
node_or_name: The node to remove or its name
|
||||
|
||||
Returns:
|
||||
bool: True if the node was removed, False if not found
|
||||
"""
|
||||
if isinstance(node_or_name, DependencyNode):
|
||||
name = node_or_name.name
|
||||
else:
|
||||
name = node_or_name
|
||||
|
||||
if name in self._nodes:
|
||||
del self._nodes[name]
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_all_nodes(self) -> List[DependencyNode]:
|
||||
"""
|
||||
Get all nodes stored in the handler.
|
||||
|
||||
Returns:
|
||||
List[DependencyNode]: List of all nodes
|
||||
"""
|
||||
return list(self._nodes.values())
|
||||
|
||||
def __iter__(self) -> Iterator[DependencyNode]:
|
||||
"""
|
||||
Iterate over all nodes.
|
||||
|
||||
Returns:
|
||||
Iterator[DependencyNode]: Iterator over all nodes
|
||||
"""
|
||||
return iter(self._nodes.values())
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""
|
||||
Get the number of nodes in the handler.
|
||||
|
||||
Returns:
|
||||
int: The number of nodes
|
||||
"""
|
||||
return len(self._nodes)
|
||||
|
||||
def __getitem__(self, name: str) -> DependencyNode:
|
||||
"""
|
||||
Get a node by name using dictionary-style access.
|
||||
|
||||
Args:
|
||||
name: The name of the node to retrieve
|
||||
|
||||
Returns:
|
||||
DependencyNode: The node with the given name
|
||||
|
||||
Raises:
|
||||
KeyError: If no node with the given name exists
|
||||
|
||||
Example:
|
||||
node = handler["some-dep_node_name"]
|
||||
"""
|
||||
if name not in self._nodes:
|
||||
raise KeyError(f"No node with name '{name}' found")
|
||||
return self._nodes[name]
|
||||
|
||||
@property
|
||||
def nodes(self):
|
||||
return self._nodes
|
||||
363
pythonbpf/vmlinux_parser/dependency_node.py
Normal file
363
pythonbpf/vmlinux_parser/dependency_node.py
Normal file
@ -0,0 +1,363 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, Any, Optional
|
||||
import ctypes
|
||||
|
||||
|
||||
# TODO: FIX THE FUCKING TYPE NAME CONVENTION.
|
||||
@dataclass
|
||||
class Field:
|
||||
"""Represents a field in a dependency node with its type and readiness state."""
|
||||
|
||||
name: str
|
||||
type: type
|
||||
ctype_complex_type: Optional[Any]
|
||||
containing_type: Optional[Any]
|
||||
type_size: Optional[int]
|
||||
bitfield_size: Optional[int]
|
||||
offset: int
|
||||
value: Any = None
|
||||
ready: bool = False
|
||||
|
||||
def set_ready(self, is_ready: bool = True) -> None:
|
||||
"""Set the readiness state of this field."""
|
||||
self.ready = is_ready
|
||||
|
||||
def set_value(self, value: Any, mark_ready: bool = False) -> None:
|
||||
"""Set the value of this field and optionally mark it as ready."""
|
||||
self.value = value
|
||||
if mark_ready:
|
||||
self.ready = True
|
||||
|
||||
def set_type(self, given_type, mark_ready: bool = False) -> None:
|
||||
"""Set value of the type field and mark as ready"""
|
||||
self.type = given_type
|
||||
if mark_ready:
|
||||
self.ready = True
|
||||
|
||||
def set_containing_type(
|
||||
self, containing_type: Optional[Any], mark_ready: bool = False
|
||||
) -> None:
|
||||
"""Set the containing_type of this field and optionally mark it as ready."""
|
||||
self.containing_type = containing_type
|
||||
if mark_ready:
|
||||
self.ready = True
|
||||
|
||||
def set_type_size(self, type_size: Any, mark_ready: bool = False) -> None:
|
||||
"""Set the type_size of this field and optionally mark it as ready."""
|
||||
self.type_size = type_size
|
||||
if mark_ready:
|
||||
self.ready = True
|
||||
|
||||
def set_ctype_complex_type(
|
||||
self, ctype_complex_type: Any, mark_ready: bool = False
|
||||
) -> None:
|
||||
"""Set the ctype_complex_type of this field and optionally mark it as ready."""
|
||||
self.ctype_complex_type = ctype_complex_type
|
||||
if mark_ready:
|
||||
self.ready = True
|
||||
|
||||
def set_bitfield_size(self, bitfield_size: Any, mark_ready: bool = False) -> None:
|
||||
"""Set the bitfield_size of this field and optionally mark it as ready."""
|
||||
self.bitfield_size = bitfield_size
|
||||
if mark_ready:
|
||||
self.ready = True
|
||||
|
||||
def set_offset(self, offset: int) -> None:
|
||||
"""Set the offset of this field"""
|
||||
self.offset = offset
|
||||
|
||||
|
||||
@dataclass
|
||||
class DependencyNode:
|
||||
"""
|
||||
A node with typed fields and readiness tracking.
|
||||
|
||||
Example usage:
|
||||
# Create a dependency node for a Person
|
||||
somestruct = DependencyNode(name="struct_1")
|
||||
|
||||
# Add fields with their types
|
||||
somestruct.add_field("field_1", str)
|
||||
somestruct.add_field("field_2", int)
|
||||
somestruct.add_field("field_3", str)
|
||||
|
||||
# Check if the node is ready (should be False initially)
|
||||
print(f"Is node ready? {somestruct.is_ready}") # False
|
||||
|
||||
# Set some field values
|
||||
somestruct.set_field_value("field_1", "someproperty")
|
||||
somestruct.set_field_value("field_2", 30)
|
||||
|
||||
# Check if the node is ready (still False because email is not ready)
|
||||
print(f"Is node ready? {somestruct.is_ready}") # False
|
||||
|
||||
# Set the last field and make the node ready
|
||||
somestruct.set_field_value("field_3", "anotherproperty")
|
||||
|
||||
# Now the node should be ready
|
||||
print(f"Is node ready? {somestruct.is_ready}") # True
|
||||
|
||||
# You can also mark a field as not ready
|
||||
somestruct.set_field_ready("field_3", False)
|
||||
|
||||
# Now the node is not ready again
|
||||
print(f"Is node ready? {somestruct.is_ready}") # False
|
||||
|
||||
# Get all field values
|
||||
print(somestruct.get_field_values()) # {'field_1': 'someproperty', 'field_2': 30, 'field_3': 'anotherproperty'}
|
||||
|
||||
# Get only ready fields
|
||||
ready_fields = somestruct.get_ready_fields()
|
||||
print(f"Ready fields: {[field.name for field in ready_fields.values()]}") # ['field_1', 'field_2']
|
||||
"""
|
||||
|
||||
name: str
|
||||
depends_on: Optional[list[str]] = None
|
||||
fields: Dict[str, Field] = field(default_factory=dict)
|
||||
_ready_cache: Optional[bool] = field(default=None, repr=False)
|
||||
current_offset: int = 0
|
||||
ctype_struct: Optional[Any] = field(default=None, repr=False)
|
||||
|
||||
def add_field(
|
||||
self,
|
||||
name: str,
|
||||
field_type: type,
|
||||
initial_value: Any = None,
|
||||
containing_type: Optional[Any] = None,
|
||||
type_size: Optional[int] = None,
|
||||
ctype_complex_type: Optional[int] = None,
|
||||
bitfield_size: Optional[int] = None,
|
||||
ready: bool = False,
|
||||
offset: int = 0,
|
||||
) -> None:
|
||||
"""Add a field to the node with an optional initial value and readiness state."""
|
||||
if self.depends_on is None:
|
||||
self.depends_on = []
|
||||
self.fields[name] = Field(
|
||||
name=name,
|
||||
type=field_type,
|
||||
value=initial_value,
|
||||
ready=ready,
|
||||
containing_type=containing_type,
|
||||
type_size=type_size,
|
||||
ctype_complex_type=ctype_complex_type,
|
||||
bitfield_size=bitfield_size,
|
||||
offset=offset,
|
||||
)
|
||||
# Invalidate readiness cache
|
||||
self._ready_cache = None
|
||||
|
||||
def set_ctype_struct(self, ctype_struct: Any) -> None:
|
||||
"""Set the ctypes structure for automatic offset calculation."""
|
||||
self.ctype_struct = ctype_struct
|
||||
|
||||
def __sizeof__(self):
|
||||
# If we have a ctype_struct, use its size
|
||||
if self.ctype_struct is not None:
|
||||
return ctypes.sizeof(self.ctype_struct)
|
||||
return self.current_offset
|
||||
|
||||
def get_field(self, name: str) -> Field:
|
||||
"""Get a field by name."""
|
||||
return self.fields[name]
|
||||
|
||||
def set_field_value(self, name: str, value: Any, mark_ready: bool = False) -> None:
|
||||
"""Set a field's value and optionally mark it as ready."""
|
||||
if name not in self.fields:
|
||||
raise KeyError(f"Field '{name}' does not exist in node '{self.name}'")
|
||||
|
||||
self.fields[name].set_value(value, mark_ready)
|
||||
# Invalidate readiness cache
|
||||
self._ready_cache = None
|
||||
|
||||
def set_field_type(self, name: str, type: Any, mark_ready: bool = False) -> None:
|
||||
"""Set a field's type and optionally mark it as ready."""
|
||||
if name not in self.fields:
|
||||
raise KeyError(f"Field '{name}' does not exist in node '{self.name}'")
|
||||
|
||||
self.fields[name].set_type(type, mark_ready)
|
||||
# Invalidate readiness cache
|
||||
self._ready_cache = None
|
||||
|
||||
def set_field_containing_type(
|
||||
self, name: str, containing_type: Any, mark_ready: bool = False
|
||||
) -> None:
|
||||
"""Set a field's containing_type and optionally mark it as ready."""
|
||||
if name not in self.fields:
|
||||
raise KeyError(f"Field '{name}' does not exist in node '{self.name}'")
|
||||
|
||||
self.fields[name].set_containing_type(containing_type, mark_ready)
|
||||
# Invalidate readiness cache
|
||||
self._ready_cache = None
|
||||
|
||||
def set_field_type_size(
|
||||
self, name: str, type_size: Any, mark_ready: bool = False
|
||||
) -> None:
|
||||
"""Set a field's type_size and optionally mark it as ready."""
|
||||
if name not in self.fields:
|
||||
raise KeyError(f"Field '{name}' does not exist in node '{self.name}'")
|
||||
|
||||
self.fields[name].set_type_size(type_size, mark_ready)
|
||||
# Invalidate readiness cache
|
||||
self._ready_cache = None
|
||||
|
||||
def set_field_ctype_complex_type(
|
||||
self, name: str, ctype_complex_type: Any, mark_ready: bool = False
|
||||
) -> None:
|
||||
"""Set a field's ctype_complex_type and optionally mark it as ready."""
|
||||
if name not in self.fields:
|
||||
raise KeyError(f"Field '{name}' does not exist in node '{self.name}'")
|
||||
|
||||
self.fields[name].set_ctype_complex_type(ctype_complex_type, mark_ready)
|
||||
# Invalidate readiness cache
|
||||
self._ready_cache = None
|
||||
|
||||
def set_field_bitfield_size(
|
||||
self, name: str, bitfield_size: Any, mark_ready: bool = False
|
||||
) -> None:
|
||||
"""Set a field's bitfield_size and optionally mark it as ready."""
|
||||
if name not in self.fields:
|
||||
raise KeyError(f"Field '{name}' does not exist in node '{self.name}'")
|
||||
|
||||
self.fields[name].set_bitfield_size(bitfield_size, mark_ready)
|
||||
# Invalidate readiness cache
|
||||
self._ready_cache = None
|
||||
|
||||
def set_field_ready(
|
||||
self,
|
||||
name: str,
|
||||
is_ready: bool = False,
|
||||
size_of_containing_type: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Mark a field as ready or not ready."""
|
||||
if name not in self.fields:
|
||||
raise KeyError(f"Field '{name}' does not exist in node '{self.name}'")
|
||||
|
||||
self.fields[name].set_ready(is_ready)
|
||||
|
||||
# Use ctypes built-in offset if available
|
||||
if self.ctype_struct is not None:
|
||||
try:
|
||||
self.fields[name].set_offset(getattr(self.ctype_struct, name).offset)
|
||||
except AttributeError:
|
||||
# Fallback to manual calculation if field not found in ctype_struct
|
||||
self.fields[name].set_offset(self.current_offset)
|
||||
self.current_offset += self._calculate_size(
|
||||
name, size_of_containing_type
|
||||
)
|
||||
else:
|
||||
# Manual offset calculation when no ctype_struct is available
|
||||
self.fields[name].set_offset(self.current_offset)
|
||||
self.current_offset += self._calculate_size(name, size_of_containing_type)
|
||||
|
||||
# Invalidate readiness cache
|
||||
self._ready_cache = None
|
||||
|
||||
def _calculate_size(
|
||||
self, name: str, size_of_containing_type: Optional[int] = None
|
||||
) -> int:
|
||||
processing_field = self.fields[name]
|
||||
# size_of_field will be in bytes
|
||||
if processing_field.type.__module__ == ctypes.__name__:
|
||||
size_of_field = ctypes.sizeof(processing_field.type)
|
||||
return size_of_field
|
||||
elif processing_field.type.__module__ == "vmlinux":
|
||||
if processing_field.ctype_complex_type is not None:
|
||||
if issubclass(processing_field.ctype_complex_type, ctypes.Array):
|
||||
if processing_field.containing_type.__module__ == ctypes.__name__:
|
||||
if (
|
||||
processing_field.containing_type is not None
|
||||
and processing_field.type_size is not None
|
||||
):
|
||||
size_of_field = (
|
||||
ctypes.sizeof(processing_field.containing_type)
|
||||
* processing_field.type_size
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"{processing_field} has no containing_type or type_size"
|
||||
)
|
||||
return size_of_field
|
||||
elif processing_field.containing_type.__module__ == "vmlinux":
|
||||
if (
|
||||
size_of_containing_type is not None
|
||||
and processing_field.type_size is not None
|
||||
):
|
||||
size_of_field = (
|
||||
size_of_containing_type * processing_field.type_size
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"{processing_field} has no containing_type or type_size"
|
||||
)
|
||||
return size_of_field
|
||||
elif issubclass(processing_field.ctype_complex_type, ctypes._Pointer):
|
||||
return ctypes.sizeof(ctypes.c_void_p)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"This subclass of ctype not supported yet"
|
||||
)
|
||||
elif processing_field.type_size is not None:
|
||||
# Handle vmlinux types with type_size but no ctype_complex_type
|
||||
# This means it's a direct vmlinux struct field (not array/pointer wrapped)
|
||||
# The type_size should already contain the full size of the struct
|
||||
# But if there's a containing_type from vmlinux, we need that size
|
||||
if processing_field.containing_type is not None:
|
||||
if processing_field.containing_type.__module__ == "vmlinux":
|
||||
# For vmlinux containing types, we need the pre-calculated size
|
||||
if size_of_containing_type is not None:
|
||||
return size_of_containing_type * processing_field.type_size
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Field {name}: vmlinux containing_type requires size_of_containing_type"
|
||||
)
|
||||
else:
|
||||
raise ModuleNotFoundError(
|
||||
f"Containing type module {processing_field.containing_type.__module__} not supported"
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("Wrong type found with no containing type")
|
||||
else:
|
||||
# No ctype_complex_type and no type_size, must rely on size_of_containing_type
|
||||
if size_of_containing_type is None:
|
||||
raise RuntimeError(
|
||||
f"Size of containing type {size_of_containing_type} is None"
|
||||
)
|
||||
return size_of_containing_type
|
||||
|
||||
else:
|
||||
raise ModuleNotFoundError("Module is not supported for the operation")
|
||||
raise RuntimeError("control should not reach here")
|
||||
|
||||
@property
|
||||
def is_ready(self) -> bool:
|
||||
"""Check if the node is ready (all fields are ready)."""
|
||||
# Use cached value if available
|
||||
if self._ready_cache is not None:
|
||||
return self._ready_cache
|
||||
|
||||
# Calculate readiness only when needed
|
||||
if not self.fields:
|
||||
self._ready_cache = True
|
||||
return True
|
||||
|
||||
self._ready_cache = all(elem.ready for elem in self.fields.values())
|
||||
return self._ready_cache
|
||||
|
||||
def get_field_values(self) -> Dict[str, Any]:
|
||||
"""Get a dictionary of field names to their values."""
|
||||
return {name: elem.value for name, elem in self.fields.items()}
|
||||
|
||||
def get_ready_fields(self) -> Dict[str, Field]:
|
||||
"""Get all fields that are marked as ready."""
|
||||
return {name: elem for name, elem in self.fields.items() if elem.ready}
|
||||
|
||||
def get_not_ready_fields(self) -> Dict[str, Field]:
|
||||
"""Get all fields that are marked as not ready."""
|
||||
return {name: elem for name, elem in self.fields.items() if not elem.ready}
|
||||
|
||||
def add_dependent(self, dep_type):
|
||||
if dep_type in self.depends_on:
|
||||
return
|
||||
else:
|
||||
self.depends_on.append(dep_type)
|
||||
147
pythonbpf/vmlinux_parser/import_detector.py
Normal file
147
pythonbpf/vmlinux_parser/import_detector.py
Normal file
@ -0,0 +1,147 @@
|
||||
import ast
|
||||
import logging
|
||||
from typing import List, Tuple, Any
|
||||
import importlib
|
||||
import inspect
|
||||
|
||||
from .dependency_handler import DependencyHandler
|
||||
from .ir_gen import IRGenerator
|
||||
from .class_handler import process_vmlinux_class
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def detect_import_statement(tree: ast.AST) -> List[Tuple[str, ast.ImportFrom]]:
|
||||
"""
|
||||
Parse AST and detect import statements from vmlinux.
|
||||
|
||||
Returns a list of tuples (module_name, imported_item) for vmlinux imports.
|
||||
Raises SyntaxError for invalid import patterns.
|
||||
|
||||
Args:
|
||||
tree: The AST to parse
|
||||
|
||||
Returns:
|
||||
List of tuples containing (module_name, imported_item) for each vmlinux import
|
||||
|
||||
Raises:
|
||||
SyntaxError: If multiple imports from vmlinux are attempted or import * is used
|
||||
"""
|
||||
vmlinux_imports = []
|
||||
|
||||
for node in ast.walk(tree):
|
||||
# Handle "from vmlinux import ..." statements
|
||||
if isinstance(node, ast.ImportFrom):
|
||||
if node.module == "vmlinux":
|
||||
# Check for wildcard import: from vmlinux import *
|
||||
if any(alias.name == "*" for alias in node.names):
|
||||
raise SyntaxError(
|
||||
"Wildcard imports from vmlinux are not supported. "
|
||||
"Please import specific types explicitly."
|
||||
)
|
||||
|
||||
# Check for multiple imports: from vmlinux import A, B, C
|
||||
if len(node.names) > 1:
|
||||
imported_names = [alias.name for alias in node.names]
|
||||
raise SyntaxError(
|
||||
f"Multiple imports from vmlinux are not supported. "
|
||||
f"Found: {', '.join(imported_names)}. "
|
||||
f"Please use separate import statements for each type."
|
||||
)
|
||||
|
||||
# Check if no specific import is specified (should not happen with valid Python)
|
||||
if len(node.names) == 0:
|
||||
raise SyntaxError(
|
||||
"Import from vmlinux must specify at least one type."
|
||||
)
|
||||
|
||||
# Valid single import
|
||||
for alias in node.names:
|
||||
import_name = alias.name
|
||||
# Use alias if provided, otherwise use the original name (commented)
|
||||
# as_name = alias.asname if alias.asname else alias.name
|
||||
vmlinux_imports.append(("vmlinux", node))
|
||||
logger.info(f"Found vmlinux import: {import_name}")
|
||||
|
||||
# Handle "import vmlinux" statements (not typical but should be rejected)
|
||||
elif isinstance(node, ast.Import):
|
||||
for alias in node.names:
|
||||
if alias.name == "vmlinux" or alias.name.startswith("vmlinux."):
|
||||
raise SyntaxError(
|
||||
"Direct import of vmlinux module is not supported. "
|
||||
"Use 'from vmlinux import <type>' instead."
|
||||
)
|
||||
|
||||
logger.info(f"Total vmlinux imports detected: {len(vmlinux_imports)}")
|
||||
return vmlinux_imports
|
||||
|
||||
|
||||
def vmlinux_proc(tree: ast.AST, module):
|
||||
import_statements = detect_import_statement(tree)
|
||||
|
||||
# initialise dependency handler
|
||||
handler = DependencyHandler()
|
||||
# initialise assignment dictionary of name to type
|
||||
assignments: dict[str, tuple[type, Any]] = {}
|
||||
|
||||
if not import_statements:
|
||||
logger.info("No vmlinux imports found")
|
||||
return
|
||||
|
||||
# Import vmlinux module directly
|
||||
try:
|
||||
vmlinux_mod = importlib.import_module("vmlinux")
|
||||
except ImportError:
|
||||
logger.warning("Could not import vmlinux module")
|
||||
return
|
||||
|
||||
source_file = inspect.getsourcefile(vmlinux_mod)
|
||||
if source_file is None:
|
||||
logger.warning("Cannot find source for vmlinux module")
|
||||
return
|
||||
|
||||
with open(source_file, "r") as f:
|
||||
mod_ast = ast.parse(f.read(), filename=source_file)
|
||||
|
||||
for import_mod, import_node in import_statements:
|
||||
for alias in import_node.names:
|
||||
imported_name = alias.name
|
||||
found = False
|
||||
for mod_node in mod_ast.body:
|
||||
if (
|
||||
isinstance(mod_node, ast.ClassDef)
|
||||
and mod_node.name == imported_name
|
||||
):
|
||||
process_vmlinux_class(mod_node, module, handler)
|
||||
found = True
|
||||
break
|
||||
if isinstance(mod_node, ast.Assign):
|
||||
for target in mod_node.targets:
|
||||
if isinstance(target, ast.Name) and target.id == imported_name:
|
||||
process_vmlinux_assign(mod_node, module, assignments)
|
||||
found = True
|
||||
break
|
||||
if found:
|
||||
break
|
||||
if not found:
|
||||
logger.info(
|
||||
f"{imported_name} not found as ClassDef or Assign in vmlinux"
|
||||
)
|
||||
|
||||
IRGenerator(module, handler)
|
||||
return assignments
|
||||
|
||||
|
||||
def process_vmlinux_assign(node, module, assignments: dict[str, tuple[type, Any]]):
|
||||
# Check if this is a simple assignment with a constant value
|
||||
if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
|
||||
target_name = node.targets[0].id
|
||||
if isinstance(node.value, ast.Constant):
|
||||
assignments[target_name] = (type(node.value.value), node.value.value)
|
||||
logger.info(
|
||||
f"Added assignment: {target_name} = {node.value.value!r} of type {type(node.value.value)}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported assignment type for {target_name}")
|
||||
else:
|
||||
raise ValueError("Not a simple assignment")
|
||||
3
pythonbpf/vmlinux_parser/ir_gen/__init__.py
Normal file
3
pythonbpf/vmlinux_parser/ir_gen/__init__.py
Normal file
@ -0,0 +1,3 @@
|
||||
from .ir_generation import IRGenerator
|
||||
|
||||
__all__ = ["IRGenerator"]
|
||||
15
pythonbpf/vmlinux_parser/ir_gen/debug_info_gen.py
Normal file
15
pythonbpf/vmlinux_parser/ir_gen/debug_info_gen.py
Normal file
@ -0,0 +1,15 @@
|
||||
from pythonbpf.debuginfo import DebugInfoGenerator
|
||||
|
||||
|
||||
def debug_info_generation(struct, llvm_module):
|
||||
generator = DebugInfoGenerator(llvm_module)
|
||||
# this is sample debug info generation
|
||||
# i64type = generator.get_uint64_type()
|
||||
|
||||
struct_type = generator.create_struct_type([], 64 * 4, is_distinct=True)
|
||||
|
||||
global_var = generator.create_global_var_debug_info(
|
||||
struct.name, struct_type, is_local=False
|
||||
)
|
||||
|
||||
return global_var
|
||||
161
pythonbpf/vmlinux_parser/ir_gen/ir_generation.py
Normal file
161
pythonbpf/vmlinux_parser/ir_gen/ir_generation.py
Normal file
@ -0,0 +1,161 @@
|
||||
import ctypes
|
||||
import logging
|
||||
from ..dependency_handler import DependencyHandler
|
||||
from .debug_info_gen import debug_info_generation
|
||||
from ..dependency_node import DependencyNode
|
||||
import llvmlite.ir as ir
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IRGenerator:
|
||||
# get the assignments dict and add this stuff to it.
|
||||
def __init__(self, llvm_module, handler: DependencyHandler, assignment=None):
|
||||
self.llvm_module = llvm_module
|
||||
self.handler: DependencyHandler = handler
|
||||
self.generated: list[str] = []
|
||||
if not handler.is_ready:
|
||||
raise ImportError(
|
||||
"Semantic analysis of vmlinux imports failed. Cannot generate IR"
|
||||
)
|
||||
for struct in handler:
|
||||
self.struct_processor(struct)
|
||||
|
||||
def struct_processor(self, struct, processing_stack=None):
|
||||
# Initialize processing stack on first call
|
||||
if processing_stack is None:
|
||||
processing_stack = set()
|
||||
|
||||
# If already generated, skip
|
||||
if struct.name in self.generated:
|
||||
return
|
||||
|
||||
# Detect circular dependency
|
||||
if struct.name in processing_stack:
|
||||
logger.info(
|
||||
f"Circular dependency detected for {struct.name}, skipping recursive processing"
|
||||
)
|
||||
# For circular dependencies, we can either:
|
||||
# 1. Use forward declarations (opaque pointers)
|
||||
# 2. Mark as incomplete and process later
|
||||
# 3. Generate a placeholder type
|
||||
# Here we'll just skip and let it be processed in its own call
|
||||
return
|
||||
|
||||
logger.info(f"IR generating for {struct.name}")
|
||||
|
||||
# Add to processing stack before processing dependencies
|
||||
processing_stack.add(struct.name)
|
||||
|
||||
try:
|
||||
# Process all dependencies first
|
||||
if struct.depends_on is None:
|
||||
pass
|
||||
else:
|
||||
for dependency in struct.depends_on:
|
||||
if dependency not in self.generated:
|
||||
# Check if dependency exists in handler
|
||||
if dependency in self.handler.nodes:
|
||||
dep_node_from_dependency = self.handler[dependency]
|
||||
# Pass the processing_stack down to track circular refs
|
||||
self.struct_processor(
|
||||
dep_node_from_dependency, processing_stack
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Warning: Dependency {dependency} not found in handler"
|
||||
)
|
||||
|
||||
# Actual processor logic here after dependencies are resolved
|
||||
self.gen_ir(struct)
|
||||
self.generated.append(struct.name)
|
||||
|
||||
finally:
|
||||
# Remove from processing stack after we're done
|
||||
processing_stack.discard(struct.name)
|
||||
|
||||
def gen_ir(self, struct):
|
||||
# TODO: we add the btf_ama attribute by monkey patching in the end of compilation, but once llvmlite
|
||||
# accepts our issue, we will resort to normal accessed attribute based attribute addition
|
||||
# currently we generate all possible field accesses for CO-RE and put into the assignment table
|
||||
debug_info = debug_info_generation(struct, self.llvm_module)
|
||||
field_index = 0
|
||||
for field_name, field in struct.fields.items():
|
||||
# does not take arrays and similar types into consideration yet.
|
||||
if field.ctype_complex_type is not None and issubclass(
|
||||
field.ctype_complex_type, ctypes.Array
|
||||
):
|
||||
array_size = field.type_size
|
||||
containing_type = field.containing_type
|
||||
if containing_type.__module__ == ctypes.__name__:
|
||||
containing_type_size = ctypes.sizeof(containing_type)
|
||||
for i in range(0, array_size):
|
||||
field_co_re_name = self._struct_name_generator(
|
||||
struct, field, field_index, True, i, containing_type_size
|
||||
)
|
||||
globvar = ir.GlobalVariable(
|
||||
self.llvm_module, ir.IntType(64), name=field_co_re_name
|
||||
)
|
||||
globvar.linkage = "external"
|
||||
globvar.set_metadata("llvm.preserve.access.index", debug_info)
|
||||
field_index += 1
|
||||
elif field.type_size is not None:
|
||||
array_size = field.type_size
|
||||
containing_type = field.containing_type
|
||||
if containing_type.__module__ == "vmlinux":
|
||||
containing_type_size = self.handler[
|
||||
containing_type.__name__
|
||||
].current_offset
|
||||
for i in range(0, array_size):
|
||||
field_co_re_name = self._struct_name_generator(
|
||||
struct, field, field_index, True, i, containing_type_size
|
||||
)
|
||||
globvar = ir.GlobalVariable(
|
||||
self.llvm_module, ir.IntType(64), name=field_co_re_name
|
||||
)
|
||||
globvar.linkage = "external"
|
||||
globvar.set_metadata("llvm.preserve.access.index", debug_info)
|
||||
field_index += 1
|
||||
else:
|
||||
field_co_re_name = self._struct_name_generator(
|
||||
struct, field, field_index
|
||||
)
|
||||
field_index += 1
|
||||
globvar = ir.GlobalVariable(
|
||||
self.llvm_module, ir.IntType(64), name=field_co_re_name
|
||||
)
|
||||
globvar.linkage = "external"
|
||||
globvar.set_metadata("llvm.preserve.access.index", debug_info)
|
||||
|
||||
def _struct_name_generator(
|
||||
self,
|
||||
struct: DependencyNode,
|
||||
field,
|
||||
field_index: int,
|
||||
is_indexed: bool = False,
|
||||
index: int = 0,
|
||||
containing_type_size: int = 0,
|
||||
) -> str:
|
||||
if is_indexed:
|
||||
name = (
|
||||
"llvm."
|
||||
+ struct.name.removeprefix("struct_")
|
||||
+ f":0:{field.offset + index * containing_type_size}"
|
||||
+ "$"
|
||||
+ f"0:{field_index}:{index}"
|
||||
)
|
||||
return name
|
||||
elif struct.name.startswith("struct_"):
|
||||
name = (
|
||||
"llvm."
|
||||
+ struct.name.removeprefix("struct_")
|
||||
+ f":0:{field.offset}"
|
||||
+ "$"
|
||||
+ f"0:{field_index}"
|
||||
)
|
||||
return name
|
||||
else:
|
||||
print(self.handler[struct.name])
|
||||
raise TypeError(
|
||||
"Name generation cannot occur due to type name not starting with struct"
|
||||
)
|
||||
@ -1,11 +1,10 @@
|
||||
#include <linux/bpf.h>
|
||||
#include "vmlinux.h"
|
||||
#include <bpf/bpf_helpers.h>
|
||||
#define u64 unsigned long long
|
||||
#define u32 unsigned int
|
||||
#include <bpf/bpf_endian.h>
|
||||
|
||||
SEC("xdp")
|
||||
int hello(struct xdp_md *ctx) {
|
||||
bpf_printk("Hello, World!\n");
|
||||
bpf_printk("Hello, World! %ud \n", ctx->data);
|
||||
return XDP_PASS;
|
||||
}
|
||||
|
||||
|
||||
@ -1,23 +1,9 @@
|
||||
// SPDX-License-Identifier: GPL-2.0
|
||||
|
||||
#include <linux/bpf.h>
|
||||
#include "vmlinux.h"
|
||||
#include <bpf/bpf_helpers.h>
|
||||
#include <bpf/bpf_tracing.h>
|
||||
|
||||
struct trace_entry {
|
||||
short unsigned int type;
|
||||
unsigned char flags;
|
||||
unsigned char preempt_count;
|
||||
int pid;
|
||||
};
|
||||
|
||||
struct trace_event_raw_sys_enter {
|
||||
struct trace_entry ent;
|
||||
long int id;
|
||||
long unsigned int args[6];
|
||||
char __data[0];
|
||||
};
|
||||
|
||||
struct event {
|
||||
__u32 pid;
|
||||
__u32 uid;
|
||||
|
||||
@ -11,6 +11,7 @@ from pythonbpf.maps import HashMap
|
||||
# We can allow bitcasts in cases where the width of the types is the same in
|
||||
# the future. But for now, we do not allow any re-interpretation of variables.
|
||||
|
||||
|
||||
@bpf
|
||||
@map
|
||||
def last() -> HashMap:
|
||||
|
||||
@ -3,16 +3,19 @@ import logging
|
||||
from pythonbpf import compile, bpf, section, bpfglobal, compile_to_ir
|
||||
from ctypes import c_void_p, c_int64, c_int32
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def somevalue() -> c_int32:
|
||||
return c_int32(42)
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def somevalue2() -> c_int64:
|
||||
return c_int64(69)
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def somevalue1() -> c_int32:
|
||||
@ -21,12 +24,14 @@ def somevalue1() -> c_int32:
|
||||
|
||||
# --- Passing examples ---
|
||||
|
||||
|
||||
# Simple constant return
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def g1() -> c_int64:
|
||||
return c_int64(42)
|
||||
|
||||
|
||||
# Constructor with one constant argument
|
||||
@bpf
|
||||
@bpfglobal
|
||||
@ -62,15 +67,17 @@ def g2() -> c_int64:
|
||||
# def g6() -> c_int64:
|
||||
# return c_int64(CONST)
|
||||
|
||||
|
||||
# Constructor with multiple args
|
||||
#TODO: this is not working. should it work ?
|
||||
# TODO: this is not working. should it work ?
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def g7() -> c_int64:
|
||||
return c_int64(1)
|
||||
|
||||
|
||||
# Dataclass call
|
||||
#TODO: fails with dataclass
|
||||
# TODO: fails with dataclass
|
||||
# @dataclass
|
||||
# class Point:
|
||||
# x: c_int64
|
||||
@ -91,6 +98,7 @@ def sometag(ctx: c_void_p) -> c_int64:
|
||||
print(f"{somevalue}")
|
||||
return c_int64(1)
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def LICENSE() -> str:
|
||||
|
||||
@ -11,6 +11,7 @@ from ctypes import c_void_p, c_int64
|
||||
# We cannot allocate space for the intermediate type now.
|
||||
# We probably need to track the ref/deref chain for each variable.
|
||||
|
||||
|
||||
@bpf
|
||||
@map
|
||||
def count() -> HashMap:
|
||||
|
||||
@ -3,6 +3,7 @@ import logging
|
||||
from pythonbpf import compile, bpf, section, bpfglobal, compile_to_ir
|
||||
from ctypes import c_void_p, c_int64
|
||||
|
||||
|
||||
# This should not pass as somevalue is not declared at all.
|
||||
@bpf
|
||||
@section("tracepoint/syscalls/sys_enter_execve")
|
||||
@ -11,6 +12,7 @@ def sometag(ctx: c_void_p) -> c_int64:
|
||||
print(f"{somevalue}") # noqa: F821
|
||||
return c_int64(1)
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def LICENSE() -> str:
|
||||
|
||||
54
tests/failing_tests/xdp_pass.py
Normal file
54
tests/failing_tests/xdp_pass.py
Normal file
@ -0,0 +1,54 @@
|
||||
from pythonbpf import bpf, map, section, bpfglobal, compile_to_ir
|
||||
from pythonbpf.maps import HashMap
|
||||
from pythonbpf.helper import XDP_PASS
|
||||
from vmlinux import TASK_COMM_LEN # noqa: F401
|
||||
|
||||
from vmlinux import struct_qspinlock # noqa: F401
|
||||
|
||||
# from vmlinux import struct_trace_event_raw_sys_enter # noqa: F401
|
||||
# from vmlinux import struct_posix_cputimers # noqa: F401
|
||||
from vmlinux import struct_xdp_md
|
||||
|
||||
# from vmlinux import struct_trace_event_raw_sys_enter # noqa: F401
|
||||
# from vmlinux import struct_ring_buffer_per_cpu # noqa: F401
|
||||
# from vmlinux import struct_request # noqa: F401
|
||||
from ctypes import c_int64
|
||||
|
||||
# Instructions to how to run this program
|
||||
# 1. Install PythonBPF: pip install pythonbpf
|
||||
# 2. Run the program: python examples/xdp_pass.py
|
||||
# 3. Run the program with sudo: sudo tools/check.sh run examples/xdp_pass.o
|
||||
# 4. Attach object file to any network device with something like ./check.sh xdp examples/xdp_pass.o tailscale0
|
||||
# 5. send traffic through the device and observe effects
|
||||
|
||||
|
||||
@bpf
|
||||
@map
|
||||
def count() -> HashMap:
|
||||
return HashMap(key=c_int64, value=c_int64, max_entries=1)
|
||||
|
||||
|
||||
@bpf
|
||||
@section("xdp")
|
||||
def hello_world(ctx: struct_xdp_md) -> c_int64:
|
||||
key = 0
|
||||
one = 1
|
||||
prev = count().lookup(key)
|
||||
if prev:
|
||||
prevval = prev + 1
|
||||
print(f"count: {prevval}")
|
||||
count().update(key, prevval)
|
||||
return XDP_PASS
|
||||
else:
|
||||
count().update(key, one)
|
||||
|
||||
return XDP_PASS
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def LICENSE() -> str:
|
||||
return "GPL"
|
||||
|
||||
|
||||
compile_to_ir("xdp_pass.py", "xdp_pass.ll")
|
||||
@ -9,6 +9,7 @@ from pythonbpf.helper import ktime
|
||||
# the compiler, it is recommended to use named variables to reduce the amount of
|
||||
# scratch space that needs to be allocated.
|
||||
|
||||
|
||||
@bpf
|
||||
@struct
|
||||
class data_t:
|
||||
@ -48,8 +49,12 @@ def hello_world(ctx: c_void_p) -> c_int64:
|
||||
a = last.lookup(0)
|
||||
print(f"a is {a}")
|
||||
last.update(9, 9)
|
||||
last.update(0, last.lookup(last.lookup(0)) +
|
||||
last.lookup(last.lookup(0)) + last.lookup(last.lookup(0)))
|
||||
last.update(
|
||||
0,
|
||||
last.lookup(last.lookup(0))
|
||||
+ last.lookup(last.lookup(0))
|
||||
+ last.lookup(last.lookup(0)),
|
||||
)
|
||||
z = last.lookup(0)
|
||||
print(f"new map val at index 0 is {z}")
|
||||
else:
|
||||
|
||||
@ -6,8 +6,8 @@ from ctypes import c_void_p, c_int32
|
||||
@section("tracepoint/syscalls/sys_enter_execve")
|
||||
def hello_world(ctx: c_void_p) -> c_int32:
|
||||
print("Hello, World!")
|
||||
a = 1 # int64
|
||||
return c_int32(a) # typecast to int32
|
||||
a = 1 # int64
|
||||
return c_int32(a) # typecast to int32
|
||||
|
||||
|
||||
@bpf
|
||||
|
||||
@ -26,8 +26,13 @@ import tempfile
|
||||
|
||||
|
||||
class BTFConverter:
|
||||
def __init__(self, btf_source="/sys/kernel/btf/vmlinux", output_file="vmlinux.py",
|
||||
keep_intermediate=False, verbose=False):
|
||||
def __init__(
|
||||
self,
|
||||
btf_source="/sys/kernel/btf/vmlinux",
|
||||
output_file="vmlinux.py",
|
||||
keep_intermediate=False,
|
||||
verbose=False,
|
||||
):
|
||||
self.btf_source = btf_source
|
||||
self.output_file = output_file
|
||||
self.keep_intermediate = keep_intermediate
|
||||
@ -44,11 +49,7 @@ class BTFConverter:
|
||||
self.log(f"{description}...")
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
shell=True,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True
|
||||
cmd, shell=True, check=True, capture_output=True, text=True
|
||||
)
|
||||
if self.verbose and result.stdout:
|
||||
print(result.stdout)
|
||||
@ -69,51 +70,55 @@ class BTFConverter:
|
||||
"""Step 1.5: Preprocess enum definitions."""
|
||||
self.log("Preprocessing enum definitions...")
|
||||
|
||||
with open(input_file, 'r') as f:
|
||||
with open(input_file, "r") as f:
|
||||
original_code = f.read()
|
||||
|
||||
# Extract anonymous enums
|
||||
enums = re.findall(
|
||||
r'(?<!typedef\s)(enum\s*\{[^}]*\})\s*(\w+)\s*(?::\s*\d+)?\s*;',
|
||||
original_code
|
||||
r"(?<!typedef\s)(enum\s*\{[^}]*\})\s*(\w+)\s*(?::\s*\d+)?\s*;",
|
||||
original_code,
|
||||
)
|
||||
enum_defs = [enum_block + ';' for enum_block, _ in enums]
|
||||
enum_defs = [enum_block + ";" for enum_block, _ in enums]
|
||||
|
||||
# Replace anonymous enums with int declarations
|
||||
processed_code = re.sub(
|
||||
r'(?<!typedef\s)enum\s*\{[^}]*\}\s*(\w+)\s*(?::\s*\d+)?\s*;',
|
||||
r'int \1;',
|
||||
original_code
|
||||
r"(?<!typedef\s)enum\s*\{[^}]*\}\s*(\w+)\s*(?::\s*\d+)?\s*;",
|
||||
r"int \1;",
|
||||
original_code,
|
||||
)
|
||||
|
||||
# Prepend enum definitions
|
||||
if enum_defs:
|
||||
enum_text = '\n'.join(enum_defs) + '\n\n'
|
||||
enum_text = "\n".join(enum_defs) + "\n\n"
|
||||
processed_code = enum_text + processed_code
|
||||
|
||||
output_file = os.path.join(self.temp_dir, "vmlinux_processed.h")
|
||||
with open(output_file, 'w') as f:
|
||||
with open(output_file, "w") as f:
|
||||
f.write(processed_code)
|
||||
|
||||
return output_file
|
||||
|
||||
def step2_5_process_kioctx(self, input_file):
|
||||
#TODO: this is a very bad bug and design decision. A single struct has an issue mostly.
|
||||
# TODO: this is a very bad bug and design decision. A single struct has an issue mostly.
|
||||
"""Step 2.5: Process struct kioctx to extract nested anonymous structs."""
|
||||
self.log("Processing struct kioctx nested structs...")
|
||||
|
||||
with open(input_file, 'r') as f:
|
||||
with open(input_file, "r") as f:
|
||||
content = f.read()
|
||||
|
||||
# Pattern to match struct kioctx with its full body (handles multiple nesting levels)
|
||||
kioctx_pattern = r'struct\s+kioctx\s*\{(?:[^{}]|\{(?:[^{}]|\{[^{}]*\})*\})*\}\s*;'
|
||||
kioctx_pattern = (
|
||||
r"struct\s+kioctx\s*\{(?:[^{}]|\{(?:[^{}]|\{[^{}]*\})*\})*\}\s*;"
|
||||
)
|
||||
|
||||
def process_kioctx_replacement(match):
|
||||
full_struct = match.group(0)
|
||||
self.log(f"Found struct kioctx, length: {len(full_struct)} chars")
|
||||
|
||||
# Extract the struct body (everything between outermost { and })
|
||||
body_match = re.search(r'struct\s+kioctx\s*\{(.*)\}\s*;', full_struct, re.DOTALL)
|
||||
body_match = re.search(
|
||||
r"struct\s+kioctx\s*\{(.*)\}\s*;", full_struct, re.DOTALL
|
||||
)
|
||||
if not body_match:
|
||||
return full_struct
|
||||
|
||||
@ -121,7 +126,7 @@ class BTFConverter:
|
||||
|
||||
# Find all anonymous structs within the body
|
||||
# Pattern: struct { ... } followed by ; (not a member name)
|
||||
anon_struct_pattern = r'struct\s*\{[^}]*\}'
|
||||
# anon_struct_pattern = r"struct\s*\{[^}]*\}"
|
||||
|
||||
anon_structs = []
|
||||
anon_counter = 4 # Start from 4, counting down to 1
|
||||
@ -131,7 +136,9 @@ class BTFConverter:
|
||||
anon_struct_content = m.group(0)
|
||||
|
||||
# Extract the body of the anonymous struct
|
||||
anon_body_match = re.search(r'struct\s*\{(.*)\}', anon_struct_content, re.DOTALL)
|
||||
anon_body_match = re.search(
|
||||
r"struct\s*\{(.*)\}", anon_struct_content, re.DOTALL
|
||||
)
|
||||
if not anon_body_match:
|
||||
return anon_struct_content
|
||||
|
||||
@ -154,7 +161,7 @@ class BTFConverter:
|
||||
processed_body = body
|
||||
|
||||
# Find all occurrences and process them
|
||||
pattern_with_semicolon = r'struct\s*\{([^}]*)\}\s*;'
|
||||
pattern_with_semicolon = r"struct\s*\{([^}]*)\}\s*;"
|
||||
matches = list(re.finditer(pattern_with_semicolon, body, re.DOTALL))
|
||||
|
||||
if not matches:
|
||||
@ -178,14 +185,16 @@ class BTFConverter:
|
||||
|
||||
# Replace in the body
|
||||
replacement = f"struct {anon_name} {member_name};"
|
||||
processed_body = processed_body[:start_pos] + replacement + processed_body[end_pos:]
|
||||
processed_body = (
|
||||
processed_body[:start_pos] + replacement + processed_body[end_pos:]
|
||||
)
|
||||
|
||||
anon_counter -= 1
|
||||
|
||||
# Rebuild the complete definition
|
||||
if anon_structs:
|
||||
# Prepend the anonymous struct definitions
|
||||
anon_definitions = '\n'.join(anon_structs) + '\n\n'
|
||||
anon_definitions = "\n".join(anon_structs) + "\n\n"
|
||||
new_struct = f"struct kioctx {{{processed_body}}};"
|
||||
return anon_definitions + new_struct
|
||||
else:
|
||||
@ -193,14 +202,11 @@ class BTFConverter:
|
||||
|
||||
# Apply the transformation
|
||||
processed_content = re.sub(
|
||||
kioctx_pattern,
|
||||
process_kioctx_replacement,
|
||||
content,
|
||||
flags=re.DOTALL
|
||||
kioctx_pattern, process_kioctx_replacement, content, flags=re.DOTALL
|
||||
)
|
||||
|
||||
output_file = os.path.join(self.temp_dir, "vmlinux_kioctx_processed.h")
|
||||
with open(output_file, 'w') as f:
|
||||
with open(output_file, "w") as f:
|
||||
f.write(processed_content)
|
||||
|
||||
self.log(f"Saved kioctx-processed output to {output_file}")
|
||||
@ -218,7 +224,7 @@ class BTFConverter:
|
||||
output_file = os.path.join(self.temp_dir, "vmlinux_raw.py")
|
||||
cmd = (
|
||||
f"clang2py {input_file} -o {output_file} "
|
||||
f"--clang-args=\"-fno-ms-extensions -I/usr/include -I/usr/include/linux\""
|
||||
f'--clang-args="-fno-ms-extensions -I/usr/include -I/usr/include/linux"'
|
||||
)
|
||||
self.run_command(cmd, "Converting to Python ctypes")
|
||||
return output_file
|
||||
@ -234,14 +240,21 @@ class BTFConverter:
|
||||
data = re.sub(r"\('_[0-9]+',\s*ctypes\.[a-zA-Z0-9_]+,\s*0\),?\s*\n?", "", data)
|
||||
|
||||
# Replace ('_20', ctypes.c_uint64, 64) → ('_20', ctypes.c_uint64)
|
||||
data = re.sub(r"\('(_[0-9]+)',\s*(ctypes\.[a-zA-Z0-9_]+),\s*[0-9]+\)", r"('\1', \2)", data)
|
||||
data = re.sub(
|
||||
r"\('(_[0-9]+)',\s*(ctypes\.[a-zA-Z0-9_]+),\s*[0-9]+\)", r"('\1', \2)", data
|
||||
)
|
||||
|
||||
# Replace ('_20', ctypes.c_char, 8) with ('_20', ctypes.c_uint8, 8)
|
||||
data = re.sub(
|
||||
r"(ctypes\.c_char)(\s*,\s*\d+\))",
|
||||
r"ctypes.c_uint8\2",
|
||||
data
|
||||
)
|
||||
data = re.sub(r"(ctypes\.c_char)(\s*,\s*\d+\))", r"ctypes.c_uint8\2", data)
|
||||
|
||||
# below to replace those c_bool with bitfield greater than 8
|
||||
def repl(m):
|
||||
name, bits = m.groups()
|
||||
return (
|
||||
f"('{name}', ctypes.c_uint32, {bits})" if int(bits) > 8 else m.group(0)
|
||||
)
|
||||
|
||||
data = re.sub(r"\('([^']+)',\s*ctypes\.c_bool,\s*(\d+)\)", repl, data)
|
||||
|
||||
# Remove ctypes. prefix from invalid entries
|
||||
invalid_ctypes = ["bpf_iter_state", "_cache_type", "fs_context_purpose"]
|
||||
@ -258,6 +271,7 @@ class BTFConverter:
|
||||
if not self.keep_intermediate and self.temp_dir != ".":
|
||||
self.log(f"Cleaning up temporary directory: {self.temp_dir}")
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def convert(self):
|
||||
@ -281,6 +295,7 @@ class BTFConverter:
|
||||
except Exception as e:
|
||||
print(f"\n✗ Error during conversion: {e}", file=sys.stderr)
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
finally:
|
||||
@ -293,18 +308,13 @@ class BTFConverter:
|
||||
dependencies = {
|
||||
"bpftool": "bpftool --version",
|
||||
"clang": "clang --version",
|
||||
"clang2py": "clang2py --version"
|
||||
"clang2py": "clang2py --version",
|
||||
}
|
||||
|
||||
missing = []
|
||||
for tool, cmd in dependencies.items():
|
||||
try:
|
||||
subprocess.run(
|
||||
cmd,
|
||||
shell=True,
|
||||
check=True,
|
||||
capture_output=True
|
||||
)
|
||||
subprocess.run(cmd, shell=True, check=True, capture_output=True)
|
||||
except subprocess.CalledProcessError:
|
||||
missing.append(tool)
|
||||
|
||||
@ -326,31 +336,31 @@ Examples:
|
||||
%(prog)s
|
||||
%(prog)s -o kernel_types.py
|
||||
%(prog)s --btf-source /sys/kernel/btf/custom_module -k -v
|
||||
"""
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--btf-source",
|
||||
default="/sys/kernel/btf/vmlinux",
|
||||
help="Path to BTF source (default: /sys/kernel/btf/vmlinux)"
|
||||
help="Path to BTF source (default: /sys/kernel/btf/vmlinux)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-o", "--output",
|
||||
"-o",
|
||||
"--output",
|
||||
default="vmlinux.py",
|
||||
help="Output Python file (default: vmlinux.py)"
|
||||
help="Output Python file (default: vmlinux.py)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-k", "--keep-intermediate",
|
||||
"-k",
|
||||
"--keep-intermediate",
|
||||
action="store_true",
|
||||
help="Keep intermediate files (vmlinux.h, vmlinux_processed.h, etc.)"
|
||||
help="Keep intermediate files (vmlinux.h, vmlinux_processed.h, etc.)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-v", "--verbose",
|
||||
action="store_true",
|
||||
help="Enable verbose output"
|
||||
"-v", "--verbose", action="store_true", help="Enable verbose output"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
@ -359,7 +369,7 @@ Examples:
|
||||
btf_source=args.btf_source,
|
||||
output_file=args.output,
|
||||
keep_intermediate=args.keep_intermediate,
|
||||
verbose=args.verbose
|
||||
verbose=args.verbose,
|
||||
)
|
||||
|
||||
converter.convert()
|
||||
|
||||
Reference in New Issue
Block a user