diff --git a/pythonbpf/codegen.py b/pythonbpf/codegen.py index 8acbf10..de3d5a8 100644 --- a/pythonbpf/codegen.py +++ b/pythonbpf/codegen.py @@ -5,14 +5,16 @@ from .functions_pass import functions_processing from .constants_pass import constants_processing from .globals_pass import globals_processing + def processor(source_code, filename, module): tree = ast.parse(source_code, filename) - print(ast.dump(tree)) + print(ast.dump(tree, indent=4)) constants_processing(tree, module) license_processing(tree, module) globals_processing(tree, module) functions_processing(tree, module) + def compile_to_ir(filename: str, output: str): with open(filename) as f: source = f.read() @@ -22,7 +24,6 @@ def compile_to_ir(filename: str, output: str): module.triple = "bpf" processor(source, filename, module) - wchar_size = module.add_metadata([ir.Constant(ir.IntType(32), 1), "wchar_size", ir.Constant(ir.IntType(32), 4)]) diff --git a/pythonbpf/constants_pass.py b/pythonbpf/constants_pass.py index cd5ce1f..d5e323c 100644 --- a/pythonbpf/constants_pass.py +++ b/pythonbpf/constants_pass.py @@ -1,12 +1,14 @@ from llvmlite import ir import ast + def emit_constants(module: ir.Module, constant_str: str, name: str): constant_bytes = constant_str.encode("utf8") + b"\x00" elems = [ir.Constant(ir.IntType(8), b) for b in constant_bytes] ty = ir.ArrayType(ir.IntType(8), len(elems)) gvar = ir.GlobalVariable(module, ty, name=name) + print("constant emitted:", name) gvar.initializer = ir.Constant(ty, elems) # type: ignore @@ -16,6 +18,7 @@ def emit_constants(module: ir.Module, constant_str: str, name: str): return gvar + def constants_processing(tree, module): """Process string constants in the given AST tree and emit them to rodata""" constant_count = 0 @@ -37,7 +40,8 @@ def constants_processing(tree, module): if constant_count == 0: constant_name = f"{current_function}.____fmt" else: - constant_name = f"{current_function}.____fmt.{constant_count}" + constant_name = f"""{current_function}.____fmt.{ + constant_count}""" emit_constants(module, node.value, constant_name) constant_count += 1 self.generic_visit(node) @@ -45,4 +49,4 @@ def constants_processing(tree, module): visitor = ConstantVisitor() visitor.visit(tree) - return constant_count \ No newline at end of file + return constant_count diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index e3a6102..3754918 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -1,5 +1,78 @@ from llvmlite import ir import ast + +def emit_function(module: ir.Module, name: str): + ret_type = ir.IntType(32) + ptr_type = ir.PointerType() + func_ty = ir.FunctionType(ret_type, [ptr_type]) + + func = ir.Function(module, func_ty, name) + + param = func.args[0] + param.add_attribute("nocapture") + + func.attributes.add("nounwind") +# func.attributes.add("\"frame-pointer\"=\"all\"") +# func.attributes.add("no-trapping-math", "true") +# func.attributes.add("stack-protector-buffer-size", "8") + + block = func.append_basic_block(name="entry") + builder = ir.IRBuilder(block) + fmt_gvar = module.get_global("hello.____fmt") + + if fmt_gvar is None: + # If you haven't created the format string global yet + print("Warning: Format string global not found") + else: + # Cast integer 6 to function pointer type + fn_type = ir.FunctionType(ir.IntType( + 64), [ptr_type, ir.IntType(32)], var_arg=True) + fn_ptr_type = ir.PointerType(fn_type) + fn_addr = ir.Constant(ir.IntType(64), 6) + fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) + # Call the function + builder.call(fn_ptr, [fmt_gvar, ir.Constant(ir.IntType(32), 14)]) + + builder.ret(ir.Constant(ret_type, 0)) + + func.return_value.add_attribute("noundef") + func.linkage = "dso_local" + func.section = "kprobe/sys_clone" + print("function emitted:", name) + return func + + def functions_processing(tree, module): - pass + bpf_functions = [] + helper_functions = [] + for node in tree.body: + section_name = "" + if isinstance(node, ast.FunctionDef): + if len(node.decorator_list) == 1: + bpf_functions.append(node) + node.end_lineno + else: + # IDK why this check is needed, but whatever + if 'helper_functions' not in locals(): + helper_functions.append(node) + + # TODO: implement helpers first + + for func in bpf_functions: + dec = func.decorator_list[0] + if ( + isinstance(dec, ast.Call) + and isinstance(dec.func, ast.Name) + and dec.func.id == "section" + and len(dec.args) == 1 + and isinstance(dec.args[0], ast.Constant) + and isinstance(dec.args[0].value, str) + ): + section_name = dec.args[0].value + else: + print(f"ERROR: Invalid decorator for function {func.name}") + continue + + # TODO: parse arguments and return type + emit_function(module, func.name + "func") diff --git a/pythonbpf/globals_pass.py b/pythonbpf/globals_pass.py index 22413cc..df32327 100644 --- a/pythonbpf/globals_pass.py +++ b/pythonbpf/globals_pass.py @@ -1,6 +1,7 @@ from llvmlite import ir import ast + def emit_globals(module: ir.Module, names: list[str]): """ Emit the @llvm.compiler.used global given a list of function/global names. @@ -15,7 +16,9 @@ def emit_globals(module: ir.Module, names: list[str]): g = module.get_global(name) else: g = ir.GlobalValue(module, ptr_ty, name) - + print("global emitted:", name) + print(isinstance(g, ir.GlobalVariable)) + print(isinstance(g, ir.Function)) elems.append(g.bitcast(ptr_ty)) gv = ir.GlobalVariable(module, used_array_ty, "llvm.compiler.used")