poem_grants_proc_macro/
lib.rs

1extern crate proc_macro;
2use proc_macro::TokenStream;
3use quote::ToTokens;
4use syn::{parse_macro_input, AttributeArgs, ImplItem, ImplItemMethod, ItemFn, ItemImpl, Meta};
5
6use crate::expand::{FnType, ProtectEndpoint};
7
8mod expand;
9
10const HAS_AUTHORITIES: &str = "has_permissions";
11const HAS_ANY_AUTHORITY: &str = "has_any_permission";
12
13const HAS_ROLES: &str = "has_roles";
14const HAS_ANY_ROLE: &str = "has_any_role";
15
16macro_rules! unwrap_result {
17    ($result:expr) => {
18        match $result {
19            Ok(result) => result,
20            Err(err) => return syn::Error::from(err).to_compile_error().into(),
21        }
22    };
23}
24
25/// Macro to сheck that the user has all the specified permissions.
26/// Allow to add a conditional restriction based on handlers parameters.
27/// Add the `secure` attribute followed by the the boolean expression to validate based on parameters
28///
29/// Attention: these macros have to be above of (`[poem::handler]`) or `oai`
30///
31/// Also you can use you own types instead of Strings, just add `type` attribute with path to type
32/// # Examples
33/// ```
34/// use poem::{Response, http::StatusCode};
35///
36/// // User should be ADMIN with OP_GET_SECRET permission
37/// #[poem_grants::has_permissions["ROLE_ADMIN", "OP_GET_SECRET"]]
38/// #[poem::handler]
39/// async fn macro_secured() -> Response {
40///     Response::builder().status(StatusCode::OK).body("some secured info")
41/// }
42///
43/// // User should be ADMIN with OP_GET_SECRET permission and the user.id param should be equal
44/// // to the path parameter {user_id}
45/// struct User {id: i32}
46/// #[poem_grants::has_permissions["ROLE_ADMIN", "OP_GET_SECRET", secure="*user_id == user.id"]]
47/// #[poem::handler]
48/// async fn macro_secured_params(user_id: web::Path<i32>, user: web::Data<User>) -> Response {
49///     Response::builder().status(StatusCode::OK).body("some secured info with user_id path equal to user.id")
50///}
51///
52/// // User must have MyPermissionEnum::OP_GET_SECRET (you own enum example)
53/// #[poem_grants::has_permissions["OP_GET_SECRET", type = "MyPermissionEnum"]]
54/// #[poem::handler]
55/// async fn macro_enum_secured() -> Response {
56///     Response::builder().status(StatusCode::OK).body("some secured info")
57/// }
58///
59///```
60#[proc_macro_attribute]
61pub fn has_permissions(args: TokenStream, input: TokenStream) -> TokenStream {
62    check_permissions(HAS_AUTHORITIES, args, input)
63}
64
65/// Macro to сheck that the user has any of the specified permissions.
66///
67/// # Examples
68/// ```
69/// use poem::{Response, http::StatusCode};
70///
71/// // User should be ADMIN or MANAGER
72/// #[poem_grants::has_any_permission["ROLE_ADMIN", "ROLE_MANAGER"]]
73/// #[poem::handler]
74/// async fn macro_secured() -> Response {
75///     Response::builder().status(StatusCode::OK).body("some secured info")
76/// }
77/// ```
78#[proc_macro_attribute]
79pub fn has_any_permission(args: TokenStream, input: TokenStream) -> TokenStream {
80    check_permissions(HAS_ANY_AUTHORITY, args, input)
81}
82
83/// Macro to сheck that the user has all the specified roles.
84/// Role - is permission with prefix "ROLE_".
85///
86/// # Examples
87/// ```
88/// use poem::{Response, http::StatusCode};
89///
90/// // User should be ADMIN and MANAGER
91/// #[poem_grants::has_roles["ADMIN", "MANAGER"]]
92/// #[poem::handler]
93/// async fn macro_secured() -> Response {
94///     Response::builder().status(StatusCode::OK).body("some secured info")
95/// }
96/// ```
97#[proc_macro_attribute]
98pub fn has_roles(args: TokenStream, input: TokenStream) -> TokenStream {
99    check_permissions(HAS_ROLES, args, input)
100}
101
102/// Macro to сheck that the user has any the specified roles.
103/// Role - is permission with prefix "ROLE_".
104///
105/// # Examples
106/// ```
107/// use poem::{Response, http::StatusCode};
108///
109/// // User should be ADMIN or MANAGER
110/// #[poem_grants::has_any_role["ADMIN", "MANAGER"]]
111/// #[poem::handler]
112/// async fn macro_secured() -> Response {
113///     Response::builder().status(StatusCode::OK).body("some secured info")
114/// }
115/// ```
116#[proc_macro_attribute]
117pub fn has_any_role(args: TokenStream, input: TokenStream) -> TokenStream {
118    check_permissions(HAS_ANY_ROLE, args, input)
119}
120
121fn check_permissions(check_fn_name: &str, args: TokenStream, input: TokenStream) -> TokenStream {
122    let args = parse_macro_input!(args as AttributeArgs);
123    let func = parse_macro_input!(input as ItemFn);
124
125    unwrap_result!(ProtectEndpoint::new(check_fn_name, args, FnType::Fn(func)))
126        .into_token_stream()
127        .into()
128}
129
130/// Macro for `poem-openapi` support
131/// Add macro `#[poem_grants::open_api]` above of `#[poem_openapi::OpenApi]` and mark all needed methods with necessary security-methods:
132/// One of [`has_permissions`, `has_any_permission`, `has_roles`, `has_any_role`]
133///
134/// # Examples
135/// ```
136/// use poem_openapi::payload::PlainText;
137///
138/// struct Api;
139///
140/// #[poem_grants::open_api]
141/// #[poem_openapi::OpenApi]
142/// impl Api {
143///     // An example of protection via `proc-macro`.
144///     // Just use the necessary name of macro provided by `poem-grants` without crate-name:
145///     #[has_permissions("OP_READ_ADMIN_INFO")]
146///     #[oai(path = "/admin", method = "get")]
147///     async fn macro_secured(&self) -> PlainText<String> {
148///         PlainText("ADMIN_RESPONSE".to_string())
149///     }
150/// }
151/// ```
152#[proc_macro_attribute]
153pub fn open_api(_args: TokenStream, input: TokenStream) -> TokenStream {
154    let mut item_impl = parse_macro_input!(input as ItemImpl);
155    let mut methods = Vec::new();
156    for (idx, item) in item_impl.items.iter().enumerate() {
157        if let ImplItem::Method(method) = item {
158            for grants_attr in method
159                .attrs
160                .iter()
161                .filter(|attr| is_poem_grants_attr(*attr))
162            {
163                let args = match unwrap_result!(grants_attr.parse_meta()) {
164                    Meta::List(list) => list.nested.into_iter().collect::<Vec<syn::NestedMeta>>(),
165                    _ => {
166                        return syn::Error::new_spanned(
167                            grants_attr,
168                            "Expected endpoint-attribute to be a list",
169                        )
170                        .to_compile_error()
171                        .into()
172                    }
173                };
174
175                let generated = unwrap_result!(ProtectEndpoint::new(
176                    &grants_attr
177                        .path
178                        .get_ident()
179                        .expect("validated by condition above")
180                        .to_string(),
181                    args,
182                    FnType::Method(method.clone()),
183                ))
184                .into_token_stream()
185                .into();
186
187                let mut gen_method = parse_macro_input!(generated as ImplItemMethod);
188
189                gen_method.attrs.retain(|attr| attr != grants_attr);
190
191                methods.push((idx, gen_method));
192            }
193        }
194    }
195
196    for (idx, method) in methods {
197        let _ = std::mem::replace(&mut item_impl.items[idx], ImplItem::Method(method));
198    }
199
200    let res = quote::quote! {
201        #item_impl
202    };
203
204    res.into()
205}
206
207fn is_poem_grants_attr(attr: &syn::Attribute) -> bool {
208    attr.path.is_ident(HAS_ANY_AUTHORITY)
209        || attr.path.is_ident(HAS_AUTHORITIES)
210        || attr.path.is_ident(HAS_ANY_ROLE)
211        || attr.path.is_ident(HAS_ROLES)
212}
poem_grants_proc_macro/lib.rs

poem_grants_proc_macro/
lib.rs
