mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
100 lines
3.2 KiB
Python
100 lines
3.2 KiB
Python
import ast
|
|
import logging
|
|
from llvmlite import ir
|
|
from .type_deducer import ctypes_to_ir
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# TODO: Shall we allow the following syntax:
|
|
# struct MyStruct:
|
|
# field1: int
|
|
# field2: str(32)
|
|
# Where int is mapped to c_uint64?
|
|
# Shall we just int64, int32 and uint32 similarly?
|
|
|
|
|
|
def structs_proc(tree, module, chunks):
|
|
""" Process all class definitions to find BPF structs """
|
|
structs_sym_tab = {}
|
|
for cls_node in chunks:
|
|
if is_bpf_struct(cls_node):
|
|
print(f"Found BPF struct: {cls_node.name}")
|
|
struct_info = process_bpf_struct(cls_node, module)
|
|
structs_sym_tab[cls_node.name] = struct_info
|
|
return structs_sym_tab
|
|
|
|
|
|
def is_bpf_struct(cls_node):
|
|
return any(
|
|
isinstance(decorator, ast.Name) and decorator.id == "struct"
|
|
for decorator in cls_node.decorator_list
|
|
)
|
|
|
|
|
|
def process_bpf_struct(cls_node, module):
|
|
""" Process a single BPF struct definition """
|
|
|
|
fields = parse_struct_fields(cls_node)
|
|
total_size = calc_struct_size(fields.values())
|
|
struct_type = ir.LiteralStructType(fields.values())
|
|
logger.info(f"Created struct {cls_node.name} with fields {fields.keys()}")
|
|
return {
|
|
"type": struct_type,
|
|
"fields": fields,
|
|
"size": total_size,
|
|
}
|
|
|
|
|
|
def parse_struct_fields(cls_node):
|
|
""" Parse fields of a struct class node """
|
|
fields = {}
|
|
|
|
for item in cls_node.body:
|
|
if isinstance(item, ast.AnnAssign) and \
|
|
isinstance(item.target, ast.Name):
|
|
fields[item.target.id] = get_type_from_ann(item.annotation)
|
|
else:
|
|
logger.error(f"Unsupported struct field: {ast.dump(item)}")
|
|
raise TypeError(f"Unsupported field in {ast.dump(cls_node)}")
|
|
return fields
|
|
|
|
|
|
def get_type_from_ann(annotation):
|
|
""" Convert an AST annotation node to an LLVM IR type for struct fields"""
|
|
if isinstance(annotation, ast.Call) and \
|
|
isinstance(annotation.func, ast.Name):
|
|
if annotation.func.id == "str":
|
|
# Char array
|
|
# Assumes constant integer argument
|
|
length = annotation.args[0].value
|
|
return ir.ArrayType(ir.IntType(8), length)
|
|
elif isinstance(annotation, ast.Name):
|
|
# Int type, written as c_int64, c_uint32, etc.
|
|
return ctypes_to_ir(annotation.id)
|
|
|
|
raise TypeError(f"Unsupported annotation type: {ast.dump(annotation)}")
|
|
|
|
|
|
def calc_struct_size(field_types):
|
|
""" Calculate total size of the struct with alignment and padding """
|
|
curr_offset = 0
|
|
for ftype in field_types:
|
|
if isinstance(ftype, ir.IntType):
|
|
fsize = ftype.width // 8
|
|
alignment = fsize
|
|
elif isinstance(ftype, ir.ArrayType):
|
|
fsize = ftype.count * (ftype.element.width // 8)
|
|
alignment = ftype.element.width // 8
|
|
elif isinstance(ftype, ir.PointerType):
|
|
# We won't encounter this rn, but for the future
|
|
fsize = 8
|
|
alignment = 8
|
|
else:
|
|
raise TypeError(f"Unsupported field type: {ftype}")
|
|
|
|
padding = (alignment - (curr_offset % alignment)) % alignment
|
|
curr_offset += padding + fsize
|
|
|
|
final_padding = (8 - (curr_offset % 8)) % 8
|
|
return curr_offset + final_padding
|