From b9082ffbeb15bee0c1f555604c3f384a591a0f65 Mon Sep 17 00:00:00 2001 From: mahesh bhatiya Date: Thu, 12 Mar 2026 23:16:13 +0530 Subject: [PATCH] feat(ctx): add ctx.load_bytes for copying context memory into event buffers --- src/ast/unit.rs | 1 + src/emit/ebpf_c/tracepoint.rs | 24 +++++++++- src/ir/ctx.rs | 17 ++++---- src/ir/instruction.rs | 2 + src/ir/unit.rs | 82 +++++++++++++++++++++++++++++------ src/lexer/mod.rs | 2 +- 6 files changed, 103 insertions(+), 25 deletions(-) diff --git a/src/ast/unit.rs b/src/ast/unit.rs index 68165d5..6f51b6f 100644 --- a/src/ast/unit.rs +++ b/src/ast/unit.rs @@ -223,6 +223,7 @@ impl ProgramKind { LoadI16, LoadI32, LoadI64, + LoadBytes, GetPidTgid, GetUidGid, GetCurrentComm, diff --git a/src/emit/ebpf_c/tracepoint.rs b/src/emit/ebpf_c/tracepoint.rs index 69aa810..08c7305 100644 --- a/src/emit/ebpf_c/tracepoint.rs +++ b/src/emit/ebpf_c/tracepoint.rs @@ -196,7 +196,8 @@ pub fn emit_tracepoint(out: &mut String, unit: &UnitIr, sec: &str) -> Result<(), if let Some(operand) = inst.operands.get(0) { let ptr = format_operand(operand); let res = res_name.unwrap(); - writeln!(out, " {} = *((__u64 *)(void *)({}));", res, ptr).map_err(err)?; + writeln!(out, " {} = *((__u64 *)(void *)({}));", res, ptr) + .map_err(err)?; } } @@ -204,7 +205,14 @@ pub fn emit_tracepoint(out: &mut String, unit: &UnitIr, sec: &str) -> Result<(), if inst.operands.len() >= 2 { let ptr = format_operand(&inst.operands[0]); let val = format_operand(&inst.operands[1]); - writeln!(out, " *((__u{} *)(void *)({})) = {};", size * 8, ptr, val).map_err(err)?; + writeln!( + out, + " *((__u{} *)(void *)({})) = {};", + size * 8, + ptr, + val + ) + .map_err(err)?; } } @@ -358,6 +366,18 @@ pub fn emit_tracepoint(out: &mut String, unit: &UnitIr, sec: &str) -> Result<(), writeln!(out, " bpf_ringbuf_submit({}, 0);", ptr).map_err(err)?; } } + Opcode::CopyCtxToMem { offset, size } => { + let dest = match &inst.operands[0] { + Operand::Var(v) => format!("v{}", v.0), + _ => panic!("CopyCtxToMem expects dest var"), + }; + + writeln!( + out, + "bpf_probe_read_kernel({}, {}, (void *)ctx + {});", + dest, size, offset + ); + } } inst_index += 1; diff --git a/src/ir/ctx.rs b/src/ir/ctx.rs index a315b6d..04a92df 100644 --- a/src/ir/ctx.rs +++ b/src/ir/ctx.rs @@ -11,28 +11,28 @@ pub enum CtxMethod { LoadI16, LoadI32, LoadI64, + LoadBytes, // ================================ // Process / task metadata helpers // ================================ - GetPidTgid, // bpf_get_current_pid_tgid (helper 14) - GetUidGid, // bpf_get_current_uid_gid - GetCurrentComm, // bpf_get_current_comm - GetCurrentTask, // bpf_get_current_task + GetPidTgid, // bpf_get_current_pid_tgid (helper 14) + GetUidGid, // bpf_get_current_uid_gid + GetCurrentComm, // bpf_get_current_comm + GetCurrentTask, // bpf_get_current_task // ================================ // Time helpers // ================================ - GetKtimeNs, // bpf_ktime_get_ns + GetKtimeNs, // bpf_ktime_get_ns // ================================ // Memory probe helpers // ================================ - ProbeReadUserStr, // bpf_probe_read_user_str (helper 202) + ProbeReadUserStr, // bpf_probe_read_user_str (helper 202) ProbeReadKernelStr, // bpf_probe_read_kernel_str (helper 204) } - impl CtxMethod { pub fn from_str(name: &str) -> Option { match name { @@ -45,6 +45,7 @@ impl CtxMethod { "load_i16" => Some(Self::LoadI16), "load_i32" => Some(Self::LoadI32), "load_i64" => Some(Self::LoadI64), + "load_bytes" => Some(Self::LoadBytes), // Metadata helpers "get_pid_tgid" => Some(Self::GetPidTgid), @@ -62,4 +63,4 @@ impl CtxMethod { _ => None, } } -} \ No newline at end of file +} diff --git a/src/ir/instruction.rs b/src/ir/instruction.rs index 1dc8798..adc8008 100644 --- a/src/ir/instruction.rs +++ b/src/ir/instruction.rs @@ -23,6 +23,8 @@ pub enum Opcode { LoadCtx { offset: i32, size: u8 }, LoadPacket { offset: i32, size: u8 }, + CopyCtxToMem { offset: i32, size: u32 }, + // checks NullCheck, diff --git a/src/ir/unit.rs b/src/ir/unit.rs index 2cf207a..a69446f 100644 --- a/src/ir/unit.rs +++ b/src/ir/unit.rs @@ -72,7 +72,14 @@ impl UnitIr { }; for stmt in &unit.body { - lower_statement(stmt, &mut ctx, &mut ir, &mut current_block, events, event_decls)?; + lower_statement( + stmt, + &mut ctx, + &mut ir, + &mut current_block, + events, + event_decls, + )?; } ir.blocks.push(current_block); @@ -117,6 +124,7 @@ fn lower_ctx_helper( | CtxMethod::LoadI16 | CtxMethod::LoadI32 | CtxMethod::LoadI64 + | CtxMethod::LoadBytes | CtxMethod::ProbeReadUserStr | CtxMethod::ProbeReadKernelStr => { return Err(LoweringError::UnitLowering(format!( @@ -194,7 +202,8 @@ fn lower_statement( ))); } - let key = lower_expr(&call.arg[0], ctx, ir, block, events, event_decls)?; + let key = + lower_expr(&call.arg[0], ctx, ir, block, events, event_decls)?; let result = ir.alloc_var(crate::ast::Type::U64); block.instructions.push(Instruction { @@ -384,7 +393,7 @@ fn lower_statement( // Handle field assignment like evt.pid = value // First, get the base pointer/variable let base_operand = lower_expr(base, ctx, ir, block, events, event_decls)?; - + // Get the variable ID (should be a pointer from reserve) let base_var = if let Operand::Var(v) = base_operand { v @@ -397,12 +406,14 @@ fn lower_statement( // Look up field offset across all events let mut field_offset: Option = None; for (_, event_decl) in event_decls { - if let Some(offset) = crate::sema::event::compute_field_offset(event_decl, field) { + if let Some(offset) = + crate::sema::event::compute_field_offset(event_decl, field) + { field_offset = Some(offset); break; } } - + let _offset = field_offset.ok_or_else(|| { LoweringError::UnitLowering(format!("Unknown field: {}", field)) })?; @@ -410,7 +421,8 @@ fn lower_statement( // For now, we only support simple field assignments (no compound ops on fields) if !matches!(assign.op, crate::ast::AssignmentOp::Assign) { return Err(LoweringError::UnitLowering( - "Compound assignment operators not supported on struct fields".to_string(), + "Compound assignment operators not supported on struct fields" + .to_string(), )); } @@ -589,7 +601,8 @@ fn lower_expr( ))); } - let offset_expr = lower_expr(&call.arg[0], ctx, ir, block, events, event_decls)?; + let offset_expr = + lower_expr(&call.arg[0], ctx, ir, block, events, event_decls)?; let offset = match offset_expr { Operand::Immediate(n) => n as i32, _ => { @@ -639,7 +652,8 @@ fn lower_expr( } let dest = lower_expr(&call.arg[0], ctx, ir, block, events, event_decls)?; - let size_expr = lower_expr(&call.arg[1], ctx, ir, block, events, event_decls)?; + let size_expr = + lower_expr(&call.arg[1], ctx, ir, block, events, event_decls)?; let src = lower_expr(&call.arg[2], ctx, ir, block, events, event_decls)?; let size = match size_expr { @@ -672,7 +686,8 @@ fn lower_expr( } let dest = lower_expr(&call.arg[0], ctx, ir, block, events, event_decls)?; - let size_expr = lower_expr(&call.arg[1], ctx, ir, block, events, event_decls)?; + let size_expr = + lower_expr(&call.arg[1], ctx, ir, block, events, event_decls)?; let src = lower_expr(&call.arg[2], ctx, ir, block, events, event_decls)?; let size = match size_expr { @@ -695,6 +710,46 @@ fn lower_expr( Ok(Operand::Var(result)) } + CtxMethod::LoadBytes => { + if call.arg.len() != 3 { + return Err(LoweringError::UnitLowering( + "ctx.load_bytes expects (offset, dest, size)".to_string(), + )); + } + + let offset_expr = + lower_expr(&call.arg[0], ctx, ir, block, events, event_decls)?; + let dest = lower_expr(&call.arg[1], ctx, ir, block, events, event_decls)?; + let size_expr = + lower_expr(&call.arg[2], ctx, ir, block, events, event_decls)?; + + let offset = match offset_expr { + Operand::Immediate(n) => n as i32, + _ => { + return Err(LoweringError::UnitLowering( + "offset must be immediate".to_string(), + )) + } + }; + + let size = match size_expr { + Operand::Immediate(n) => n as u32, + _ => { + return Err(LoweringError::UnitLowering( + "size must be immediate".to_string(), + )) + } + }; + + block.instructions.push(Instruction { + result: ir.alloc_var(crate::ast::Type::U64), + opcode: Opcode::CopyCtxToMem { offset, size }, + operands: vec![dest.clone()], + result_type: crate::ast::Type::U64, + }); + + Ok(dest) + } // ctx helper methods (0 args) _ => { if !call.arg.is_empty() { @@ -908,7 +963,7 @@ fn lower_expr( ExprKind::FieldAccess { base, field } => { // Handle field access like evt.filename let base_operand = lower_expr(base, ctx, ir, block, events, event_decls)?; - + // Get the variable ID (should be a pointer from reserve) let base_var = if let Operand::Var(v) = base_operand { v @@ -926,10 +981,9 @@ fn lower_expr( break; } } - - let _offset = field_offset.ok_or_else(|| { - LoweringError::UnitLowering(format!("Unknown field: {}", field)) - })?; + + let _offset = field_offset + .ok_or_else(|| LoweringError::UnitLowering(format!("Unknown field: {}", field)))?; // Create a load instruction from the field // TODO: Currently we load from base pointer with fixed size diff --git a/src/lexer/mod.rs b/src/lexer/mod.rs index 88a591e..6e268d2 100644 --- a/src/lexer/mod.rs +++ b/src/lexer/mod.rs @@ -132,7 +132,7 @@ impl<'src> Lexer<'src> { "guard" => TokenKind::KeywordGuard, "heap" => TokenKind::KeywordHeap, "event" => TokenKind::KeywordEvent, - "bytes" => TokenKind::KeywordBytes, // ADD THIS + "bytes" => TokenKind::KeywordBytes, "u32" => TokenKind::TypeU32, "u64" => TokenKind::TypeU64, "i32" => TokenKind::TypeI32,