cgroup_sock.rs 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. use std::borrow::Cow;
  2. use proc_macro2::TokenStream;
  3. use proc_macro_error::abort;
  4. use quote::quote;
  5. use syn::{Ident, ItemFn, Result};
  6. pub(crate) struct CgroupSock {
  7. item: ItemFn,
  8. attach_type: String,
  9. }
  10. impl CgroupSock {
  11. pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<CgroupSock> {
  12. if attrs.is_empty() {
  13. abort!(attrs, "missing attach type")
  14. }
  15. let item: ItemFn = syn::parse2(item)?;
  16. let attach_type: Ident = syn::parse2(attrs)?;
  17. match attach_type.to_string().as_str() {
  18. "post_bind4" | "post_bind6" | "sock_create" | "sock_release" => (),
  19. _ => abort!(attach_type, "invalid attach type"),
  20. }
  21. Ok(CgroupSock {
  22. item,
  23. attach_type: attach_type.to_string(),
  24. })
  25. }
  26. pub(crate) fn expand(&self) -> Result<TokenStream> {
  27. let section_name: Cow<'_, _> = format!("cgroup/{}", self.attach_type).into();
  28. let fn_vis = &self.item.vis;
  29. let fn_name = self.item.sig.ident.clone();
  30. let item = &self.item;
  31. Ok(quote! {
  32. #[no_mangle]
  33. #[link_section = #section_name]
  34. #fn_vis fn #fn_name(ctx: *mut ::aya_ebpf::bindings::bpf_sock) -> i32 {
  35. return #fn_name(::aya_ebpf::programs::SockContext::new(ctx));
  36. #item
  37. }
  38. })
  39. }
  40. }
  41. #[cfg(test)]
  42. mod tests {
  43. use syn::parse_quote;
  44. use super::*;
  45. #[test]
  46. fn cgroup_sock_post_bind4() {
  47. let prog = CgroupSock::parse(
  48. parse_quote!(post_bind4),
  49. parse_quote!(
  50. fn foo(ctx: SockContext) -> i32 {
  51. 0
  52. }
  53. ),
  54. )
  55. .unwrap();
  56. let expanded = prog.expand().unwrap();
  57. let expected = quote! {
  58. #[no_mangle]
  59. #[link_section = "cgroup/post_bind4"]
  60. fn foo(ctx: *mut ::aya_ebpf::bindings::bpf_sock) -> i32 {
  61. return foo(::aya_ebpf::programs::SockContext::new(ctx));
  62. fn foo(ctx: SockContext) -> i32 {
  63. 0
  64. }
  65. }
  66. };
  67. assert_eq!(expected.to_string(), expanded.to_string());
  68. }
  69. #[test]
  70. fn cgroup_sock_post_bind6() {
  71. let prog = CgroupSock::parse(
  72. parse_quote!(post_bind6),
  73. parse_quote!(
  74. fn foo(ctx: SockContext) -> i32 {
  75. 0
  76. }
  77. ),
  78. )
  79. .unwrap();
  80. let expanded = prog.expand().unwrap();
  81. let expected = quote! {
  82. #[no_mangle]
  83. #[link_section = "cgroup/post_bind6"]
  84. fn foo(ctx: *mut ::aya_ebpf::bindings::bpf_sock) -> i32 {
  85. return foo(::aya_ebpf::programs::SockContext::new(ctx));
  86. fn foo(ctx: SockContext) -> i32 {
  87. 0
  88. }
  89. }
  90. };
  91. assert_eq!(expected.to_string(), expanded.to_string());
  92. }
  93. #[test]
  94. fn cgroup_sock_sock_create() {
  95. let prog = CgroupSock::parse(
  96. parse_quote!(sock_create),
  97. parse_quote!(
  98. fn foo(ctx: SockContext) -> i32 {
  99. 0
  100. }
  101. ),
  102. )
  103. .unwrap();
  104. let expanded = prog.expand().unwrap();
  105. let expected = quote! {
  106. #[no_mangle]
  107. #[link_section = "cgroup/sock_create"]
  108. fn foo(ctx: *mut ::aya_ebpf::bindings::bpf_sock) -> i32 {
  109. return foo(::aya_ebpf::programs::SockContext::new(ctx));
  110. fn foo(ctx: SockContext) -> i32 {
  111. 0
  112. }
  113. }
  114. };
  115. assert_eq!(expected.to_string(), expanded.to_string());
  116. }
  117. #[test]
  118. fn cgroup_sock_sock_release() {
  119. let prog = CgroupSock::parse(
  120. parse_quote!(sock_release),
  121. parse_quote!(
  122. fn foo(ctx: SockContext) -> i32 {
  123. 0
  124. }
  125. ),
  126. )
  127. .unwrap();
  128. let expanded = prog.expand().unwrap();
  129. let expected = quote! {
  130. #[no_mangle]
  131. #[link_section = "cgroup/sock_release"]
  132. fn foo(ctx: *mut ::aya_ebpf::bindings::bpf_sock) -> i32 {
  133. return foo(::aya_ebpf::programs::SockContext::new(ctx));
  134. fn foo(ctx: SockContext) -> i32 {
  135. 0
  136. }
  137. }
  138. };
  139. assert_eq!(expected.to_string(), expanded.to_string());
  140. }
  141. }