1use proc_macro::TokenStream;
2use quote::{ToTokens, quote};
3use syn::meta::ParseNestedMeta;
4use syn::{FnArg, Ident, ItemFn, LitInt, LitStr, PatType, ReturnType, Type};
5use syn::{parse_macro_input, parse_quote};
6
7#[derive(PartialEq, Eq)]
8enum EventType {
9 Apply,
10 Delete,
11}
12
13impl TryFrom<String> for EventType {
14 type Error = String;
15
16 fn try_from(value: String) -> Result<Self, Self::Error> {
17 match value.as_str() {
18 "Apply" => Ok(EventType::Apply),
19 "Delete" => Ok(EventType::Delete),
20 _ => Err(format!("Invalid event type {value}")),
21 }
22 }
23}
24
25impl ToTokens for EventType {
26 fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
27 let ty: Type = match self {
28 EventType::Apply => parse_quote!(::kubus::EventType::Apply),
29 EventType::Delete => parse_quote!(::kubus::EventType::Delete),
30 };
31 ty.to_tokens(tokens);
32 }
33}
34
35fn extract_generic_arg(ty: &Type, pos: usize) -> Option<&Type> {
36 if let syn::Type::Path(type_path) = ty
37 && let Some(last_segment) = type_path.path.segments.last()
38 && let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments
39 && let Some(syn::GenericArgument::Type(inner_type)) = args.args.get(pos)
40 {
41 return Some(inner_type);
42 }
43 None
44}
45
46fn extract_return_type(func: &ItemFn) -> Option<&Type> {
47 match func.sig.output {
48 ReturnType::Type(_, ref ty) => Some(ty),
49 ReturnType::Default => None,
50 }
51}
52
53fn extract_func_arg_type(arg: &FnArg) -> Option<&PatType> {
54 match arg {
55 FnArg::Typed(pat) => Some(pat),
56 FnArg::Receiver(_) => None,
57 }
58}
59
60fn extract_resource_type(func: &ItemFn) -> Option<&Type> {
61 func.sig
62 .inputs
63 .first()
64 .and_then(|arg| extract_func_arg_type(arg))
65 .and_then(|pat| extract_generic_arg(&pat.ty, 0))
66}
67
68fn extract_context_type(func: &ItemFn) -> Option<&Type> {
69 func.sig
70 .inputs
71 .iter()
72 .nth(1)
73 .and_then(|arg| extract_func_arg_type(arg))
74 .and_then(|arc| extract_generic_arg(&arc.ty, 0))
75 .and_then(|ctx| extract_generic_arg(ctx, 0))
76}
77
78fn extract_function_return_error_type(func: &ItemFn) -> Option<&Type> {
79 extract_return_type(func).and_then(|ty| extract_generic_arg(ty, 1))
80}
81
82fn internal_prefix(ident: Ident) -> Ident {
83 let value = ident.to_string();
84 let prefixed = format!("__kubus_{value}");
85 Ident::new(&prefixed, ident.span())
86}
87
88fn quote_option<T>(value: Option<T>) -> proc_macro2::TokenStream
89where
90 T: ToTokens,
91{
92 match value {
93 Some(inner) => quote! { Some(#inner) },
94 None => quote! { None },
95 }
96}
97
98#[derive(Default)]
99struct EventHandlerAttrs {
100 event: Option<EventType>,
101 finalizer: Option<LitStr>,
102 label_selector: Option<LitStr>,
103 field_selector: Option<LitStr>,
104 requeue_interval: Option<LitInt>,
105}
106
107impl EventHandlerAttrs {
108 fn parse(&mut self, meta: ParseNestedMeta) -> syn::parse::Result<()> {
109 if meta.path.is_ident("event") {
110 let str: Ident = meta.value()?.parse()?;
111 self.event = Some(
112 str.to_string()
113 .try_into()
114 .map_err(|err| syn::parse::Error::new(str.span(), err))?,
115 );
116 Ok(())
117 } else if meta.path.is_ident("finalizer") {
118 self.finalizer = Some(meta.value()?.parse()?);
119 Ok(())
120 } else if meta.path.is_ident("label_selector") {
121 self.label_selector = Some(meta.value()?.parse()?);
122 Ok(())
123 } else if meta.path.is_ident("field_selector") {
124 self.field_selector = Some(meta.value()?.parse()?);
125 Ok(())
126 } else if meta.path.is_ident("requeue_interval") {
127 self.requeue_interval = Some(meta.value()?.parse()?);
128 Ok(())
129 } else {
130 Err(meta.error("unsupported kubus property"))
131 }
132 }
133}
134
135#[proc_macro_attribute]
194pub fn kubus(args: TokenStream, input: TokenStream) -> TokenStream {
195 let mut attrs = EventHandlerAttrs::default();
196 let attr_parser = syn::meta::parser(|meta| attrs.parse(meta));
197 parse_macro_input!(args with attr_parser);
198
199 let event = attrs.event.expect("event kubus attribute missing");
200
201 let func = parse_macro_input!(input as ItemFn);
202 let resource_ty = extract_resource_type(&func)
203 .expect("unable to extract resource type from handler function");
204 let context_ty =
205 extract_context_type(&func).expect("unable to extract resource type from handler function");
206 let error_ty = extract_function_return_error_type(&func)
207 .cloned()
208 .unwrap_or_else(|| parse_quote! { ::kubus::HandlerError });
209
210 let internal_func = {
211 let mut func = func.clone();
212 func.sig.ident = internal_prefix(func.sig.ident);
213 func
214 };
215
216 let struct_name = func.sig.ident.clone();
217 let internal_func_name = internal_func.sig.ident.clone();
218
219 let field_selector = quote_option(attrs.field_selector);
220 let label_selector = quote_option(attrs.label_selector);
221 let requeue_interval = attrs.requeue_interval.unwrap_or_else(|| parse_quote!(30));
222
223 let update_finalizer = attrs
224 .finalizer
225 .map(|finalizer| {
226 let update_func = match event {
227 EventType::Apply => quote! { ::kubus::apply_finalizer },
228 EventType::Delete => quote! { ::kubus::remove_finalizer },
229 };
230 quote! {
231 let namespace = resource.namespace();
232 let client = context.client.clone();
233 let api: ::kube::Api<#resource_ty> =
234 <#resource_ty as ::kube::Resource>::Scope::api(client, namespace);
235
236 #update_func(&api, #finalizer, resource).await?;
237 }
238 })
239 .unwrap_or_default();
240
241 let handler_name = LitStr::new(&struct_name.to_string(), struct_name.span());
242
243 quote! {
244 #[allow(non_snake_case)]
245 #internal_func
246
247 #[allow(non_camel_case_types)]
248 #[doc(hidden)]
249 pub struct #struct_name;
250
251 #[::async_trait::async_trait]
252 impl ::kubus::Handler<#resource_ty, #context_ty, #error_ty> for #struct_name {
253 const NAME: &'static str = #handler_name;
254 const FIELD_SELECTOR: Option<&'static str> = #field_selector;
255 const LABEL_SELECTOR: Option<&'static str> = #label_selector;
256
257 async fn handle(
258 resource: ::std::sync::Arc<#resource_ty>,
259 context: ::std::sync::Arc<::kubus::Context<#context_ty>>,
260 ) -> ::std::result::Result<::kube::runtime::controller::Action, #error_ty> {
261 use ::kubus::ScopeExt;
262 use ::kube::{Resource, ResourceExt};
263 let requeue = ::kube::runtime::controller::Action::requeue(
264 ::std::time::Duration::from_secs(#requeue_interval)
265 );
266
267 if let (::kubus::EventType::Apply, Some(_)) | (::kubus::EventType::Delete, None) =
268 (#event, resource.meta().deletion_timestamp.as_ref())
269 {
270 return Ok(requeue);
271 }
272
273 #internal_func_name(resource.clone(), context.clone())
274 .await
275 .map_err(|err| ::kubus::Error::Handler(Box::new(err)))?;
276
277 #update_finalizer
278
279 Ok(requeue)
280 }
281 }
282
283 #[::async_trait::async_trait]
284 impl ::kubus::Runnable<#context_ty, #error_ty> for #struct_name {
285 fn name(&self) -> &'static str {
286 <#struct_name as ::kubus::Handler<#resource_ty, #context_ty, #error_ty>>::NAME
287 }
288
289 async fn run(
290 &self,
291 client: ::kube::Client,
292 context: ::std::sync::Arc<::kubus::Context<#context_ty>>,
293 ) -> ::std::result::Result<(), #error_ty> {
294 <#struct_name as ::kubus::Handler<#resource_ty, #context_ty, #error_ty>>::run(self, client, context).await
295 }
296 }
297 }
298 .into()
299}
300
301enum AdmissionKind {
302 Validating,
303 Mutating,
304}
305
306#[derive(Default)]
307struct AdmissionHandlerAttrs {
308 kind: Option<AdmissionKind>,
309}
310
311impl AdmissionHandlerAttrs {
312 fn parse(&mut self, meta: ParseNestedMeta) -> syn::parse::Result<()> {
313 if meta.path.is_ident("validating") && self.kind.is_none() {
314 self.kind = Some(AdmissionKind::Validating);
315 Ok(())
316 } else if meta.path.is_ident("mutating") && self.kind.is_none() {
317 self.kind = Some(AdmissionKind::Mutating);
318 Ok(())
319 } else {
320 Err(meta.error("unsupported kubus property"))
321 }
322 }
323}
324
325#[proc_macro_attribute]
326pub fn admission(args: TokenStream, input: TokenStream) -> TokenStream {
327 let mut attrs = AdmissionHandlerAttrs::default();
328 let attr_parser = syn::meta::parser(|meta| attrs.parse(meta));
329 parse_macro_input!(args with attr_parser);
330
331 let kind = attrs
332 .kind
333 .expect("admission attribute must specify either 'mutating' or 'validating'");
334
335 let func = parse_macro_input!(input as ItemFn);
336 let func_name = &func.sig.ident;
337 let name_string = LitStr::new(&func_name.to_string(), func_name.span());
338
339 let error_ty = extract_function_return_error_type(&func)
341 .cloned()
342 .unwrap_or_else(|| parse_quote! { ::kubus::HandlerError });
343
344 let internal_func = {
345 let mut func = func.clone();
346 func.sig.ident = internal_prefix(func.sig.ident);
347 func
348 };
349 let internal_func_name = &internal_func.sig.ident;
350
351 let (trait_name, method_name): (Type, Ident) = match kind {
352 AdmissionKind::Mutating => (
353 parse_quote! { ::kubus::admission::MutatingAdmissionHandler },
354 Ident::new("mutate", func_name.span()),
355 ),
356 AdmissionKind::Validating => (
357 parse_quote! { ::kubus::admission::ValidatingAdmissionHandler },
358 Ident::new("validate", func_name.span()),
359 ),
360 };
361
362 quote! {
363 #[allow(non_snake_case)]
364 #internal_func
365
366 #[allow(non_camel_case_types)]
367 #[doc(hidden)]
368 pub struct #func_name;
369
370 #[::async_trait::async_trait]
371 impl #trait_name for #func_name {
372 type Err = #error_ty;
373
374 fn name(&self) -> &'static str {
375 #name_string
376 }
377
378 async fn #method_name(
379 &self,
380 req: &::kube::core::admission::AdmissionRequest<::kube::api::DynamicObject>,
381 ) -> ::std::result::Result<::kube::core::admission::AdmissionResponse, Self::Err> {
382 #internal_func_name(req).await
383 }
384 }
385 }
386 .into()
387}