python-bpf/pythonbpf/codegen.py

import ast
from llvmlite import ir
from .license_pass import license_processing
from .functions_pass import func_proc
from .maps import maps_proc
from .structs import structs_proc
from .globals_pass import globals_processing
from .debuginfo import DW_LANG_C11, DwarfBehaviorEnum
import os
import subprocess
import inspect
from pathlib import Path
from pylibbpf import BpfProgram
import tempfile

VERSION = "v0.1.3"


def find_bpf_chunks(tree):
    """Find all functions decorated with @bpf in the AST."""
    bpf_functions = []
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef) or isinstance(node, ast.ClassDef):
            for decorator in node.decorator_list:
                if isinstance(decorator, ast.Name) and decorator.id == "bpf":
                    bpf_functions.append(node)
                    break
    return bpf_functions


def processor(source_code, filename, module):
    tree = ast.parse(source_code, filename)
    print(ast.dump(tree, indent=4))

    bpf_chunks = find_bpf_chunks(tree)
    for func_node in bpf_chunks:
        print(f"Found BPF function/struct: {func_node.name}")

    structs_sym_tab = structs_proc(tree, module, bpf_chunks)
    map_sym_tab = maps_proc(tree, module, bpf_chunks)
    func_proc(tree, module, bpf_chunks, map_sym_tab, structs_sym_tab)

    license_processing(tree, module)
    globals_processing(tree, module)


def compile_to_ir(filename: str, output: str):
    with open(filename) as f:
        source = f.read()

    module = ir.Module(name=filename)
    module.data_layout = "e-m:e-p:64:64-i64:64-i128:128-n32:64-S128"
    module.triple = "bpf"

    if not hasattr(module, "_debug_compile_unit"):
        module._file_metadata = module.add_debug_info(
            "DIFile",
            {  # type: ignore
                "filename": filename,
                "directory": os.path.dirname(filename),
            },
        )

        module._debug_compile_unit = module.add_debug_info(
            "DICompileUnit",
            {  # type: ignore
                "language": DW_LANG_C11,
                "file": module._file_metadata,  # type: ignore
                "producer": f"PythonBPF {VERSION}",
                "isOptimized": True,  # TODO: This is probably not true
                # TODO: add a global field here that keeps track of all the globals. Works without it, but I think it might
                # be required for kprobes.
                "runtimeVersion": 0,
                "emissionKind": 1,
                "splitDebugInlining": False,
                "nameTableKind": 0,
            },
            is_distinct=True,
        )

        module.add_named_metadata("llvm.dbg.cu", module._debug_compile_unit)  # type: ignore

    processor(source, filename, module)

    wchar_size = module.add_metadata(
        [
            DwarfBehaviorEnum.ERROR_IF_MISMATCH,
            "wchar_size",
            ir.Constant(ir.IntType(32), 4),
        ]
    )
    frame_pointer = module.add_metadata(
        [
            DwarfBehaviorEnum.OVERRIDE_USE_LARGEST,
            "frame-pointer",
            ir.Constant(ir.IntType(32), 2),
        ]
    )
    # Add Debug Info Version (3 = DWARF v3, which LLVM expects)
    debug_info_version = module.add_metadata(
        [
            DwarfBehaviorEnum.WARNING_IF_MISMATCH,
            "Debug Info Version",
            ir.Constant(ir.IntType(32), 3),
        ]
    )

    # Add explicit DWARF version 5
    dwarf_version = module.add_metadata(
        [
            DwarfBehaviorEnum.OVERRIDE_USE_LARGEST,
            "Dwarf Version",
            ir.Constant(ir.IntType(32), 5),
        ]
    )

    module.add_named_metadata("llvm.module.flags", wchar_size)
    module.add_named_metadata("llvm.module.flags", frame_pointer)
    module.add_named_metadata("llvm.module.flags", debug_info_version)
    module.add_named_metadata("llvm.module.flags", dwarf_version)

    module.add_named_metadata("llvm.ident", [f"PythonBPF {VERSION}"])

    print(f"IR written to {output}")
    with open(output, "w") as f:
        f.write(f'source_filename = "{filename}"\n')
        f.write(str(module))
        f.write("\n")

    return output


def compile() -> bool:
    # Look one level up the stack to the caller of this function
    caller_frame = inspect.stack()[1]
    caller_file = Path(caller_frame.filename).resolve()

    ll_file = Path("/tmp") / caller_file.with_suffix(".ll").name
    o_file = caller_file.with_suffix(".o")

    success = True
    success = compile_to_ir(str(caller_file), str(ll_file)) and success

    success = (
        subprocess.run(
            [
                "llc",
                "-march=bpf",
                "-filetype=obj",
                "-O2",
                str(ll_file),
                "-o",
                str(o_file),
            ],
            check=True,
        )
        and success
    )

    print(f"Object written to {o_file}")
    return success


def BPF() -> BpfProgram:
    caller_frame = inspect.stack()[1]
    src = inspect.getsource(caller_frame.frame)
    with tempfile.NamedTemporaryFile(
        mode="w+", delete=True, suffix=".py"
    ) as f, tempfile.NamedTemporaryFile(
        mode="w+", delete=True, suffix=".ll"
    ) as inter, tempfile.NamedTemporaryFile(
        mode="w+", delete=False, suffix=".o"
    ) as obj_file:
        f.write(src)
        f.flush()
        source = f.name
        compile_to_ir(source, str(inter.name))
        subprocess.run(
            [
                "llc",
                "-march=bpf",
                "-filetype=obj",
                "-O2",
                str(inter.name),
                "-o",
                str(obj_file.name),
            ],
            check=True,
        )

        return BpfProgram(str(obj_file.name))