From bf7fdff1cef28961f096d1c1e00181e0a0c2d14e Mon Sep 17 00:00:00 2001 From: Dave Tucker Date: Wed, 12 Oct 2022 14:27:00 +0000 Subject: [PATCH 1/3] aya: Find programs using the symbol table This makes a few changes to the way that Aya reads the ELF object files. 1. To find programs in a section, we use the symbols table. This allows for cases where multiple programs could appear in the same section. 2. When parsing our ELF file we build symbols_by_section_index as an optimization as we use it for legacy maps, BTF maps and now programs. As a result of theses changes the "NAME" used in `bpf.prog_mut("NAME")` is now ALWAYS the same as the function name in the eBPF code, making the user experience more consistent. Signed-off-by: Dave Tucker --- aya-obj/src/obj.rs | 487 ++++++++++++++++------- test/integration-ebpf/src/name_test.rs | 2 +- test/integration-test/bpf/ext.bpf.c | 2 +- test/integration-test/bpf/main.bpf.c | 2 +- test/integration-test/src/tests/load.rs | 26 +- test/integration-test/src/tests/rbpf.rs | 6 +- test/integration-test/src/tests/smoke.rs | 9 +- xtask/public-api/aya-obj.txt | 6 +- 8 files changed, 375 insertions(+), 165 deletions(-) diff --git a/aya-obj/src/obj.rs b/aya-obj/src/obj.rs index a36eb684..3a9eb038 100644 --- a/aya-obj/src/obj.rs +++ b/aya-obj/src/obj.rs @@ -123,6 +123,7 @@ pub struct Object { pub functions: BTreeMap<(usize, u64), Function>, pub(crate) relocations: HashMap>, pub(crate) symbol_table: HashMap, + pub(crate) symbols_by_section: HashMap>, pub(crate) section_infos: HashMap, // symbol_offset_by_name caches symbols that could be referenced from a // BTF VAR type so the offsets can be fixed up @@ -600,7 +601,13 @@ impl Object { kind: symbol.kind(), }; bpf_obj.symbol_table.insert(symbol.index().0, sym); - + if let Some(section_idx) = symbol.section().index() { + bpf_obj + .symbols_by_section + .entry(section_idx) + .or_default() + .push(symbol.index().0); + } if symbol.is_global() || symbol.kind() == SymbolKind::Data { bpf_obj.symbol_offset_by_name.insert(name, symbol.address()); } @@ -642,6 +649,7 @@ impl Object { functions: BTreeMap::new(), relocations: HashMap::new(), symbol_table: HashMap::new(), + symbols_by_section: HashMap::new(), section_infos: HashMap::new(), symbol_offset_by_name: HashMap::new(), } @@ -711,10 +719,42 @@ impl Object { Ok(()) } - fn parse_program(&self, section: &Section) -> Result<(Program, Function), ParseError> { - let prog_sec = ProgramSection::from_str(section.name)?; - let name = prog_sec.name().to_owned(); + fn parse_programs(&mut self, section: &Section) -> Result<(), ParseError> { + let program_section = ProgramSection::from_str(section.name)?; + let syms = + self.symbols_by_section + .get(§ion.index) + .ok_or(ParseError::NoSymbolsForSection { + section_name: section.name.to_string(), + })?; + for symbol_index in syms { + let symbol = self + .symbol_table + .get(symbol_index) + .expect("all symbols in symbols_by_section are also in symbol_table"); + + let Some(name) = symbol.name.as_ref() else { + continue; + }; + if name.is_empty() { + continue; + } + let (p, f) = + self.parse_program(section, program_section.clone(), name.to_string(), symbol)?; + let key = p.function_key(); + self.programs.insert(f.name.clone(), p); + self.functions.insert(key, f); + } + Ok(()) + } + fn parse_program( + &self, + section: &Section, + program_section: ProgramSection, + name: String, + symbol: &Symbol, + ) -> Result<(Program, Function), ParseError> { let (func_info, line_info, func_info_rec_size, line_info_rec_size) = if let Some(btf_ext) = &self.btf_ext { let func_info = btf_ext.func_info.get(section.name); @@ -729,12 +769,15 @@ impl Object { (FuncSecInfo::default(), LineSecInfo::default(), 0, 0) }; + let start = symbol.address as usize; + let end = (symbol.address + symbol.size) as usize; + let function = Function { - name, - address: section.address, + name: name.to_owned(), + address: symbol.address, section_index: section.index, - section_offset: 0, - instructions: copy_instructions(section.data)?, + section_offset: start, + instructions: copy_instructions(§ion.data[start..end])?, func_info, line_info, func_info_rec_size, @@ -745,9 +788,9 @@ impl Object { Program { license: self.license.clone(), kernel_version: self.kernel_version, - section: prog_sec, - section_index: function.section_index.0, - address: function.address, + section: program_section.clone(), + section_index: section.index.0, + address: symbol.address, }, function, )) @@ -845,15 +888,23 @@ impl Object { Ok(()) } - fn parse_btf_maps( - &mut self, - section: &Section, - symbols: HashMap, - ) -> Result<(), ParseError> { + fn parse_btf_maps(&mut self, section: &Section) -> Result<(), ParseError> { if self.btf.is_none() { return Err(ParseError::NoBTF); } let btf = self.btf.as_ref().unwrap(); + let maps: HashMap<&String, usize> = self + .symbols_by_section + .get(§ion.index) + .ok_or(ParseError::NoSymbolsForSection { + section_name: section.name.to_owned(), + })? + .iter() + .filter_map(|s| { + let symbol = self.symbol_table.get(s).unwrap(); + symbol.name.as_ref().map(|name| (name, symbol.index)) + }) + .collect(); for t in btf.types() { if let BtfType::DataSec(datasec) = &t { @@ -865,18 +916,17 @@ impl Object { // each btf_var_secinfo contains a map for info in &datasec.entries { let (map_name, def) = parse_btf_map_def(btf, info)?; - let symbol_index = symbols - .get(&map_name) - .ok_or_else(|| ParseError::SymbolNotFound { - name: map_name.to_string(), - })? - .index; + let symbol_index = + maps.get(&map_name) + .ok_or_else(|| ParseError::SymbolNotFound { + name: map_name.to_string(), + })?; self.maps.insert( map_name, Map::Btf(BtfMap { def, section_index: section.index.0, - symbol_index, + symbol_index: *symbol_index, data: Vec::new(), }), ); @@ -887,6 +937,50 @@ impl Object { Ok(()) } + // Parses multiple map definition contained in a single `maps` section (which is + // different from `.maps` which is used for BTF). We can tell where each map is + // based on the symbol table. + fn parse_maps_section<'a, I: Iterator>( + &self, + maps: &mut HashMap, + section: &Section, + symbols: I, + ) -> Result<(), ParseError> { + let mut have_symbols = false; + // each symbol in the section is a separate map + for i in symbols { + let sym = self.symbol_table.get(i).ok_or(ParseError::SymbolNotFound { + name: i.to_string(), + })?; + let start = sym.address as usize; + let end = start + sym.size as usize; + let data = §ion.data[start..end]; + let name = sym + .name + .as_ref() + .ok_or(ParseError::MapSymbolNameNotFound { i: *i })?; + let def = parse_map_def(name, data)?; + maps.insert( + name.to_string(), + Map::Legacy(LegacyMap { + section_index: section.index.0, + section_kind: section.kind, + symbol_index: Some(sym.index), + def, + data: Vec::new(), + }), + ); + have_symbols = true; + } + if !have_symbols { + return Err(ParseError::NoSymbolsForSection { + section_name: section.name.to_owned(), + }); + } + + Ok(()) + } + fn parse_section(&mut self, section: Section) -> Result<(), ParseError> { self.section_infos .insert(section.name.to_owned(), (section.index, section.size)); @@ -898,22 +992,7 @@ impl Object { BpfSectionKind::Text => self.parse_text_section(section)?, BpfSectionKind::Btf => self.parse_btf(§ion)?, BpfSectionKind::BtfExt => self.parse_btf_ext(§ion)?, - BpfSectionKind::BtfMaps => { - let symbols: HashMap = self - .symbol_table - .values() - .filter(|s| { - if let Some(idx) = s.section_index { - idx == section.index.0 && s.name.is_some() - } else { - false - } - }) - .cloned() - .map(|s| (s.name.as_ref().unwrap().to_string(), s)) - .collect(); - self.parse_btf_maps(§ion, symbols)? - } + BpfSectionKind::BtfMaps => self.parse_btf_maps(§ion)?, BpfSectionKind::Maps => { // take out self.maps so we can borrow the iterator below // without cloning or collecting @@ -921,13 +1000,15 @@ impl Object { // extract the symbols for the .maps section, we'll need them // during parsing - let symbols = self.symbol_table.values().filter(|s| { - s.section_index - .map(|idx| idx == section.index.0) - .unwrap_or(false) - }); + let symbols = self + .symbols_by_section + .get(§ion.index) + .ok_or(ParseError::NoSymbolsForSection { + section_name: section.name.to_owned(), + })? + .iter(); - let res = parse_maps_section(&mut maps, §ion, symbols); + let res = self.parse_maps_section(&mut maps, §ion, symbols); // put the maps back self.maps = maps; @@ -935,10 +1016,7 @@ impl Object { res? } BpfSectionKind::Program => { - let (program, function) = self.parse_program(§ion)?; - self.functions.insert(program.function_key(), function); - self.programs - .insert(program.section.name().to_owned(), program); + self.parse_programs(§ion)?; if !section.relocations.is_empty() { self.relocations.insert( section.index, @@ -1003,45 +1081,6 @@ impl Function { } } -// Parses multiple map definition contained in a single `maps` section (which is -// different from `.maps` which is used for BTF). We can tell where each map is -// based on the symbol table. -fn parse_maps_section<'a, I: Iterator>( - maps: &mut HashMap, - section: &Section, - symbols: I, -) -> Result<(), ParseError> { - let mut have_symbols = false; - - // each symbol in the section is a separate map - for (i, sym) in symbols.enumerate() { - let start = sym.address as usize; - let end = start + sym.size as usize; - let data = §ion.data[start..end]; - let name = sym - .name - .as_ref() - .ok_or(ParseError::MapSymbolNameNotFound { i })?; - let def = parse_map_def(name, data)?; - maps.insert( - name.to_string(), - Map::Legacy(LegacyMap { - section_index: section.index.0, - section_kind: section.kind, - symbol_index: Some(sym.index), - def, - data: Vec::new(), - }), - ); - have_symbols = true; - } - if !have_symbols { - return Err(ParseError::NoSymbolsForMapsSection); - } - - Ok(()) -} - /// Errors caught during parsing the object file #[derive(Debug, thiserror::Error)] #[allow(missing_docs)] @@ -1105,8 +1144,8 @@ pub enum ParseError { #[error("the map number {i} in the `maps` section doesn't have a symbol name")] MapSymbolNameNotFound { i: usize }, - #[error("no symbols for `maps` section, can't parse maps")] - NoSymbolsForMapsSection, + #[error("no symbols found in the {section_name} section")] + NoSymbolsForSection { section_name: String }, /// No BTF parsed for object #[error("no BTF parsed for object")] @@ -1480,9 +1519,17 @@ mod tests { use super::*; use crate::maps::PinningType; - fn fake_section<'a>(kind: BpfSectionKind, name: &'a str, data: &'a [u8]) -> Section<'a> { + const FAKE_INS_LEN: u64 = 8; + + fn fake_section<'a>( + kind: BpfSectionKind, + name: &'a str, + data: &'a [u8], + index: Option, + ) -> Section<'a> { + let idx = index.unwrap_or(0); Section { - index: SectionIndex(0), + index: SectionIndex(idx), kind, address: 0, name, @@ -1513,9 +1560,13 @@ mod tests { address, size, is_definition: false, - kind: SymbolKind::Data, + kind: SymbolKind::Text, }, ); + obj.symbols_by_section + .entry(SectionIndex(section_index)) + .or_default() + .push(idx + 1); } fn bytes_of(val: &T) -> &[u8] { @@ -1637,6 +1688,7 @@ mod tests { BpfSectionKind::Data, ".bss", map_data, + None, ), ), Ok(Map::Legacy(LegacyMap { @@ -1670,6 +1722,7 @@ mod tests { &[ 159, 235, 1, 0, 24, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, ], + None, )) .unwrap(); obj.parse_section(fake_section( @@ -1679,6 +1732,7 @@ mod tests { 159, 235, 1, 0, 24, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 4, 0, 0, 0, 4, 0, 0, 0, 8, 0, 0, 0, 16, 0, 0, 0, ], + None, )) .unwrap(); @@ -1688,36 +1742,54 @@ mod tests { #[test] fn test_parse_program_error() { - let obj = fake_obj(); - + let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", 1); assert_matches!( - obj.parse_program(&fake_section( + obj.parse_programs(&fake_section( BpfSectionKind::Program, "kprobe/foo", &42u32.to_ne_bytes(), - )), + None, + ),), Err(ParseError::InvalidProgramCode) ); } #[test] fn test_parse_program() { - let obj = fake_obj(); + let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + + obj.parse_programs(&fake_section( + BpfSectionKind::Program, + "kprobe/foo", + bytes_of(&fake_ins()), + None, + )) + .unwrap(); + + let prog_foo = obj.programs.get("foo").unwrap(); assert_matches!( - obj.parse_program(&fake_section(BpfSectionKind::Program,"kprobe/foo", bytes_of(&fake_ins()))), - Ok((Program { + prog_foo, + Program { license, kernel_version: None, section: ProgramSection::KProbe { .. }, - .. }, Function { - name, - address: 0, - section_index: SectionIndex(0), - section_offset: 0, - instructions, - ..})) if license.to_string_lossy() == "GPL" && name == "foo" && instructions.len() == 1 + .. + } if license.to_str().unwrap() == "GPL" ); + + assert_matches!( + obj.functions.get(&prog_foo.function_key()), + Some(Function { + name, + address: 0, + section_index: SectionIndex(0), + section_offset: 0, + instructions, + ..}) if name == "foo" && instructions.len() == 1 + ) } #[test] @@ -1735,13 +1807,72 @@ mod tests { max_entries: 4, map_flags: 5, ..Default::default() - }) + }), + None, )), Ok(()) ); assert!(obj.maps.get("foo").is_some()); } + #[test] + fn test_parse_multiple_program_in_same_section() { + let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + fake_sym(&mut obj, 0, FAKE_INS_LEN, "bar", FAKE_INS_LEN); + + let insns = [fake_ins(), fake_ins()]; + let data = bytes_of(&insns); + + obj.parse_programs(&fake_section(BpfSectionKind::Program, "kprobe", data, None)) + .unwrap(); + + let prog_foo = obj.programs.get("foo").unwrap(); + let function_foo = obj.functions.get(&prog_foo.function_key()).unwrap(); + let prog_bar = obj.programs.get("bar").unwrap(); + let function_bar = obj.functions.get(&prog_bar.function_key()).unwrap(); + + assert_matches!(prog_foo, + Program { + license, + kernel_version: None, + section: ProgramSection::KProbe { .. }, + .. + } if license.to_string_lossy() == "GPL" + ); + assert_matches!( + function_foo, + Function { + name, + address: 0, + section_index: SectionIndex(0), + section_offset: 0, + instructions, + .. + } if name == "foo" && instructions.len() == 1 + ); + + assert_matches!(prog_bar, + Program { + license, + kernel_version: None, + section: ProgramSection::KProbe { .. }, + .. + } if license.to_string_lossy() == "GPL" + ); + assert_matches!( + function_bar, + Function { + name, + address: 8, + section_index: SectionIndex(0), + section_offset: 8, + instructions, + .. + } if name == "bar" && instructions.len() == 1 + ); + } + #[test] fn test_parse_section_multiple_maps() { let mut obj = fake_obj(); @@ -1764,7 +1895,12 @@ mod tests { buf.extend([0, 0, 0, 0]); buf.extend(&map_data); assert_matches!( - obj.parse_section(fake_section(BpfSectionKind::Maps, "maps", buf.as_slice(),)), + obj.parse_section(fake_section( + BpfSectionKind::Maps, + "maps", + buf.as_slice(), + None + )), Ok(()) ); assert!(obj.maps.get("foo").is_some()); @@ -1783,13 +1919,23 @@ mod tests { fn test_parse_section_data() { let mut obj = fake_obj(); assert_matches!( - obj.parse_section(fake_section(BpfSectionKind::Data, ".bss", b"map data")), + obj.parse_section(fake_section( + BpfSectionKind::Data, + ".bss", + b"map data", + None + )), Ok(()) ); assert!(obj.maps.get(".bss").is_some()); assert_matches!( - obj.parse_section(fake_section(BpfSectionKind::Data, ".rodata", b"map data")), + obj.parse_section(fake_section( + BpfSectionKind::Data, + ".rodata", + b"map data", + None + )), Ok(()) ); assert!(obj.maps.get(".rodata").is_some()); @@ -1798,20 +1944,31 @@ mod tests { obj.parse_section(fake_section( BpfSectionKind::Data, ".rodata.boo", - b"map data" + b"map data", + None )), Ok(()) ); assert!(obj.maps.get(".rodata.boo").is_some()); assert_matches!( - obj.parse_section(fake_section(BpfSectionKind::Data, ".data", b"map data")), + obj.parse_section(fake_section( + BpfSectionKind::Data, + ".data", + b"map data", + None + )), Ok(()) ); assert!(obj.maps.get(".data").is_some()); assert_matches!( - obj.parse_section(fake_section(BpfSectionKind::Data, ".data.boo", b"map data")), + obj.parse_section(fake_section( + BpfSectionKind::Data, + ".data.boo", + b"map data", + None + )), Ok(()) ); assert!(obj.maps.get(".data.boo").is_some()); @@ -1820,12 +1977,14 @@ mod tests { #[test] fn test_parse_section_kprobe() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "kprobe/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -1841,12 +2000,14 @@ mod tests { #[test] fn test_parse_section_uprobe() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "uprobe/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -1862,12 +2023,15 @@ mod tests { #[test] fn test_parse_section_trace_point() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + fake_sym(&mut obj, 1, 0, "bar", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "tracepoint/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -1883,12 +2047,13 @@ mod tests { obj.parse_section(fake_section( BpfSectionKind::Program, "tp/foo/bar", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + Some(1), )), Ok(()) ); assert_matches!( - obj.programs.get("foo/bar"), + obj.programs.get("bar"), Some(Program { section: ProgramSection::TracePoint { .. }, .. @@ -1899,12 +2064,14 @@ mod tests { #[test] fn test_parse_section_socket_filter() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "socket/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -1920,12 +2087,14 @@ mod tests { #[test] fn test_parse_section_xdp() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "xdp/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -1941,12 +2110,14 @@ mod tests { #[test] fn test_parse_section_xdp_frags() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "xdp.frags/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -1962,12 +2133,15 @@ mod tests { #[test] fn test_parse_section_raw_tp() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + fake_sym(&mut obj, 1, 0, "bar", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "raw_tp/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -1983,7 +2157,8 @@ mod tests { obj.parse_section(fake_section( BpfSectionKind::Program, "raw_tracepoint/bar", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + Some(1) )), Ok(()) ); @@ -1999,12 +2174,14 @@ mod tests { #[test] fn test_parse_section_lsm() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "lsm/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2023,12 +2200,14 @@ mod tests { #[test] fn test_parse_section_lsm_sleepable() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "lsm.s/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2047,12 +2226,14 @@ mod tests { #[test] fn test_parse_section_btf_tracepoint() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "tp_btf/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2068,12 +2249,14 @@ mod tests { #[test] fn test_parse_section_skskb_unnamed() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "stream_parser", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "sk_skb/stream_parser", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2089,12 +2272,14 @@ mod tests { #[test] fn test_parse_section_skskb_named() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "my_parser", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "sk_skb/stream_parser/my_parser", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2110,12 +2295,14 @@ mod tests { #[test] fn test_parse_section_fentry() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "fentry/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2131,12 +2318,14 @@ mod tests { #[test] fn test_parse_section_fexit() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "fexit/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2152,12 +2341,14 @@ mod tests { #[test] fn test_parse_section_cgroup_skb_ingress_unnamed() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "ingress", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "cgroup_skb/ingress", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2173,12 +2364,14 @@ mod tests { #[test] fn test_parse_section_cgroup_skb_ingress_named() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "cgroup_skb/ingress/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2194,12 +2387,14 @@ mod tests { #[test] fn test_parse_section_cgroup_skb_no_direction_unamed() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "skb", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "cgroup/skb", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2215,12 +2410,14 @@ mod tests { #[test] fn test_parse_section_cgroup_skb_no_direction_named() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "cgroup/skb/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2236,12 +2433,14 @@ mod tests { #[test] fn test_parse_section_sock_addr_named() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "cgroup/connect4/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2260,12 +2459,14 @@ mod tests { #[test] fn test_parse_section_sock_addr_unnamed() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "connect4", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "cgroup/connect4", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2284,12 +2485,14 @@ mod tests { #[test] fn test_parse_section_sockopt_named() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "cgroup/getsockopt/foo", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2308,12 +2511,14 @@ mod tests { #[test] fn test_parse_section_sockopt_unnamed() { let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "getsockopt", FAKE_INS_LEN); assert_matches!( obj.parse_section(fake_section( BpfSectionKind::Program, "cgroup/getsockopt", - bytes_of(&fake_ins()) + bytes_of(&fake_ins()), + None )), Ok(()) ); @@ -2455,10 +2660,10 @@ mod tests { 0x2E, 0x6D, 0x61, 0x70, 0x73, 0x00, 0x6C, 0x69, 0x63, 0x65, 0x6E, 0x73, 0x65, 0x00, ]; - let btf_section = fake_section(BpfSectionKind::Btf, ".BTF", data); + let btf_section = fake_section(BpfSectionKind::Btf, ".BTF", data, None); obj.parse_section(btf_section).unwrap(); - let map_section = fake_section(BpfSectionKind::BtfMaps, ".maps", &[]); + let map_section = fake_section(BpfSectionKind::BtfMaps, ".maps", &[], None); obj.parse_section(map_section).unwrap(); let map = obj.maps.get("map_1").unwrap(); diff --git a/test/integration-ebpf/src/name_test.rs b/test/integration-ebpf/src/name_test.rs index 1a5f870c..a9208e95 100644 --- a/test/integration-ebpf/src/name_test.rs +++ b/test/integration-ebpf/src/name_test.rs @@ -3,7 +3,7 @@ use aya_bpf::{bindings::xdp_action, macros::xdp, programs::XdpContext}; -#[xdp(name = "ihaveaverylongname")] +#[xdp] pub fn ihaveaverylongname(ctx: XdpContext) -> u32 { match unsafe { try_pass(ctx) } { Ok(ret) => ret, diff --git a/test/integration-test/bpf/ext.bpf.c b/test/integration-test/bpf/ext.bpf.c index 6c77127a..dc82f5cf 100644 --- a/test/integration-test/bpf/ext.bpf.c +++ b/test/integration-test/bpf/ext.bpf.c @@ -1,7 +1,7 @@ #include #include -SEC("xdp/drop") +SEC("xdp") int xdp_drop(struct xdp_md *ctx) { return XDP_DROP; diff --git a/test/integration-test/bpf/main.bpf.c b/test/integration-test/bpf/main.bpf.c index 79c82041..2a30577f 100644 --- a/test/integration-test/bpf/main.bpf.c +++ b/test/integration-test/bpf/main.bpf.c @@ -1,7 +1,7 @@ #include #include -SEC("xdp/pass") +SEC("xdp") int xdp_pass(struct xdp_md *ctx) { return XDP_PASS; diff --git a/test/integration-test/src/tests/load.rs b/test/integration-test/src/tests/load.rs index 8958ae59..4c366e46 100644 --- a/test/integration-test/src/tests/load.rs +++ b/test/integration-test/src/tests/load.rs @@ -36,7 +36,7 @@ fn multiple_btf_maps() { let map_1: Array<_, u64> = bpf.take_map("map_1").unwrap().try_into().unwrap(); let map_2: Array<_, u64> = bpf.take_map("map_2").unwrap().try_into().unwrap(); - let prog: &mut TracePoint = bpf.program_mut("tracepoint").unwrap().try_into().unwrap(); + let prog: &mut TracePoint = bpf.program_mut("bpf_prog").unwrap().try_into().unwrap(); prog.load().unwrap(); prog.attach("sched", "sched_switch").unwrap(); @@ -110,26 +110,26 @@ macro_rules! assert_loaded { #[test] fn unload_xdp() { let mut bpf = Bpf::load(crate::TEST).unwrap(); - let prog: &mut Xdp = bpf.program_mut("test_xdp").unwrap().try_into().unwrap(); + let prog: &mut Xdp = bpf.program_mut("pass").unwrap().try_into().unwrap(); prog.load().unwrap(); - assert_loaded!("test_xdp", true); + assert_loaded!("pass", true); let link = prog.attach("lo", XdpFlags::default()).unwrap(); { let _link_owned = prog.take_link(link).unwrap(); prog.unload().unwrap(); - assert_loaded_and_linked!("test_xdp", true); + assert_loaded_and_linked!("pass", true); }; - assert_loaded!("test_xdp", false); + assert_loaded!("pass", false); prog.load().unwrap(); - assert_loaded!("test_xdp", true); + assert_loaded!("pass", true); prog.attach("lo", XdpFlags::default()).unwrap(); - assert_loaded!("test_xdp", true); + assert_loaded!("pass", true); prog.unload().unwrap(); - assert_loaded!("test_xdp", false); + assert_loaded!("pass", false); } #[test] @@ -224,26 +224,26 @@ fn pin_link() { } let mut bpf = Bpf::load(crate::TEST).unwrap(); - let prog: &mut Xdp = bpf.program_mut("test_xdp").unwrap().try_into().unwrap(); + let prog: &mut Xdp = bpf.program_mut("pass").unwrap().try_into().unwrap(); prog.load().unwrap(); let link_id = prog.attach("lo", XdpFlags::default()).unwrap(); let link = prog.take_link(link_id).unwrap(); - assert_loaded!("test_xdp", true); + assert_loaded!("pass", true); let fd_link: FdLink = link.try_into().unwrap(); let pinned = fd_link.pin("/sys/fs/bpf/aya-xdp-test-lo").unwrap(); // because of the pin, the program is still attached prog.unload().unwrap(); - assert_loaded!("test_xdp", true); + assert_loaded!("pass", true); // delete the pin, but the program is still attached let new_link = pinned.unpin().unwrap(); - assert_loaded!("test_xdp", true); + assert_loaded!("pass", true); // finally when new_link is dropped we're detached drop(new_link); - assert_loaded!("test_xdp", false); + assert_loaded!("pass", false); } #[test] diff --git a/test/integration-test/src/tests/rbpf.rs b/test/integration-test/src/tests/rbpf.rs index 790b2fdc..c12bbce8 100644 --- a/test/integration-test/src/tests/rbpf.rs +++ b/test/integration-test/src/tests/rbpf.rs @@ -36,10 +36,10 @@ fn use_map_with_rbpf() { assert_eq!(object.programs.len(), 1); assert_matches::assert_matches!( - object.programs["tracepoint"].section, + object.programs["bpf_prog"].section, ProgramSection::TracePoint { .. } ); - assert_eq!(object.programs["tracepoint"].section.name(), "tracepoint"); + assert_eq!(object.programs["bpf_prog"].section.name(), "tracepoint"); // Initialize maps: // - fd: 0xCAFE00 or 0xCAFE01 (the 0xCAFE00 part is used to distinguish fds from indices), @@ -83,7 +83,7 @@ fn use_map_with_rbpf() { assert_eq!(object.programs.len(), 1); let instructions = &object .functions - .get(&object.programs["tracepoint"].function_key()) + .get(&object.programs["bpf_prog"].function_key()) .unwrap() .instructions; let data = unsafe { diff --git a/test/integration-test/src/tests/smoke.rs b/test/integration-test/src/tests/smoke.rs index 30a97a39..89f7d0a1 100644 --- a/test/integration-test/src/tests/smoke.rs +++ b/test/integration-test/src/tests/smoke.rs @@ -32,11 +32,14 @@ fn extension() { return; } let mut bpf = Bpf::load(crate::MAIN).unwrap(); - let pass: &mut Xdp = bpf.program_mut("pass").unwrap().try_into().unwrap(); + let pass: &mut Xdp = bpf.program_mut("xdp_pass").unwrap().try_into().unwrap(); pass.load().unwrap(); pass.attach("lo", XdpFlags::default()).unwrap(); - let mut bpf = BpfLoader::new().extension("drop").load(crate::EXT).unwrap(); - let drop_: &mut Extension = bpf.program_mut("drop").unwrap().try_into().unwrap(); + let mut bpf = BpfLoader::new() + .extension("xdp_drop") + .load(crate::EXT) + .unwrap(); + let drop_: &mut Extension = bpf.program_mut("xdp_drop").unwrap().try_into().unwrap(); drop_.load(pass.fd().unwrap(), "xdp_pass").unwrap(); } diff --git a/xtask/public-api/aya-obj.txt b/xtask/public-api/aya-obj.txt index 53fb9207..8144037f 100644 --- a/xtask/public-api/aya-obj.txt +++ b/xtask/public-api/aya-obj.txt @@ -5764,7 +5764,8 @@ pub aya_obj::obj::ParseError::MapSymbolNameNotFound::i: usize pub aya_obj::obj::ParseError::MissingLicenseNullTerminator pub aya_obj::obj::ParseError::MissingLicenseNullTerminator::data: alloc::vec::Vec pub aya_obj::obj::ParseError::NoBTF -pub aya_obj::obj::ParseError::NoSymbolsForMapsSection +pub aya_obj::obj::ParseError::NoSymbolsForSection +pub aya_obj::obj::ParseError::NoSymbolsForSection::section_name: alloc::string::String pub aya_obj::obj::ParseError::SectionError pub aya_obj::obj::ParseError::SectionError::error: object::read::Error pub aya_obj::obj::ParseError::SectionError::index: usize @@ -6511,7 +6512,8 @@ pub aya_obj::ParseError::MapSymbolNameNotFound::i: usize pub aya_obj::ParseError::MissingLicenseNullTerminator pub aya_obj::ParseError::MissingLicenseNullTerminator::data: alloc::vec::Vec pub aya_obj::ParseError::NoBTF -pub aya_obj::ParseError::NoSymbolsForMapsSection +pub aya_obj::ParseError::NoSymbolsForSection +pub aya_obj::ParseError::NoSymbolsForSection::section_name: alloc::string::String pub aya_obj::ParseError::SectionError pub aya_obj::ParseError::SectionError::error: object::read::Error pub aya_obj::ParseError::SectionError::index: usize From c72aab5f7b809b48273cd2fed4554ddff6b66bdf Mon Sep 17 00:00:00 2001 From: Dave Tucker Date: Tue, 4 Jul 2023 17:13:40 +0100 Subject: [PATCH 2/3] aya-bpf-macros: Refactor for ease of testing The aya-bpf-macros needed refactoring for: 1. Ease of testing 2. To be consistent with when we use K/V args vs. idents 3. To deprecate the use of `name` to change the exported name of a function - we now use the symbol table. Signed-off-by: Dave Tucker --- Cargo.toml | 1 + aya-bpf-macros/Cargo.toml | 1 + aya-bpf-macros/src/args.rs | 69 ++ aya-bpf-macros/src/btf_tracepoint.rs | 89 ++ aya-bpf-macros/src/cgroup_device.rs | 65 ++ aya-bpf-macros/src/cgroup_skb.rs | 210 +++++ aya-bpf-macros/src/cgroup_sock.rs | 150 ++++ aya-bpf-macros/src/cgroup_sock_addr.rs | 365 ++++++++ aya-bpf-macros/src/cgroup_sockopt.rs | 103 +++ aya-bpf-macros/src/cgroup_sysctl.rs | 65 ++ aya-bpf-macros/src/expand.rs | 1044 ----------------------- aya-bpf-macros/src/fentry.rs | 84 ++ aya-bpf-macros/src/fexit.rs | 84 ++ aya-bpf-macros/src/kprobe.rs | 213 +++++ aya-bpf-macros/src/lib.rs | 452 +++++----- aya-bpf-macros/src/lsm.rs | 85 ++ aya-bpf-macros/src/map.rs | 75 ++ aya-bpf-macros/src/perf_event.rs | 67 ++ aya-bpf-macros/src/raw_tracepoint.rs | 76 ++ aya-bpf-macros/src/sk_lookup.rs | 65 ++ aya-bpf-macros/src/sk_msg.rs | 65 ++ aya-bpf-macros/src/sk_skb.rs | 115 +++ aya-bpf-macros/src/sock_ops.rs | 65 ++ aya-bpf-macros/src/socket_filter.rs | 65 ++ aya-bpf-macros/src/tc.rs | 65 ++ aya-bpf-macros/src/tracepoint.rs | 84 ++ aya-bpf-macros/src/uprobe.rs | 232 +++++ aya-bpf-macros/src/xdp.rs | 104 +++ test/integration-ebpf/src/map_test.rs | 2 +- test/integration-ebpf/src/pass.rs | 2 +- test/integration-ebpf/src/test.rs | 2 +- test/integration-test/src/tests/rbpf.rs | 7 +- 32 files changed, 2931 insertions(+), 1240 deletions(-) create mode 100644 aya-bpf-macros/src/args.rs create mode 100644 aya-bpf-macros/src/btf_tracepoint.rs create mode 100644 aya-bpf-macros/src/cgroup_device.rs create mode 100644 aya-bpf-macros/src/cgroup_skb.rs create mode 100644 aya-bpf-macros/src/cgroup_sock.rs create mode 100644 aya-bpf-macros/src/cgroup_sock_addr.rs create mode 100644 aya-bpf-macros/src/cgroup_sockopt.rs create mode 100644 aya-bpf-macros/src/cgroup_sysctl.rs delete mode 100644 aya-bpf-macros/src/expand.rs create mode 100644 aya-bpf-macros/src/fentry.rs create mode 100644 aya-bpf-macros/src/fexit.rs create mode 100644 aya-bpf-macros/src/kprobe.rs create mode 100644 aya-bpf-macros/src/lsm.rs create mode 100644 aya-bpf-macros/src/map.rs create mode 100644 aya-bpf-macros/src/perf_event.rs create mode 100644 aya-bpf-macros/src/raw_tracepoint.rs create mode 100644 aya-bpf-macros/src/sk_lookup.rs create mode 100644 aya-bpf-macros/src/sk_msg.rs create mode 100644 aya-bpf-macros/src/sk_skb.rs create mode 100644 aya-bpf-macros/src/sock_ops.rs create mode 100644 aya-bpf-macros/src/socket_filter.rs create mode 100644 aya-bpf-macros/src/tc.rs create mode 100644 aya-bpf-macros/src/tracepoint.rs create mode 100644 aya-bpf-macros/src/uprobe.rs create mode 100644 aya-bpf-macros/src/xdp.rs diff --git a/Cargo.toml b/Cargo.toml index ce579306..2dfb461c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -76,6 +76,7 @@ num_enum = { version = "0.6", default-features = false } object = { version = "0.31", default-features = false } parking_lot = { version = "0.12.0", default-features = false } proc-macro2 = { version = "1", default-features = false } +proc-macro-error = { version = "1.0", default-features = false } public-api = { version = "0.31.2", default-features = false } quote = { version = "1", default-features = false } rbpf = { version = "0.2.0", default-features = false } diff --git a/aya-bpf-macros/Cargo.toml b/aya-bpf-macros/Cargo.toml index fd569341..f238b7fb 100644 --- a/aya-bpf-macros/Cargo.toml +++ b/aya-bpf-macros/Cargo.toml @@ -9,6 +9,7 @@ proc-macro = true [dependencies] proc-macro2 = { workspace = true } +proc-macro-error = { workspace = true } quote = { workspace = true } syn = { workspace = true, default-features = true, features = ["full"] } diff --git a/aya-bpf-macros/src/args.rs b/aya-bpf-macros/src/args.rs new file mode 100644 index 00000000..44a0c2cd --- /dev/null +++ b/aya-bpf-macros/src/args.rs @@ -0,0 +1,69 @@ +use syn::{ + parse::{Parse, ParseStream}, + punctuated::{Pair, Punctuated}, + token::Eq, + Error, Ident, LitStr, Result, Token, +}; + +pub(crate) struct NameValue { + name: Ident, + _eq: Eq, + value: LitStr, +} + +pub(crate) struct Args { + pub(crate) args: Vec, +} + +impl Parse for Args { + fn parse(input: ParseStream) -> Result { + let args = Punctuated::::parse_terminated_with(input, |input| { + Ok(NameValue { + name: input.parse()?, + _eq: input.parse()?, + value: input.parse()?, + }) + })? + .into_pairs() + .map(|pair| match pair { + Pair::Punctuated(name_val, _) => name_val, + Pair::End(name_val) => name_val, + }) + .collect(); + + Ok(Args { args }) + } +} + +pub(crate) fn pop_arg(args: &mut Args, name: &str) -> Option { + match args.args.iter().position(|arg| arg.name == name) { + Some(index) => Some(args.args.remove(index).value.value()), + None => None, + } +} + +pub(crate) fn pop_required_arg(args: &mut Args, name: &str) -> Result { + let value = match args.args.iter().position(|arg| arg.name == name) { + Some(index) => Some(args.args.remove(index).value.value()), + None => None, + }; + match value { + Some(value) => Ok(value), + None => Err(Error::new_spanned( + args.args.first().unwrap().name.clone(), + format!("missing required argument `{}`", name), + )), + } +} + +pub(crate) fn err_on_unknown_args(args: &Args) -> Result<()> { + if let Some(arg) = args.args.get(0) { + return Err(Error::new_spanned(&arg.name, "invalid argument")); + } + Ok(()) +} + +pub(crate) fn name_arg(args: &mut Args) -> Result> { + let name = pop_arg(args, "name"); + Ok(name) +} diff --git a/aya-bpf-macros/src/btf_tracepoint.rs b/aya-bpf-macros/src/btf_tracepoint.rs new file mode 100644 index 00000000..f05b561e --- /dev/null +++ b/aya-bpf-macros/src/btf_tracepoint.rs @@ -0,0 +1,89 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use quote::quote; +use syn::{Error, ItemFn, Result}; + +use crate::args::{err_on_unknown_args, pop_arg, pop_required_arg, Args}; + +pub(crate) struct BtfTracePoint { + item: ItemFn, + function: String, + sleepable: bool, +} + +impl BtfTracePoint { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + let mut args: Args = syn::parse2(attrs)?; + let item = syn::parse2(item)?; + let function = pop_required_arg(&mut args, "function")?; + let mut sleepable = false; + if let Some(s) = pop_arg(&mut args, "sleepable") { + if let Ok(m) = s.parse() { + sleepable = m + } else { + return Err(Error::new_spanned( + s, + "invalid value. should be 'true' or 'false'", + )); + } + } + err_on_unknown_args(&args)?; + Ok(BtfTracePoint { + item, + function, + sleepable, + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_prefix = if self.sleepable { "tp_btf.s" } else { "tp_btf" }; + let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into(); + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 { + let _ = #fn_name(::aya_bpf::programs::BtfTracePointContext::new(ctx)); + return 0; + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_btf_tracepoint() { + let prog = BtfTracePoint::parse( + parse_quote!(function = "some_func"), + parse_quote!( + fn foo(ctx: BtfTracePointContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote!( + #[no_mangle] + #[link_section = "tp_btf/some_func"] + fn foo(ctx: *mut ::core::ffi::c_void) -> i32 { + let _ = foo(::aya_bpf::programs::BtfTracePointContext::new(ctx)); + return 0; + + fn foo(ctx: BtfTracePointContext) -> i32 { + 0 + } + } + ); + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/cgroup_device.rs b/aya-bpf-macros/src/cgroup_device.rs new file mode 100644 index 00000000..6acfa537 --- /dev/null +++ b/aya-bpf-macros/src/cgroup_device.rs @@ -0,0 +1,65 @@ +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +pub(crate) struct CgroupDevice { + item: ItemFn, +} + +impl CgroupDevice { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute") + } + let item = syn::parse2(item)?; + Ok(CgroupDevice { item }) + } + + pub(crate) fn expand(&self) -> Result { + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = "cgroup/dev"] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_cgroup_dev_ctx) -> i32 { + return #fn_name(::aya_bpf::programs::DeviceContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_cgroup_device() { + let prog = CgroupDevice::parse( + parse_quote!(), + parse_quote!( + fn foo(ctx: DeviceContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/dev"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_cgroup_dev_ctx) -> i32 { + return foo(::aya_bpf::programs::DeviceContext::new(ctx)); + + fn foo(ctx: DeviceContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/cgroup_skb.rs b/aya-bpf-macros/src/cgroup_skb.rs new file mode 100644 index 00000000..94488382 --- /dev/null +++ b/aya-bpf-macros/src/cgroup_skb.rs @@ -0,0 +1,210 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{Ident, ItemFn, Result}; + +pub(crate) struct CgroupSkb { + item: ItemFn, + attach_type: Option, +} + +impl CgroupSkb { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + let item: ItemFn = syn::parse2(item)?; + let mut attach_type = None; + if !attrs.is_empty() { + let ident: Ident = syn::parse2(attrs)?; + match ident.to_string().as_str() { + "ingress" | "egress" => (), + _ => abort!(attach_type, "invalid attach type"), + } + attach_type = Some(ident.to_string()); + } + + Ok(CgroupSkb { item, attach_type }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = if self.attach_type.is_some() { + format!("cgroup_skb/{}", self.attach_type.as_ref().unwrap()).into() + } else { + "cgroup_skb".into() + }; + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn cgroup_skb() { + let prog = CgroupSkb::parse( + parse_quote!(), + parse_quote!( + fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup_skb"] + fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return foo(::aya_bpf::programs::SkBuffContext::new(ctx)); + + fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_skb_egress() { + let prog = CgroupSkb::parse( + parse_quote!(egress), + parse_quote!( + fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup_skb/egress"] + fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return foo(::aya_bpf::programs::SkBuffContext::new(ctx)); + + fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_skb_ingress() { + let prog = CgroupSkb::parse( + parse_quote!(ingress), + parse_quote!( + fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup_skb/ingress"] + fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return foo(::aya_bpf::programs::SkBuffContext::new(ctx)); + + fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn priv_function() { + let prog = CgroupSkb::parse( + parse_quote!(egress), + parse_quote!( + fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup_skb/egress"] + fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return foo(::aya_bpf::programs::SkBuffContext::new(ctx)); + + fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn pub_function() { + let prog = CgroupSkb::parse( + parse_quote!(egress), + parse_quote!( + pub fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup_skb/egress"] + pub fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return foo(::aya_bpf::programs::SkBuffContext::new(ctx)); + + pub fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn pub_crate_function() { + let prog = CgroupSkb::parse( + parse_quote!(egress), + parse_quote!( + pub(crate) fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup_skb/egress"] + pub(crate) fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return foo(::aya_bpf::programs::SkBuffContext::new(ctx)); + + pub(crate) fn foo(ctx: SkBuffContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/cgroup_sock.rs b/aya-bpf-macros/src/cgroup_sock.rs new file mode 100644 index 00000000..dfc24561 --- /dev/null +++ b/aya-bpf-macros/src/cgroup_sock.rs @@ -0,0 +1,150 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{Ident, ItemFn, Result}; + +pub(crate) struct CgroupSock { + item: ItemFn, + attach_type: String, +} + +impl CgroupSock { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + let item: ItemFn = syn::parse2(item)?; + let attach_type: Ident = syn::parse2(attrs)?; + match attach_type.to_string().as_str() { + "post_bind4" | "post_bind6" | "sock_create" | "sock_release" => (), + _ => abort!(attach_type, "invalid attach type"), + } + Ok(CgroupSock { + item, + attach_type: attach_type.to_string(), + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = format!("cgroup/{}", self.attach_type).into(); + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 { + return #fn_name(::aya_bpf::programs::SockContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn cgroup_sock_post_bind4() { + let prog = CgroupSock::parse( + parse_quote!(post_bind4), + parse_quote!( + fn foo(ctx: SockContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/post_bind4"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 { + return foo(::aya_bpf::programs::SockContext::new(ctx)); + + fn foo(ctx: SockContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_post_bind6() { + let prog = CgroupSock::parse( + parse_quote!(post_bind6), + parse_quote!( + fn foo(ctx: SockContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/post_bind6"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 { + return foo(::aya_bpf::programs::SockContext::new(ctx)); + + fn foo(ctx: SockContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + #[test] + fn cgroup_sock_sock_create() { + let prog = CgroupSock::parse( + parse_quote!(sock_create), + parse_quote!( + fn foo(ctx: SockContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/sock_create"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 { + return foo(::aya_bpf::programs::SockContext::new(ctx)); + + fn foo(ctx: SockContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + #[test] + fn cgroup_sock_sock_release() { + let prog = CgroupSock::parse( + parse_quote!(sock_release), + parse_quote!( + fn foo(ctx: SockContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/sock_release"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 { + return foo(::aya_bpf::programs::SockContext::new(ctx)); + + fn foo(ctx: SockContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/cgroup_sock_addr.rs b/aya-bpf-macros/src/cgroup_sock_addr.rs new file mode 100644 index 00000000..6ca0f20e --- /dev/null +++ b/aya-bpf-macros/src/cgroup_sock_addr.rs @@ -0,0 +1,365 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{Ident, ItemFn, Result}; + +pub(crate) struct CgroupSockAddr { + item: ItemFn, + attach_type: String, +} + +impl CgroupSockAddr { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if attrs.is_empty() { + abort!(attrs, "missing attach type") + } + let attach_type: Ident = syn::parse2(attrs)?; + match attach_type.to_string().as_str() { + "connect4" | "connect6" | "bind4" | "bind6" | "getpeername4" | "getpeername6" + | "getsockname4" | "getsockname6" | "sendmsg4" | "sendmsg6" | "recvmsg4" + | "recvmsg6" => (), + _ => abort!(attach_type, "invalid attach type"), + } + let item = syn::parse2(item)?; + Ok(CgroupSockAddr { + item, + attach_type: attach_type.to_string(), + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = format!("cgroup/{}", self.attach_type).into(); + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return #fn_name(::aya_bpf::programs::SockAddrContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn cgroup_sock_addr_connect4() { + let prog = CgroupSockAddr::parse( + parse_quote!(connect4), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/connect4"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_connect6() { + let prog = CgroupSockAddr::parse( + parse_quote!(connect6), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/connect6"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_bind4() { + let prog = CgroupSockAddr::parse( + parse_quote!(bind4), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/bind4"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_bind6() { + let prog = CgroupSockAddr::parse( + parse_quote!(bind6), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/bind6"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_getpeername4() { + let prog = CgroupSockAddr::parse( + parse_quote!(getpeername4), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/getpeername4"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_getpeername6() { + let prog = CgroupSockAddr::parse( + parse_quote!(getpeername6), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/getpeername6"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_getsockname4() { + let prog = CgroupSockAddr::parse( + parse_quote!(getsockname4), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/getsockname4"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_getsockname6() { + let prog = CgroupSockAddr::parse( + parse_quote!(getsockname6), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/getsockname6"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_sendmsg4() { + let prog = CgroupSockAddr::parse( + parse_quote!(sendmsg4), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/sendmsg4"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_sendmsg6() { + let prog = CgroupSockAddr::parse( + parse_quote!(sendmsg6), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/sendmsg6"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_recvmsg4() { + let prog = CgroupSockAddr::parse( + parse_quote!(recvmsg4), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/recvmsg4"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sock_addr_recvmsg6() { + let prog = CgroupSockAddr::parse( + parse_quote!(recvmsg6), + parse_quote!( + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/recvmsg6"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { + return foo(::aya_bpf::programs::SockAddrContext::new(ctx)); + + fn foo(ctx: CgroupSockAddrContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/cgroup_sockopt.rs b/aya-bpf-macros/src/cgroup_sockopt.rs new file mode 100644 index 00000000..11440f12 --- /dev/null +++ b/aya-bpf-macros/src/cgroup_sockopt.rs @@ -0,0 +1,103 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{Ident, ItemFn, Result}; + +pub(crate) struct CgroupSockopt { + item: ItemFn, + attach_type: String, +} + +impl CgroupSockopt { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if attrs.is_empty() { + abort!(attrs, "expected attach type"); + } + let attach_type: Ident = syn::parse2(attrs)?; + match attach_type.to_string().as_str() { + "getsockopt" | "setsockopt" => (), + _ => abort!(attach_type, "invalid attach type"), + } + let item = syn::parse2(item)?; + Ok(CgroupSockopt { + item, + attach_type: attach_type.to_string(), + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = format!("cgroup/{}", self.attach_type).into(); + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sockopt) -> i32 { + return #fn_name(::aya_bpf::programs::SockoptContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn cgroup_sockopt_getsockopt() { + let prog = CgroupSockopt::parse( + parse_quote!(getsockopt), + parse_quote!( + fn foo(ctx: SockoptContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote!( + #[no_mangle] + #[link_section = "cgroup/getsockopt"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sockopt) -> i32 { + return foo(::aya_bpf::programs::SockoptContext::new(ctx)); + + fn foo(ctx: SockoptContext) -> i32 { + 0 + } + } + ); + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn cgroup_sockopt_setsockopt() { + let prog = CgroupSockopt::parse( + parse_quote!(setsockopt), + parse_quote!( + fn foo(ctx: SockoptContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote!( + #[no_mangle] + #[link_section = "cgroup/setsockopt"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sockopt) -> i32 { + return foo(::aya_bpf::programs::SockoptContext::new(ctx)); + + fn foo(ctx: SockoptContext) -> i32 { + 0 + } + } + ); + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/cgroup_sysctl.rs b/aya-bpf-macros/src/cgroup_sysctl.rs new file mode 100644 index 00000000..e3b11223 --- /dev/null +++ b/aya-bpf-macros/src/cgroup_sysctl.rs @@ -0,0 +1,65 @@ +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +pub(crate) struct CgroupSysctl { + item: ItemFn, +} + +impl CgroupSysctl { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute") + } + let item = syn::parse2(item)?; + Ok(CgroupSysctl { item }) + } + + pub(crate) fn expand(&self) -> Result { + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = "cgroup/sysctl"] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sysctl) -> i32 { + return #fn_name(::aya_bpf::programs::SysctlContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_cgroup_sysctl() { + let prog = CgroupSysctl::parse( + parse_quote!(), + parse_quote!( + fn foo(ctx: SysctlContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "cgroup/sysctl"] + fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sysctl) -> i32 { + return foo(::aya_bpf::programs::SysctlContext::new(ctx)); + + fn foo(ctx: SysctlContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/expand.rs b/aya-bpf-macros/src/expand.rs deleted file mode 100644 index c419172f..00000000 --- a/aya-bpf-macros/src/expand.rs +++ /dev/null @@ -1,1044 +0,0 @@ -use proc_macro2::TokenStream; -use quote::quote; -use syn::{ - parse::{Parse, ParseStream}, - punctuated::{Pair, Punctuated}, - token::Eq, - Error, Ident, ItemFn, ItemStatic, LitStr, Result, Token, -}; - -pub struct NameValue { - name: Ident, - _eq: Eq, - value: LitStr, -} - -pub struct Args { - args: Vec, -} - -impl Parse for Args { - fn parse(input: ParseStream) -> Result { - let args = Punctuated::::parse_terminated_with(input, |input| { - Ok(NameValue { - name: input.parse()?, - _eq: input.parse()?, - value: input.parse()?, - }) - })? - .into_pairs() - .map(|pair| match pair { - Pair::Punctuated(name_val, _) => name_val, - Pair::End(name_val) => name_val, - }) - .collect(); - - Ok(Args { args }) - } -} - -pub struct SockAddrArgs { - pub(crate) attach_type: Ident, - pub(crate) args: Args, -} - -impl Parse for SockAddrArgs { - fn parse(input: ParseStream) -> Result { - let attach_type: Ident = input.parse()?; - match attach_type.to_string().as_str() { - "connect4" | "connect6" | "bind4" | "bind6" | "getpeername4" | "getpeername6" - | "getsockname4" | "getsockname6" | "sendmsg4" | "sendmsg6" | "recvmsg4" - | "recvmsg6" => (), - _ => return Err(input.error("invalid attach type")), - } - let args = if input.parse::().is_ok() { - Args::parse(input)? - } else { - Args { args: vec![] } - }; - Ok(SockAddrArgs { attach_type, args }) - } -} - -pub struct SockoptArgs { - pub(crate) attach_type: Ident, - pub(crate) args: Args, -} - -impl Parse for SockoptArgs { - fn parse(input: ParseStream) -> Result { - let attach_type: Ident = input.parse()?; - match attach_type.to_string().as_str() { - "getsockopt" | "setsockopt" => (), - _ => return Err(input.error("invalid attach type")), - } - let args = if input.parse::().is_ok() { - Args::parse(input)? - } else { - Args { args: vec![] } - }; - Ok(SockoptArgs { attach_type, args }) - } -} - -pub struct Map { - item: ItemStatic, - name: String, -} - -impl Map { - pub fn from_syn(mut args: Args, item: ItemStatic) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.ident.to_string()); - Ok(Map { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = "maps"; - let name = &self.name; - let item = &self.item; - Ok(quote! { - #[link_section = #section_name] - #[export_name = #name] - #item - }) - } -} - -pub struct Probe { - kind: ProbeKind, - item: ItemFn, - name: String, -} - -impl Probe { - pub fn from_syn(kind: ProbeKind, mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); - - Ok(Probe { kind, item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = format!("{}/{}", self.kind, self.name); - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { - let _ = #fn_name(::aya_bpf::programs::ProbeContext::new(ctx)); - return 0; - - #item - } - }) - } -} - -pub struct SockOps { - item: ItemFn, - name: Option, -} - -impl SockOps { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?; - - Ok(SockOps { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name: std::borrow::Cow<'_, _> = if let Some(name) = &self.name { - format!("sockops/{name}").into() - } else { - "sockops".into() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock_ops) -> u32 { - return #fn_name(::aya_bpf::programs::SockOpsContext::new(ctx)); - - #item - } - }) - } -} - -pub struct SkMsg { - item: ItemFn, - name: String, -} - -impl SkMsg { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); - - Ok(SkMsg { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = format!("sk_msg/{}", self.name); - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::sk_msg_md) -> u32 { - return #fn_name(::aya_bpf::programs::SkMsgContext::new(ctx)); - - #item - } - }) - } -} - -pub struct Xdp { - item: ItemFn, - name: Option, - frags: bool, -} - -impl Xdp { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = pop_arg(&mut args, "name"); - let mut frags = false; - if let Some(s) = pop_arg(&mut args, "frags") { - if let Ok(m) = s.parse() { - frags = m - } else { - return Err(Error::new_spanned( - s, - "invalid value. should be 'true' or 'false'", - )); - } - } - err_on_unknown_args(&args)?; - Ok(Xdp { item, name, frags }) - } - - pub fn expand(&self) -> Result { - let section_prefix = if self.frags { "xdp.frags" } else { "xdp" }; - let section_name = if let Some(name) = &self.name { - format!("{section_prefix}/{name}") - } else { - section_prefix.to_string() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 { - return #fn_name(::aya_bpf::programs::XdpContext::new(ctx)); - - #item - } - }) - } -} - -pub struct SchedClassifier { - item: ItemFn, - name: Option, -} - -impl SchedClassifier { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?; - - Ok(SchedClassifier { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name: std::borrow::Cow<'_, _> = if let Some(name) = &self.name { - format!("classifier/{name}").into() - } else { - "classifier".into() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { - return #fn_name(::aya_bpf::programs::TcContext::new(ctx)); - - #item - } - }) - } -} - -pub struct CgroupSysctl { - item: ItemFn, - name: Option, -} - -impl CgroupSysctl { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?; - - Ok(CgroupSysctl { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = if let Some(name) = &self.name { - format!("cgroup/sysctl/{name}") - } else { - ("cgroup/sysctl").to_owned() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sysctl) -> i32 { - return #fn_name(::aya_bpf::programs::SysctlContext::new(ctx)); - - #item - } - }) - } -} - -pub struct CgroupSockopt { - item: ItemFn, - attach_type: String, - name: Option, -} - -impl CgroupSockopt { - pub fn from_syn(mut args: Args, item: ItemFn, attach_type: String) -> Result { - let name = pop_arg(&mut args, "name"); - err_on_unknown_args(&args)?; - - Ok(CgroupSockopt { - item, - attach_type, - name, - }) - } - pub fn expand(&self) -> Result { - let section_name = if let Some(name) = &self.name { - format!("cgroup/{}/{}", self.attach_type, name) - } else { - format!("cgroup/{}", self.attach_type) - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sockopt) -> i32 { - return #fn_name(::aya_bpf::programs::SockoptContext::new(ctx)); - - #item - } - }) - } -} - -pub struct CgroupSkb { - item: ItemFn, - expected_attach_type: Option, - name: Option, -} - -impl CgroupSkb { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = pop_arg(&mut args, "name"); - let expected_attach_type = pop_arg(&mut args, "attach"); - err_on_unknown_args(&args)?; - - Ok(CgroupSkb { - item, - expected_attach_type, - name, - }) - } - - pub fn expand(&self) -> Result { - let section_name = if let Some(attach) = &self.expected_attach_type { - if let Some(name) = &self.name { - format!("cgroup_skb/{attach}/{name}") - } else { - format!("cgroup_skb/{attach}") - } - } else if let Some(name) = &self.name { - format!("cgroup/skb/{name}") - } else { - ("cgroup/skb").to_owned() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { - return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx)); - - #item - } - }) - } -} - -pub struct CgroupSockAddr { - item: ItemFn, - attach_type: String, - name: Option, -} - -impl CgroupSockAddr { - pub fn from_syn(mut args: Args, item: ItemFn, attach_type: String) -> Result { - let name = pop_arg(&mut args, "name"); - err_on_unknown_args(&args)?; - - Ok(CgroupSockAddr { - item, - attach_type, - name, - }) - } - - pub fn expand(&self) -> Result { - let section_name = if let Some(name) = &self.name { - format!("cgroup/{}/{}", self.attach_type, name) - } else { - format!("cgroup/{}", self.attach_type) - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 { - return #fn_name(::aya_bpf::programs::SockAddrContext::new(ctx)); - - #item - } - }) - } -} - -pub struct CgroupSock { - item: ItemFn, - attach_type: Option, - name: Option, -} - -impl CgroupSock { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = pop_arg(&mut args, "name"); - let attach_type = pop_arg(&mut args, "attach"); - err_on_unknown_args(&args)?; - - Ok(CgroupSock { - item, - attach_type, - name, - }) - } - - pub fn expand(&self) -> Result { - let section_name: std::borrow::Cow<'_, _> = if let Some(name) = &self.name { - if let Some(attach_type) = &self.attach_type { - format!("cgroup/{attach_type}/{name}").into() - } else { - format!("cgroup/sock/{name}").into() - } - } else if let Some(attach_type) = &self.attach_type { - format!("cgroup/{attach_type}").into() - } else { - "cgroup/sock".into() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 { - return #fn_name(::aya_bpf::programs::SockContext::new(ctx)); - - #item - } - }) - } -} - -fn pop_arg(args: &mut Args, name: &str) -> Option { - match args.args.iter().position(|arg| arg.name == name) { - Some(index) => Some(args.args.remove(index).value.value()), - None => None, - } -} - -fn err_on_unknown_args(args: &Args) -> Result<()> { - if let Some(arg) = args.args.get(0) { - return Err(Error::new_spanned(&arg.name, "invalid argument")); - } - - Ok(()) -} - -fn name_arg(args: &mut Args) -> Result> { - let name = pop_arg(args, "name"); - err_on_unknown_args(args)?; - - Ok(name) -} - -#[allow(clippy::enum_variant_names)] -#[derive(Debug, Copy, Clone)] -pub enum ProbeKind { - KProbe, - KRetProbe, - UProbe, - URetProbe, -} - -impl std::fmt::Display for ProbeKind { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - use ProbeKind::*; - match self { - KProbe => write!(f, "kprobe"), - KRetProbe => write!(f, "kretprobe"), - UProbe => write!(f, "uprobe"), - URetProbe => write!(f, "uretprobe"), - } - } -} - -pub struct TracePoint { - item: ItemFn, - name: String, -} - -impl TracePoint { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); - - Ok(TracePoint { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = format!("tp/{}", self.name); - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { - let _ = #fn_name(::aya_bpf::programs::TracePointContext::new(ctx)); - return 0; - - #item - } - }) - } -} - -pub struct PerfEvent { - item: ItemFn, - name: String, -} - -impl PerfEvent { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); - - Ok(PerfEvent { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = format!("perf_event/{}", self.name); - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { - let _ = #fn_name(::aya_bpf::programs::PerfEventContext::new(ctx)); - return 0; - - #item - } - }) - } -} - -pub struct RawTracePoint { - item: ItemFn, - name: String, -} - -impl RawTracePoint { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); - - Ok(RawTracePoint { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = format!("raw_tp/{}", self.name); - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { - let _ = #fn_name(::aya_bpf::programs::RawTracePointContext::new(ctx)); - return 0; - - #item - } - }) - } -} - -pub struct Lsm { - item: ItemFn, - name: Option, - sleepable: bool, -} - -impl Lsm { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = pop_arg(&mut args, "name"); - let mut sleepable = false; - if let Some(s) = pop_arg(&mut args, "sleepable") { - if let Ok(m) = s.parse() { - sleepable = m - } else { - return Err(Error::new_spanned( - s, - "invalid value. should be 'true' or 'false'", - )); - } - } - err_on_unknown_args(&args)?; - Ok(Lsm { - item, - name, - sleepable, - }) - } - - pub fn expand(&self) -> Result { - let section_prefix = if self.sleepable { "lsm.s" } else { "lsm" }; - let section_name = if let Some(name) = &self.name { - format!("{section_prefix}/{name}") - } else { - section_prefix.to_string() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - // LSM probes need to return an integer corresponding to the correct - // policy decision. Therefore we do not simply default to a return value - // of 0 as in other program types. - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 { - return #fn_name(::aya_bpf::programs::LsmContext::new(ctx)); - - #item - } - }) - } -} - -pub struct BtfTracePoint { - item: ItemFn, - name: String, -} - -impl BtfTracePoint { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); - - Ok(BtfTracePoint { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = format!("tp_btf/{}", self.name); - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 { - let _ = #fn_name(::aya_bpf::programs::BtfTracePointContext::new(ctx)); - return 0; - - #item - } - }) - } -} - -#[allow(clippy::enum_variant_names)] -#[derive(Debug, Copy, Clone)] -pub enum SkSkbKind { - StreamVerdict, - StreamParser, -} - -impl std::fmt::Display for SkSkbKind { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - use SkSkbKind::*; - match self { - StreamVerdict => write!(f, "stream_verdict"), - StreamParser => write!(f, "stream_parser"), - } - } -} - -pub struct SkSkb { - kind: SkSkbKind, - item: ItemFn, - name: Option, -} - -impl SkSkb { - pub fn from_syn(kind: SkSkbKind, mut args: Args, item: ItemFn) -> Result { - let name = pop_arg(&mut args, "name"); - err_on_unknown_args(&args)?; - - Ok(SkSkb { item, kind, name }) - } - - pub fn expand(&self) -> Result { - let kind = &self.kind; - let section_name = if let Some(name) = &self.name { - format!("sk_skb/{kind}/{name}") - } else { - format!("sk_skb/{kind}") - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 { - return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx)); - - #item - } - }) - } -} - -pub struct SocketFilter { - item: ItemFn, - name: Option, -} - -impl SocketFilter { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?; - err_on_unknown_args(&args)?; - - Ok(SocketFilter { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name: std::borrow::Cow<'_, _> = if let Some(name) = &self.name { - format!("socket/{name}").into() - } else { - "socket".into() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i64 { - return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx)); - - #item - } - }) - } -} - -pub struct FEntry { - item: ItemFn, - name: String, -} - -impl FEntry { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); - - Ok(FEntry { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = format!("fentry/{}", self.name); - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 { - let _ = #fn_name(::aya_bpf::programs::FEntryContext::new(ctx)); - return 0; - - #item - } - }) - } -} - -pub struct FExit { - item: ItemFn, - name: String, -} - -impl FExit { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); - - Ok(FExit { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name = format!("fexit/{}", self.name); - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 { - let _ = #fn_name(::aya_bpf::programs::FExitContext::new(ctx)); - return 0; - - #item - } - }) - } -} - -pub struct SkLookup { - item: ItemFn, - name: Option, -} - -impl SkLookup { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?; - - Ok(SkLookup { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name: std::borrow::Cow<'_, _> = if let Some(name) = &self.name { - format!("sk_lookup/{name}").into() - } else { - "sk_lookup".into() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sk_lookup) -> u32 { - return #fn_name(::aya_bpf::programs::SkLookupContext::new(ctx)); - - #item - } - }) - } -} - -pub struct CgroupDevice { - item: ItemFn, - name: Option, -} - -impl CgroupDevice { - pub fn from_syn(mut args: Args, item: ItemFn) -> Result { - let name = name_arg(&mut args)?; - - Ok(CgroupDevice { item, name }) - } - - pub fn expand(&self) -> Result { - let section_name: std::borrow::Cow<'_, _> = if let Some(name) = &self.name { - format!("cgroup/dev/{name}").into() - } else { - "cgroup/dev".into() - }; - let fn_vis = &self.item.vis; - let fn_name = &self.item.sig.ident; - let item = &self.item; - Ok(quote! { - #[no_mangle] - #[link_section = #section_name] - #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_cgroup_dev_ctx) -> i32 { - return #fn_name(::aya_bpf::programs::DeviceContext::new(ctx)); - - #item - } - }) - } -} -#[cfg(test)] -mod tests { - use syn::parse_quote; - - use super::*; - - #[test] - fn cgroup_skb_with_attach_and_name() { - let prog = CgroupSkb::from_syn( - parse_quote!(name = "foo", attach = "ingress"), - parse_quote!( - fn foo(ctx: SkBuffContext) -> i32 { - 0 - } - ), - ) - .unwrap(); - let stream = prog.expand().unwrap(); - assert!(stream - .to_string() - .contains("[link_section = \"cgroup_skb/ingress/foo\"]")); - } - - #[test] - fn cgroup_skb_with_name() { - let prog = CgroupSkb::from_syn( - parse_quote!(name = "foo"), - parse_quote!( - fn foo(ctx: SkBuffContext) -> i32 { - 0 - } - ), - ) - .unwrap(); - let stream = prog.expand().unwrap(); - assert!(stream - .to_string() - .contains("[link_section = \"cgroup/skb/foo\"]")); - } - - #[test] - fn cgroup_skb_no_name() { - let prog = CgroupSkb::from_syn( - parse_quote!(), - parse_quote!( - fn foo(ctx: SkBuffContext) -> i32 { - 0 - } - ), - ) - .unwrap(); - let stream = prog.expand().unwrap(); - assert!(stream - .to_string() - .contains("[link_section = \"cgroup/skb\"]")); - } - - #[test] - fn cgroup_skb_with_attach_no_name() { - let prog = CgroupSkb::from_syn( - parse_quote!(attach = "egress"), - parse_quote!( - fn foo(ctx: SkBuffContext) -> i32 { - 0 - } - ), - ) - .unwrap(); - let stream = prog.expand().unwrap(); - assert!(stream - .to_string() - .contains("[link_section = \"cgroup_skb/egress\"]")); - } - - #[test] - fn cgroup_device_no_name() { - let prog = CgroupDevice::from_syn( - parse_quote!(), - parse_quote!( - fn foo(ctx: DeviceContext) -> i32 { - 0 - } - ), - ) - .unwrap(); - let stream = prog.expand().unwrap(); - assert!(stream - .to_string() - .contains("[link_section = \"cgroup/dev\"]")); - } - - #[test] - fn priv_function() { - let prog = CgroupSkb::from_syn( - parse_quote!(attach = "egress"), - parse_quote!( - fn foo(ctx: SkBuffContext) -> i32 { - 0 - } - ), - ) - .unwrap(); - let stream = prog.expand().unwrap(); - assert!(stream.to_string().contains("] fn foo (")); - } - - #[test] - fn pub_function() { - let prog = CgroupSkb::from_syn( - parse_quote!(attach = "egress"), - parse_quote!( - pub fn foo(ctx: SkBuffContext) -> i32 { - 0 - } - ), - ) - .unwrap(); - let stream = prog.expand().unwrap(); - assert!(stream.to_string().contains("] pub fn foo (")); - } - - #[test] - fn pub_crate_function() { - let prog = CgroupSkb::from_syn( - parse_quote!(attach = "egress"), - parse_quote!( - pub(crate) fn foo(ctx: SkBuffContext) -> i32 { - 0 - } - ), - ) - .unwrap(); - let stream = prog.expand().unwrap(); - assert!(stream.to_string().contains("] pub (crate) fn foo (")); - } -} diff --git a/aya-bpf-macros/src/fentry.rs b/aya-bpf-macros/src/fentry.rs new file mode 100644 index 00000000..d63dfd8b --- /dev/null +++ b/aya-bpf-macros/src/fentry.rs @@ -0,0 +1,84 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +use crate::args::{err_on_unknown_args, pop_required_arg}; + +pub(crate) struct FEntry { + item: ItemFn, + function: String, + sleepable: bool, +} + +impl FEntry { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if attrs.is_empty() { + abort!(attrs, "missing function name"); + } + let mut args = syn::parse2(attrs)?; + let item = syn::parse2(item)?; + let function = pop_required_arg(&mut args, "function")?; + err_on_unknown_args(&args)?; + Ok(FEntry { + item, + function, + sleepable: false, + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_prefix = if self.sleepable { "fentry.s" } else { "fentry" }; + let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into(); + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 { + let _ = #fn_name(::aya_bpf::programs::FEntryContext::new(ctx)); + return 0; + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_fentry() { + let prog = FEntry::parse( + parse_quote! { + function = "sys_clone" + }, + parse_quote! { + fn sys_clone(ctx: &mut aya_bpf::programs::FEntryContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "fentry/sys_clone"] + fn sys_clone(ctx: *mut ::core::ffi::c_void) -> i32 { + let _ = sys_clone(::aya_bpf::programs::FEntryContext::new(ctx)); + return 0; + + fn sys_clone(ctx: &mut aya_bpf::programs::FEntryContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/fexit.rs b/aya-bpf-macros/src/fexit.rs new file mode 100644 index 00000000..b81be64a --- /dev/null +++ b/aya-bpf-macros/src/fexit.rs @@ -0,0 +1,84 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +use crate::args::{err_on_unknown_args, pop_required_arg}; + +pub(crate) struct FExit { + item: ItemFn, + function: String, + sleepable: bool, +} + +impl FExit { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if attrs.is_empty() { + abort!(attrs, "missing function name"); + } + let mut args = syn::parse2(attrs)?; + let item = syn::parse2(item)?; + let function = pop_required_arg(&mut args, "function")?; + err_on_unknown_args(&args)?; + Ok(FExit { + item, + function, + sleepable: false, + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_prefix = if self.sleepable { "fexit.s" } else { "fexit" }; + let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into(); + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 { + let _ = #fn_name(::aya_bpf::programs::FExitContext::new(ctx)); + return 0; + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_fexit() { + let prog = FExit::parse( + parse_quote! { + function = "sys_clone" + }, + parse_quote! { + fn sys_clone(ctx: &mut FExitContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "fexit/sys_clone"] + fn sys_clone(ctx: *mut ::core::ffi::c_void) -> i32 { + let _ = sys_clone(::aya_bpf::programs::FExitContext::new(ctx)); + return 0; + + fn sys_clone(ctx: &mut FExitContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/kprobe.rs b/aya-bpf-macros/src/kprobe.rs new file mode 100644 index 00000000..5393187d --- /dev/null +++ b/aya-bpf-macros/src/kprobe.rs @@ -0,0 +1,213 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; + +use quote::quote; +use syn::{ItemFn, Result}; + +use crate::args::{err_on_unknown_args, pop_arg}; + +#[allow(clippy::enum_variant_names)] +#[derive(Debug, Copy, Clone)] +pub(crate) enum KProbeKind { + KProbe, + KRetProbe, +} + +impl std::fmt::Display for KProbeKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use KProbeKind::*; + match self { + KProbe => write!(f, "kprobe"), + KRetProbe => write!(f, "kretprobe"), + } + } +} + +pub(crate) struct KProbe { + kind: KProbeKind, + function: Option, + offset: Option, + item: ItemFn, +} + +impl KProbe { + pub(crate) fn parse(kind: KProbeKind, attrs: TokenStream, item: TokenStream) -> Result { + let mut function = None; + let mut offset = None; + + if !attrs.is_empty() { + let mut args = syn::parse2(attrs)?; + function = pop_arg(&mut args, "function"); + offset = pop_arg(&mut args, "offset").map(|v| v.parse::().unwrap()); + err_on_unknown_args(&args)?; + } + let item = syn::parse2(item)?; + Ok(KProbe { + kind, + item, + function, + offset, + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = if self.function.is_some() && self.offset.is_some() { + format!( + "{}/{}+{}", + self.kind, + self.function.as_ref().unwrap(), + self.offset.unwrap() + ) + .into() + } else if self.function.is_some() { + format!("{}/{}", self.kind, self.function.as_ref().unwrap()).into() + } else { + format!("{}", self.kind).into() + }; + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = #fn_name(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_kprobe() { + let kprobe = KProbe::parse( + KProbeKind::KProbe, + parse_quote! {}, + parse_quote! { + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + assert_eq!( + kprobe.expand().unwrap().to_string(), + quote! { + #[no_mangle] + #[link_section = "kprobe"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + } + } + .to_string() + ); + } + + #[test] + fn test_kprobe_with_function() { + let kprobe = KProbe::parse( + KProbeKind::KProbe, + parse_quote! { + function = "fib_lookup" + }, + parse_quote! { + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + assert_eq!( + kprobe.expand().unwrap().to_string(), + quote! { + #[no_mangle] + #[link_section = "kprobe/fib_lookup"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + } + } + .to_string() + ); + } + + #[test] + fn test_kprobe_with_function_and_offset() { + let kprobe = KProbe::parse( + KProbeKind::KProbe, + parse_quote! { + function = "fib_lookup", + offset = "10" + }, + parse_quote! { + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + assert_eq!( + kprobe.expand().unwrap().to_string(), + quote! { + #[no_mangle] + #[link_section = "kprobe/fib_lookup+10"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + } + } + .to_string() + ); + } + + #[test] + fn test_kretprobe() { + let kprobe = KProbe::parse( + KProbeKind::KRetProbe, + parse_quote! {}, + parse_quote! { + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + assert_eq!( + kprobe.expand().unwrap().to_string(), + quote! { + #[no_mangle] + #[link_section = "kretprobe"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + } + } + .to_string() + ); + } +} diff --git a/aya-bpf-macros/src/lib.rs b/aya-bpf-macros/src/lib.rs index 182a2fcd..247b636a 100644 --- a/aya-bpf-macros/src/lib.rs +++ b/aya-bpf-macros/src/lib.rs @@ -1,121 +1,187 @@ -mod expand; - -use expand::{ - Args, BtfTracePoint, CgroupDevice, CgroupSkb, CgroupSock, CgroupSockAddr, CgroupSockopt, - CgroupSysctl, FEntry, FExit, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, - SchedClassifier, SkLookup, SkMsg, SkSkb, SkSkbKind, SockAddrArgs, SockOps, SocketFilter, - SockoptArgs, TracePoint, Xdp, -}; +pub(crate) mod args; +mod btf_tracepoint; +mod cgroup_device; +mod cgroup_skb; +mod cgroup_sock; +mod cgroup_sock_addr; +mod cgroup_sockopt; +mod cgroup_sysctl; +mod fentry; +mod fexit; +mod kprobe; +mod lsm; +mod map; +mod perf_event; +mod raw_tracepoint; +mod sk_lookup; +mod sk_msg; +mod sk_skb; +mod sock_ops; +mod socket_filter; +mod tc; +mod tracepoint; +mod uprobe; +mod xdp; + +use btf_tracepoint::BtfTracePoint; +use cgroup_device::CgroupDevice; +use cgroup_skb::CgroupSkb; +use cgroup_sock::CgroupSock; +use cgroup_sock_addr::CgroupSockAddr; +use cgroup_sockopt::CgroupSockopt; +use cgroup_sysctl::CgroupSysctl; +use fentry::FEntry; +use fexit::FExit; +use kprobe::{KProbe, KProbeKind}; +use lsm::Lsm; +use map::Map; +use perf_event::PerfEvent; use proc_macro::TokenStream; -use syn::{parse_macro_input, ItemFn, ItemStatic}; - +use proc_macro_error::{abort, proc_macro_error}; +use raw_tracepoint::RawTracePoint; +use sk_lookup::SkLookup; +use sk_msg::SkMsg; +use sk_skb::{SkSkb, SkSkbKind}; +use sock_ops::SockOps; +use socket_filter::SocketFilter; +use uprobe::{UProbe, UProbeKind}; + +use tc::SchedClassifier; +use tracepoint::TracePoint; +use xdp::Xdp; +#[proc_macro_error] #[proc_macro_attribute] pub fn map(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemStatic); - - Map::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match Map::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn kprobe(attrs: TokenStream, item: TokenStream) -> TokenStream { - probe(ProbeKind::KProbe, attrs, item) + match KProbe::parse(KProbeKind::KProbe, attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn kretprobe(attrs: TokenStream, item: TokenStream) -> TokenStream { - probe(ProbeKind::KRetProbe, attrs, item) + match KProbe::parse(KProbeKind::KRetProbe, attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn uprobe(attrs: TokenStream, item: TokenStream) -> TokenStream { - probe(ProbeKind::UProbe, attrs, item) + match UProbe::parse(UProbeKind::UProbe, attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn uretprobe(attrs: TokenStream, item: TokenStream) -> TokenStream { - probe(ProbeKind::URetProbe, attrs, item) + match UProbe::parse(UProbeKind::URetProbe, attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn sock_ops(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - SockOps::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match SockOps::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } +#[proc_macro_error] #[proc_macro_attribute] pub fn sk_msg(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - SkMsg::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match SkMsg::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } +#[proc_macro_error] #[proc_macro_attribute] pub fn xdp(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - Xdp::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match Xdp::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn classifier(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - SchedClassifier::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match SchedClassifier::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } #[proc_macro_attribute] pub fn cgroup_sysctl(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - CgroupSysctl::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match CgroupSysctl::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn cgroup_sockopt(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as SockoptArgs); - let attach_type = args.attach_type.to_string(); - let item = parse_macro_input!(item as ItemFn); - - CgroupSockopt::from_syn(args.args, item, attach_type) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match CgroupSockopt::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn cgroup_skb(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - CgroupSkb::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match CgroupSkb::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as a [`CgroupSockAddr`] eBPF program. @@ -152,59 +218,51 @@ pub fn cgroup_skb(attrs: TokenStream, item: TokenStream) -> TokenStream { /// Ok(0) /// } /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn cgroup_sock_addr(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as SockAddrArgs); - let attach_type = args.attach_type.to_string(); - let item = parse_macro_input!(item as ItemFn); - - CgroupSockAddr::from_syn(args.args, item, attach_type) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match CgroupSockAddr::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } - +#[proc_macro_error] #[proc_macro_attribute] pub fn cgroup_sock(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - CgroupSock::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() -} - -fn probe(kind: ProbeKind, attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - Probe::from_syn(kind, args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match CgroupSock::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } +#[proc_macro_error] #[proc_macro_attribute] pub fn tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - TracePoint::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match TracePoint::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } +#[proc_macro_error] #[proc_macro_attribute] pub fn perf_event(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - PerfEvent::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match PerfEvent::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as a raw tracepoint eBPF program that can be attached at a @@ -223,7 +281,7 @@ pub fn perf_event(attrs: TokenStream, item: TokenStream) -> TokenStream { /// ```no_run /// use aya_bpf::{macros::raw_tracepoint, programs::RawTracePointContext}; /// -/// #[raw_tracepoint(name = "sys_enter")] +/// #[raw_tracepoint(tracepoint = "sys_enter")] /// pub fn sys_enter(ctx: RawTracePointContext) -> i32 { /// match unsafe { try_sys_enter(ctx) } { /// Ok(ret) => ret, @@ -235,20 +293,25 @@ pub fn perf_event(attrs: TokenStream, item: TokenStream) -> TokenStream { /// Ok(0) /// } /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn raw_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - RawTracePoint::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match RawTracePoint::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as an LSM program that can be attached to Linux LSM hooks. /// Used to implement security policy and audit logging. /// +/// The hook name is the first argument to the macro. +/// You may also provide `sleepable = true` to mark the program as sleepable. +/// Arguments should be comma separated. +/// /// LSM probes can be attached to the kernel's security hooks to implement mandatory /// access control policy and security auditing. /// @@ -265,7 +328,7 @@ pub fn raw_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { /// ```no_run /// use aya_bpf::{macros::lsm, programs::LsmContext}; /// -/// #[lsm(name = "file_open")] +/// #[lsm(hook = "file_open")] /// pub fn file_open(ctx: LsmContext) -> i32 { /// match unsafe { try_file_open(ctx) } { /// Ok(ret) => ret, @@ -277,15 +340,16 @@ pub fn raw_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { /// Ok(0) /// } /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn lsm(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - Lsm::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match Lsm::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as a [BTF-enabled raw tracepoint][1] eBPF program that can be attached at @@ -304,7 +368,7 @@ pub fn lsm(attrs: TokenStream, item: TokenStream) -> TokenStream { /// ```no_run /// use aya_bpf::{macros::btf_tracepoint, programs::BtfTracePointContext}; /// -/// #[btf_tracepoint(name = "sched_process_fork")] +/// #[btf_tracepoint(function = "sched_process_fork")] /// pub fn sched_process_fork(ctx: BtfTracePointContext) -> u32 { /// match unsafe { try_sched_process_fork(ctx) } { /// Ok(ret) => ret, @@ -318,15 +382,16 @@ pub fn lsm(attrs: TokenStream, item: TokenStream) -> TokenStream { /// ``` /// /// [1]: https://github.com/torvalds/linux/commit/9e15db66136a14cde3f35691f1d839d950118826 +#[proc_macro_error] #[proc_macro_attribute] pub fn btf_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - BtfTracePoint::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match BtfTracePoint::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as a SK_SKB Stream Parser eBPF program that can be attached @@ -354,6 +419,7 @@ pub fn btf_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { /// Ok(ctx.len()) ///} /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn stream_parser(attrs: TokenStream, item: TokenStream) -> TokenStream { sk_skb(SkSkbKind::StreamParser, attrs, item) @@ -384,19 +450,20 @@ pub fn stream_parser(attrs: TokenStream, item: TokenStream) -> TokenStream { /// Ok(sk_action::SK_PASS) ///} /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn stream_verdict(attrs: TokenStream, item: TokenStream) -> TokenStream { sk_skb(SkSkbKind::StreamVerdict, attrs, item) } fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - SkSkb::from_syn(kind, args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match SkSkb::parse(kind, attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as a eBPF Socket Filter program that can be attached to @@ -411,20 +478,21 @@ fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream /// ```no_run /// use aya_bpf::{macros::socket_filter, programs::SkBuffContext}; /// -/// #[socket_filter(name = "accept_all")] +/// #[socket_filter] /// pub fn accept_all(_ctx: SkBuffContext) -> i64 { /// return 0 /// } /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - SocketFilter::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match SocketFilter::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as a fentry eBPF program that can be attached to almost @@ -444,7 +512,7 @@ pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream { /// # type filename = u32; /// # type path = u32; /// -/// #[fentry(name = "filename_lookup")] +/// #[fentry(function = "filename_lookup")] /// fn filename_lookup(ctx: FEntryContext) -> i32 { /// match unsafe { try_filename_lookup(ctx) } { /// Ok(ret) => ret, @@ -459,15 +527,16 @@ pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream { /// Ok(0) /// } /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn fentry(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - FEntry::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match FEntry::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as a fexit eBPF program that can be attached to almost @@ -488,7 +557,7 @@ pub fn fentry(attrs: TokenStream, item: TokenStream) -> TokenStream { /// # type filename = u32; /// # type path = u32; /// -/// #[fexit(name = "filename_lookup")] +/// #[fexit(function = "filename_lookup")] /// fn filename_lookup(ctx: FExitContext) -> i32 { /// match unsafe { try_filename_lookup(ctx) } { /// Ok(ret) => ret, @@ -503,15 +572,16 @@ pub fn fentry(attrs: TokenStream, item: TokenStream) -> TokenStream { /// Ok(0) /// } /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - FExit::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match FExit::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as an eBPF Socket Lookup program that can be attached to @@ -526,21 +596,22 @@ pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream { /// ```no_run /// use aya_bpf::{macros::sk_lookup, programs::SkLookupContext}; /// -/// #[sk_lookup(name = "redirect")] +/// #[sk_lookup] /// pub fn accept_all(_ctx: SkLookupContext) -> u32 { /// // use sk_assign to redirect /// return 0 /// } /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - SkLookup::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match SkLookup::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } /// Marks a function as a cgroup device eBPF program that can be attached to a @@ -558,19 +629,20 @@ pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream { /// programs::DeviceContext, /// }; /// -/// #[cgroup_device(name="cgroup_dev")] +/// #[cgroup_device] /// pub fn cgroup_dev(ctx: DeviceContext) -> i32 { /// // Reject all device access /// return 0; /// } /// ``` +#[proc_macro_error] #[proc_macro_attribute] pub fn cgroup_device(attrs: TokenStream, item: TokenStream) -> TokenStream { - let args = parse_macro_input!(attrs as Args); - let item = parse_macro_input!(item as ItemFn); - - CgroupDevice::from_syn(args, item) - .and_then(|u| u.expand()) - .unwrap_or_else(|err| err.to_compile_error()) - .into() + match CgroupDevice::parse(attrs.into(), item.into()) { + Ok(prog) => prog + .expand() + .unwrap_or_else(|err| abort!(err.span(), "{}", err)) + .into(), + Err(err) => abort!(err.span(), "{}", err), + } } diff --git a/aya-bpf-macros/src/lsm.rs b/aya-bpf-macros/src/lsm.rs new file mode 100644 index 00000000..546d9385 --- /dev/null +++ b/aya-bpf-macros/src/lsm.rs @@ -0,0 +1,85 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +use crate::args::err_on_unknown_args; + +pub(crate) struct Lsm { + item: ItemFn, + hook: String, + sleepable: bool, +} + +impl Lsm { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if attrs.is_empty() { + abort!(attrs, "missing hook name"); + } + let mut args = syn::parse2(attrs)?; + let item = syn::parse2(item)?; + let hook = crate::args::pop_required_arg(&mut args, "hook")?; + err_on_unknown_args(&args)?; + Ok(Lsm { + item, + hook, + sleepable: false, + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_prefix = if self.sleepable { "lsm.s" } else { "lsm" }; + let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.hook).into(); + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + // LSM probes need to return an integer corresponding to the correct + // policy decision. Therefore we do not simply default to a return value + // of 0 as in other program types. + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 { + return #fn_name(::aya_bpf::programs::LsmContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_lsm() { + let prog = Lsm::parse( + parse_quote! { + hook = "bprm_committed_creds" + }, + parse_quote! { + fn bprm_committed_creds(ctx: &mut ::aya_bpf::programs::LsmContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "lsm/bprm_committed_creds"] + fn bprm_committed_creds(ctx: *mut ::core::ffi::c_void) -> i32 { + return bprm_committed_creds(::aya_bpf::programs::LsmContext::new(ctx)); + + fn bprm_committed_creds(ctx: &mut ::aya_bpf::programs::LsmContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/map.rs b/aya-bpf-macros/src/map.rs new file mode 100644 index 00000000..044d9395 --- /dev/null +++ b/aya-bpf-macros/src/map.rs @@ -0,0 +1,75 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use quote::quote; +use syn::Result; + +use crate::args::name_arg; + +use syn::ItemStatic; +pub(crate) struct Map { + item: ItemStatic, + name: String, +} + +impl Map { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + let mut args = syn::parse2(attrs)?; + let item: ItemStatic = syn::parse2(item)?; + let name = name_arg(&mut args)?.unwrap_or_else(|| item.ident.to_string()); + Ok(Map { item, name }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = "maps".to_string().into(); + let name = &self.name; + let item = &self.item; + Ok(quote! { + #[link_section = #section_name] + #[export_name = #name] + #item + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_map_with_name() { + let map = Map::parse( + parse_quote!(name = "foo"), + parse_quote!( + static BAR: HashMap<&'static str, u32> = HashMap::new(); + ), + ) + .unwrap(); + let expanded = map.expand().unwrap(); + let expected = quote!( + #[link_section = "maps"] + #[export_name = "foo"] + static BAR: HashMap<&'static str, u32> = HashMap::new(); + ); + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn test_map_no_name() { + let map = Map::parse( + parse_quote!(), + parse_quote!( + static BAR: HashMap<&'static str, u32> = HashMap::new(); + ), + ) + .unwrap(); + let expanded = map.expand().unwrap(); + let expected = quote!( + #[link_section = "maps"] + #[export_name = "BAR"] + static BAR: HashMap<&'static str, u32> = HashMap::new(); + ); + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/perf_event.rs b/aya-bpf-macros/src/perf_event.rs new file mode 100644 index 00000000..3756959d --- /dev/null +++ b/aya-bpf-macros/src/perf_event.rs @@ -0,0 +1,67 @@ +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +pub(crate) struct PerfEvent { + item: ItemFn, +} + +impl PerfEvent { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute") + } + let item = syn::parse2(item)?; + Ok(PerfEvent { item }) + } + + pub(crate) fn expand(&self) -> Result { + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = "perf_event"] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = #fn_name(::aya_bpf::programs::PerfEventContext::new(ctx)); + return 0; + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_perf_event() { + let prog = PerfEvent::parse( + parse_quote!(), + parse_quote!( + fn foo(ctx: PerfEventContext) -> i32 { + 0 + } + ), + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "perf_event"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::PerfEventContext::new(ctx)); + return 0; + + fn foo(ctx: PerfEventContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/raw_tracepoint.rs b/aya-bpf-macros/src/raw_tracepoint.rs new file mode 100644 index 00000000..a220251d --- /dev/null +++ b/aya-bpf-macros/src/raw_tracepoint.rs @@ -0,0 +1,76 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +use crate::args::{err_on_unknown_args, pop_required_arg}; + +pub(crate) struct RawTracePoint { + item: ItemFn, + tracepoint: String, +} + +impl RawTracePoint { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if attrs.is_empty() { + abort!(attrs, "expected `tracepoint` attribute") + } + let mut args = syn::parse2(attrs)?; + let item = syn::parse2(item)?; + let tracepoint = pop_required_arg(&mut args, "tracepoint")?; + err_on_unknown_args(&args)?; + Ok(RawTracePoint { item, tracepoint }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = format!("raw_tp/{}", self.tracepoint).into(); + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = #fn_name(::aya_bpf::programs::RawTracePointContext::new(ctx)); + return 0; + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_raw_tracepoint() { + let prog = RawTracePoint::parse( + parse_quote! { tracepoint = "sys_enter" }, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::RawTracePointContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "raw_tp/sys_enter"] + fn prog(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = prog(::aya_bpf::programs::RawTracePointContext::new(ctx)); + return 0; + + fn prog(ctx: &mut ::aya_bpf::programs::RawTracePointContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/sk_lookup.rs b/aya-bpf-macros/src/sk_lookup.rs new file mode 100644 index 00000000..ddeae3e8 --- /dev/null +++ b/aya-bpf-macros/src/sk_lookup.rs @@ -0,0 +1,65 @@ +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +pub(crate) struct SkLookup { + item: ItemFn, +} + +impl SkLookup { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute") + } + let item = syn::parse2(item)?; + Ok(SkLookup { item }) + } + + pub(crate) fn expand(&self) -> Result { + let fn_name = self.item.sig.ident.clone(); + let fn_vis = &self.item.vis; + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = "sk_lookup"] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sk_lookup) -> u32 { + return #fn_name(::aya_bpf::programs::SkLookupContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_sk_lookup() { + let prog = SkLookup::parse( + parse_quote! {}, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::SkLookupContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "sk_lookup"] + fn prog(ctx: *mut ::aya_bpf::bindings::bpf_sk_lookup) -> u32 { + return prog(::aya_bpf::programs::SkLookupContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::SkLookupContext) -> u32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/sk_msg.rs b/aya-bpf-macros/src/sk_msg.rs new file mode 100644 index 00000000..bfb6be8a --- /dev/null +++ b/aya-bpf-macros/src/sk_msg.rs @@ -0,0 +1,65 @@ +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +pub(crate) struct SkMsg { + item: ItemFn, +} + +impl SkMsg { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute") + } + let item = syn::parse2(item)?; + Ok(SkMsg { item }) + } + + pub(crate) fn expand(&self) -> Result { + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = "sk_msg"] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::sk_msg_md) -> u32 { + return #fn_name(::aya_bpf::programs::SkMsgContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_sk_msg() { + let prog = SkMsg::parse( + parse_quote! {}, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::SkMsgContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "sk_msg"] + fn prog(ctx: *mut ::aya_bpf::bindings:: sk_msg_md) -> u32 { + return prog(::aya_bpf::programs::SkMsgContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::SkMsgContext) -> u32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/sk_skb.rs b/aya-bpf-macros/src/sk_skb.rs new file mode 100644 index 00000000..72f57ff7 --- /dev/null +++ b/aya-bpf-macros/src/sk_skb.rs @@ -0,0 +1,115 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +#[allow(clippy::enum_variant_names)] +#[derive(Debug, Copy, Clone)] +pub(crate) enum SkSkbKind { + StreamVerdict, + StreamParser, +} + +impl std::fmt::Display for SkSkbKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use SkSkbKind::*; + match self { + StreamVerdict => write!(f, "stream_verdict"), + StreamParser => write!(f, "stream_parser"), + } + } +} + +pub(crate) struct SkSkb { + kind: SkSkbKind, + item: ItemFn, +} + +impl SkSkb { + pub(crate) fn parse(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute"); + } + let item = syn::parse2(item)?; + Ok(SkSkb { item, kind }) + } + + pub(crate) fn expand(&self) -> Result { + let kind = &self.kind; + let section_name: Cow<'_, _> = format!("sk_skb/{kind}").into(); + let fn_name = self.item.sig.ident.clone(); + let fn_vis = &self.item.vis; + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 { + return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_stream_parser() { + let prog = SkSkb::parse( + SkSkbKind::StreamParser, + parse_quote! {}, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "sk_skb/stream_parser"] + fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 { + return prog(::aya_bpf::programs::SkBuffContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> u32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn test_stream_verdict() { + let prog = SkSkb::parse( + SkSkbKind::StreamVerdict, + parse_quote! {}, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "sk_skb/stream_verdict"] + fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 { + return prog(::aya_bpf::programs::SkBuffContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> u32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/sock_ops.rs b/aya-bpf-macros/src/sock_ops.rs new file mode 100644 index 00000000..b47ead9c --- /dev/null +++ b/aya-bpf-macros/src/sock_ops.rs @@ -0,0 +1,65 @@ +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +pub(crate) struct SockOps { + item: ItemFn, +} + +impl SockOps { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute") + } + let item = syn::parse2(item)?; + Ok(SockOps { item }) + } + + pub(crate) fn expand(&self) -> Result { + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = "sockops"] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock_ops) -> u32 { + return #fn_name(::aya_bpf::programs::SockOpsContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_sock_ops() { + let prog = SockOps::parse( + parse_quote! {}, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::SockOpsContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "sockops"] + fn prog(ctx: *mut ::aya_bpf::bindings::bpf_sock_ops) -> u32 { + return prog(::aya_bpf::programs::SockOpsContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::SockOpsContext) -> u32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/socket_filter.rs b/aya-bpf-macros/src/socket_filter.rs new file mode 100644 index 00000000..6511d635 --- /dev/null +++ b/aya-bpf-macros/src/socket_filter.rs @@ -0,0 +1,65 @@ +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +pub(crate) struct SocketFilter { + item: ItemFn, +} + +impl SocketFilter { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute") + } + let item = syn::parse2(item)?; + Ok(SocketFilter { item }) + } + + pub(crate) fn expand(&self) -> Result { + let fn_name = self.item.sig.ident.clone(); + let fn_vis = &self.item.vis; + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = "socket"] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i64 { + return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_socket_filter() { + let prog = SocketFilter::parse( + parse_quote! {}, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> i64 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "socket"] + fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i64 { + return prog(::aya_bpf::programs::SkBuffContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> i64 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/tc.rs b/aya-bpf-macros/src/tc.rs new file mode 100644 index 00000000..9db12ccd --- /dev/null +++ b/aya-bpf-macros/src/tc.rs @@ -0,0 +1,65 @@ +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +pub(crate) struct SchedClassifier { + item: ItemFn, +} + +impl SchedClassifier { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + if !attrs.is_empty() { + abort!(attrs, "unexpected attribute") + } + let item = syn::parse2(item)?; + Ok(SchedClassifier { item }) + } + + pub(crate) fn expand(&self) -> Result { + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = "classifier"] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return #fn_name(::aya_bpf::programs::TcContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_sched_classifier() { + let prog = SchedClassifier::parse( + parse_quote! {}, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::TcContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "classifier"] + fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 { + return prog(::aya_bpf::programs::TcContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::TcContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/tracepoint.rs b/aya-bpf-macros/src/tracepoint.rs new file mode 100644 index 00000000..4a1b5f50 --- /dev/null +++ b/aya-bpf-macros/src/tracepoint.rs @@ -0,0 +1,84 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +use crate::args::{err_on_unknown_args, pop_arg}; + +pub(crate) struct TracePoint { + item: ItemFn, + category: Option, + name: Option, +} + +impl TracePoint { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + let mut args = syn::parse2(attrs)?; + let item = syn::parse2(item)?; + let name = pop_arg(&mut args, "name"); + let category = pop_arg(&mut args, "category"); + err_on_unknown_args(&args)?; + Ok(TracePoint { + item, + category, + name, + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = match (&self.category, &self.name) { + (Some(category), Some(name)) => format!("tracepoint/{}/{}", category, name).into(), + (Some(_), None) => abort!(self.item, "expected `name` and `category` arguments"), + (None, Some(_)) => abort!(self.item, "expected `name` and `category` arguments"), + _ => "tracepoint".into(), + }; + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = #fn_name(::aya_bpf::programs::TracePointContext::new(ctx)); + return 0; + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_tracepoint() { + let prog = TracePoint::parse( + parse_quote! { name = "sys_enter_bind", category = "syscalls" }, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::TracePointContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "tracepoint/syscalls/sys_enter_bind"] + fn prog(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = prog(::aya_bpf::programs::TracePointContext::new(ctx)); + return 0; + + fn prog(ctx: &mut ::aya_bpf::programs::TracePointContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/aya-bpf-macros/src/uprobe.rs b/aya-bpf-macros/src/uprobe.rs new file mode 100644 index 00000000..7f06d0d3 --- /dev/null +++ b/aya-bpf-macros/src/uprobe.rs @@ -0,0 +1,232 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{ItemFn, Result}; + +use crate::args::{err_on_unknown_args, pop_arg}; + +#[allow(clippy::enum_variant_names)] +#[derive(Debug, Copy, Clone)] +pub(crate) enum UProbeKind { + UProbe, + URetProbe, +} + +impl std::fmt::Display for UProbeKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use UProbeKind::*; + match self { + UProbe => write!(f, "uprobe"), + URetProbe => write!(f, "uretprobe"), + } + } +} + +pub(crate) struct UProbe { + kind: UProbeKind, + path: Option, + function: Option, + offset: Option, + item: ItemFn, +} + +impl UProbe { + pub(crate) fn parse(kind: UProbeKind, attrs: TokenStream, item: TokenStream) -> Result { + let mut path = None; + let mut function = None; + let mut offset = None; + if !attrs.is_empty() { + let mut args = syn::parse2(attrs)?; + path = pop_arg(&mut args, "path"); + function = pop_arg(&mut args, "function"); + offset = pop_arg(&mut args, "offset").map(|v| v.parse::().unwrap()); + err_on_unknown_args(&args)?; + } + + let item = syn::parse2(item)?; + Ok(UProbe { + kind, + item, + path, + function, + offset, + }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = if self.path.is_some() && self.offset.is_some() { + if self.function.is_none() { + abort!(self.item.sig.ident, "expected `function` attribute"); + } + let mut path = self.path.as_ref().unwrap().clone(); + if path.starts_with('/') { + path.remove(0); + } + format!( + "{}/{}:{}+{}", + self.kind, + path, + self.function.as_ref().unwrap(), + self.offset.unwrap() + ) + .into() + } else if self.path.is_some() { + if self.function.is_none() { + abort!(self.item.sig.ident, "expected `function` attribute"); + } + let mut path = self.path.as_ref().unwrap().clone(); + if path.starts_with('/') { + path.remove(0); + } + format!("{}/{}:{}", self.kind, path, self.function.as_ref().unwrap()).into() + } else { + format!("{}", self.kind).into() + }; + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = #fn_name(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn uprobe() { + let uprobe = UProbe::parse( + UProbeKind::UProbe, + parse_quote! {}, + parse_quote! { + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + assert_eq!( + uprobe.expand().unwrap().to_string(), + quote! { + #[no_mangle] + #[link_section = "uprobe"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + } + } + .to_string() + ); + } + + #[test] + fn uprobe_with_path() { + let uprobe = UProbe::parse( + UProbeKind::UProbe, + parse_quote! { + path = "/self/proc/exe", + function = "trigger_uprobe" + }, + parse_quote! { + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + assert_eq!( + uprobe.expand().unwrap().to_string(), + quote! { + #[no_mangle] + #[link_section = "uprobe/self/proc/exe:trigger_uprobe"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + } + } + .to_string() + ); + } + + #[test] + fn test_uprobe_with_path_and_offset() { + let uprobe = UProbe::parse( + UProbeKind::UProbe, + parse_quote! { + path = "/self/proc/exe", function = "foo", offset = "123" + }, + parse_quote! { + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + assert_eq!( + uprobe.expand().unwrap().to_string(), + quote! { + #[no_mangle] + #[link_section = "uprobe/self/proc/exe:foo+123"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + } + } + .to_string() + ); + } + + #[test] + fn test_uretprobe() { + let uprobe = UProbe::parse( + UProbeKind::URetProbe, + parse_quote! {}, + parse_quote! { + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + }, + ) + .unwrap(); + assert_eq!( + uprobe.expand().unwrap().to_string(), + quote! { + #[no_mangle] + #[link_section = "uretprobe"] + fn foo(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx)); + return 0; + + fn foo(ctx: ProbeContext) -> u32 { + 0 + } + } + } + .to_string() + ); + } +} diff --git a/aya-bpf-macros/src/xdp.rs b/aya-bpf-macros/src/xdp.rs new file mode 100644 index 00000000..f9b01529 --- /dev/null +++ b/aya-bpf-macros/src/xdp.rs @@ -0,0 +1,104 @@ +use std::borrow::Cow; + +use proc_macro2::TokenStream; +use proc_macro_error::abort; +use quote::quote; +use syn::{Ident, ItemFn, Result}; + +pub(crate) struct Xdp { + item: ItemFn, + frags: bool, +} + +impl Xdp { + pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { + let mut frags = false; + let item = syn::parse2(item)?; + if !attrs.is_empty() { + let ident: Ident = syn::parse2(attrs)?; + if ident != "frags" { + abort!(ident, "unexpected attribute"); + } + frags = true; + } + Ok(Xdp { item, frags }) + } + + pub(crate) fn expand(&self) -> Result { + let section_name: Cow<'_, _> = if self.frags { + "xdp.frags".into() + } else { + "xdp".into() + }; + let fn_vis = &self.item.vis; + let fn_name = self.item.sig.ident.clone(); + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 { + return #fn_name(::aya_bpf::programs::XdpContext::new(ctx)); + + #item + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use syn::parse_quote; + + #[test] + fn test_xdp() { + let prog = Xdp::parse( + parse_quote! {}, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "xdp"] + fn prog(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 { + return prog(::aya_bpf::programs::XdpContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn test_xdp_frags() { + let prog = Xdp::parse( + parse_quote! { frags }, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "xdp.frags"] + fn prog(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 { + return prog(::aya_bpf::programs::XdpContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } +} diff --git a/test/integration-ebpf/src/map_test.rs b/test/integration-ebpf/src/map_test.rs index d63d6fc6..b8a66fce 100644 --- a/test/integration-ebpf/src/map_test.rs +++ b/test/integration-ebpf/src/map_test.rs @@ -14,7 +14,7 @@ static FOO: Array = Array::::with_max_entries(10, 0); #[map(name = "BAR")] static BAZ: Array = Array::::with_max_entries(10, 0); -#[xdp(frags = "true")] +#[xdp(frags)] pub fn pass(ctx: XdpContext) -> u32 { match unsafe { try_pass(ctx) } { Ok(ret) => ret, diff --git a/test/integration-ebpf/src/pass.rs b/test/integration-ebpf/src/pass.rs index 120ef7cc..8de36333 100644 --- a/test/integration-ebpf/src/pass.rs +++ b/test/integration-ebpf/src/pass.rs @@ -5,7 +5,7 @@ use aya_bpf::{bindings::xdp_action, macros::xdp, programs::XdpContext}; // Note: the `frags` attribute causes this probe to be incompatible with kernel versions < 5.18.0. // See https://github.com/torvalds/linux/commit/c2f2cdb. -#[xdp(name = "pass", frags = "true")] +#[xdp(frags)] pub fn pass(ctx: XdpContext) -> u32 { match unsafe { try_pass(ctx) } { Ok(ret) => ret, diff --git a/test/integration-ebpf/src/test.rs b/test/integration-ebpf/src/test.rs index 187a3926..8b0f5d79 100644 --- a/test/integration-ebpf/src/test.rs +++ b/test/integration-ebpf/src/test.rs @@ -7,7 +7,7 @@ use aya_bpf::{ programs::{ProbeContext, TracePointContext, XdpContext}, }; -#[xdp(name = "test_xdp")] +#[xdp] pub fn pass(ctx: XdpContext) -> u32 { match unsafe { try_pass(ctx) } { Ok(ret) => ret, diff --git a/test/integration-test/src/tests/rbpf.rs b/test/integration-test/src/tests/rbpf.rs index c12bbce8..d29ec66e 100644 --- a/test/integration-test/src/tests/rbpf.rs +++ b/test/integration-test/src/tests/rbpf.rs @@ -1,6 +1,7 @@ use core::{mem::size_of, ptr::null_mut, slice::from_raw_parts}; use std::collections::HashMap; +use assert_matches::assert_matches; use aya_obj::{generated::bpf_insn, Object, ProgramSection}; #[test] @@ -8,8 +9,8 @@ fn run_with_rbpf() { let object = Object::parse(crate::PASS).unwrap(); assert_eq!(object.programs.len(), 1); - assert_matches::assert_matches!(object.programs["pass"].section, ProgramSection::Xdp { .. }); - assert_eq!(object.programs["pass"].section.name(), "pass"); + assert_matches!(object.programs["pass"].section, ProgramSection::Xdp { .. }); + assert_eq!(object.programs["pass"].section.name(), "xdp.frags"); let instructions = &object .functions @@ -35,7 +36,7 @@ fn use_map_with_rbpf() { let mut object = Object::parse(crate::MULTIMAP_BTF).unwrap(); assert_eq!(object.programs.len(), 1); - assert_matches::assert_matches!( + assert_matches!( object.programs["bpf_prog"].section, ProgramSection::TracePoint { .. } ); From 6d92119fbc7b3339b1b135bcf8f3de747242442f Mon Sep 17 00:00:00 2001 From: Dave Tucker Date: Mon, 24 Jul 2023 19:04:19 +0100 Subject: [PATCH 3/3] integration-test: Add test for 2 progs in same section Signed-off-by: Dave Tucker --- test/integration-ebpf/Cargo.toml | 4 ++++ test/integration-ebpf/src/two_progs.rs | 21 +++++++++++++++++++++ test/integration-test/src/lib.rs | 1 + test/integration-test/src/tests/smoke.rs | 24 +++++++++++++++++++++++- 4 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 test/integration-ebpf/src/two_progs.rs diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index bb179c59..38c097f3 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/test/integration-ebpf/Cargo.toml @@ -38,3 +38,7 @@ path = "src/relocations.rs" [[bin]] name = "bpf_probe_read" path = "src/bpf_probe_read.rs" + +[[bin]] +name = "two_progs" +path = "src/two_progs.rs" diff --git a/test/integration-ebpf/src/two_progs.rs b/test/integration-ebpf/src/two_progs.rs new file mode 100644 index 00000000..0e4b1f9a --- /dev/null +++ b/test/integration-ebpf/src/two_progs.rs @@ -0,0 +1,21 @@ +// Two programs in the same ELF section + +#![no_std] +#![no_main] + +use aya_bpf::{macros::tracepoint, programs::TracePointContext}; + +#[tracepoint] +pub fn test_tracepoint_one(_ctx: TracePointContext) -> u32 { + 0 +} +#[tracepoint] +pub fn test_tracepoint_two(_ctx: TracePointContext) -> u32 { + 0 +} + +#[cfg(not(test))] +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + loop {} +} diff --git a/test/integration-test/src/lib.rs b/test/integration-test/src/lib.rs index b8e5120d..043b0cff 100644 --- a/test/integration-test/src/lib.rs +++ b/test/integration-test/src/lib.rs @@ -13,6 +13,7 @@ pub const NAME_TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/n pub const PASS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/pass")); pub const TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/test")); pub const RELOCATIONS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/relocations")); +pub const TWO_PROGS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/two_progs")); pub const BPF_PROBE_READ: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/bpf_probe_read")); diff --git a/test/integration-test/src/tests/smoke.rs b/test/integration-test/src/tests/smoke.rs index 89f7d0a1..78614bd9 100644 --- a/test/integration-test/src/tests/smoke.rs +++ b/test/integration-test/src/tests/smoke.rs @@ -1,5 +1,5 @@ use aya::{ - programs::{Extension, Xdp, XdpFlags}, + programs::{Extension, TracePoint, Xdp, XdpFlags}, util::KernelVersion, Bpf, BpfLoader, }; @@ -22,6 +22,28 @@ fn xdp() { dispatcher.attach("lo", XdpFlags::default()).unwrap(); } +#[test] +fn two_progs() { + let mut bpf = Bpf::load(crate::TWO_PROGS).unwrap(); + + let prog_one: &mut TracePoint = bpf + .program_mut("test_tracepoint_one") + .unwrap() + .try_into() + .unwrap(); + + prog_one.load().unwrap(); + prog_one.attach("sched", "sched_switch").unwrap(); + + let prog_two: &mut TracePoint = bpf + .program_mut("test_tracepoint_two") + .unwrap() + .try_into() + .unwrap(); + prog_two.load().unwrap(); + prog_two.attach("sched", "sched_switch").unwrap(); +} + #[test] fn extension() { let _netns = NetNsGuard::new();