fadroma_proc_auth/
lib.rs

1use proc_macro::TokenStream;
2use syn::{FnArg, ItemFn, Stmt, Type, parse_macro_input, parse_quote, Pat, PatIdent};
3use syn::punctuated::Punctuated;
4use syn::token::Comma;
5use quote::quote;
6
7#[proc_macro_attribute]
8pub fn require_admin(_attr: TokenStream, func: TokenStream) -> TokenStream {
9    let mut ast = parse_macro_input!(func as ItemFn);
10
11    let (deps, env) = find_extern_arg(&ast.sig.inputs);
12
13    let stmt = create_require_admin_stmt(deps, env);
14    ast.block.stmts.insert(0, stmt);
15
16    let result = quote! {
17        #ast
18    };
19
20    TokenStream::from(result)
21}
22
23fn find_extern_arg(args: &Punctuated<FnArg, Comma>) -> (PatIdent, PatIdent) {
24    let mut deps: Option<PatIdent> = None;
25    let mut info: Option<PatIdent> = None;
26
27    for arg in args {
28        match arg {
29            FnArg::Typed(item) => {
30                match item.ty.as_ref() {
31                    Type::Path(type_path) => {
32                        let info_arg = type_path.path.segments.iter()
33                            .find(|i| {
34                                if i.ident.to_string() == "MessageInfo" {
35                                    return true
36                                }
37                    
38                                false
39                            });
40
41                        let deps_arg = type_path.path.segments.iter().find(|i| {
42                            let ident = i.ident.to_string();
43                            if ident == "Deps" || ident == "DepsMut" {
44                                return true;
45                            }
46
47                            false
48                        });
49
50                        if let Some(_) = deps_arg {
51                            if let Pat::Ident(ident) = item.pat.as_ref() {
52                                deps = Some(ident.clone())
53                            }
54                        }
55                
56                        if let Some(_) = info_arg {
57                            if let Pat::Ident(ident) = item.pat.as_ref() {
58                                info = Some(ident.clone())
59                            }
60                        }
61                    }
62                    _ => continue
63                }
64            },
65            _ => continue
66        }
67    }
68    
69    return (
70        deps.unwrap_or_else(|| panic!("Couldn't find argument of type \"Deps or DepsMut\"")),
71        info.unwrap_or_else(|| panic!("Couldn't find arguments of type \"MessageInfo\""))
72    );
73}
74
75fn create_require_admin_stmt(deps: PatIdent, info: PatIdent) -> Stmt {
76    let ref deps = deps.ident;
77    let ref info = info.ident;
78
79    let code = quote! {
80        fadroma::admin::assert(#deps.as_ref(), &#info)?;
81    };
82
83    parse_quote!(#code)
84}
fadroma_proc_auth/lib.rs

fadroma_proc_auth/
lib.rs
