mainstay_attribute_access_control/
lib.rs

1extern crate proc_macro;
2
3use quote::quote;
4use syn::parse_macro_input;
5
6/// Executes the given access control method before running the decorated
7/// instruction handler. Any method in scope of the attribute can be invoked
8/// with any arguments from the associated instruction handler.
9///
10/// # Example
11///
12/// ```ignore
13/// use mainstay_lang::prelude::*;
14///
15/// #[program]
16/// mod errors {
17///     use super::*;
18///
19///     #[access_control(Create::accounts(&ctx, bump_seed))]
20///     pub fn create(ctx: Context<Create>, bump_seed: u8) -> Result<()> {
21///       let my_account = &mut ctx.accounts.my_account;
22///       my_account.bump_seed = bump_seed;
23///     }
24/// }
25///
26/// #[derive(Accounts)]
27/// pub struct Create {
28///   #[account(init)]
29///   my_account: Account<'info, MyAccount>,
30/// }
31///
32/// impl Create {
33///   pub fn accounts(ctx: &Context<Create>, bump_seed: u8) -> Result<()> {
34///     let seeds = &[ctx.accounts.my_account.to_account_info().key.as_ref(), &[bump_seed]];
35///     Pubkey::create_program_address(seeds, ctx.program_id)
36///       .map_err(|_| ErrorCode::InvalidNonce)?;
37///     Ok(())
38///   }
39/// }
40/// ```
41///
42/// This example demonstrates a useful pattern. Not only can you use
43/// `#[access_control]` to ensure any invariants or preconditions hold prior to
44/// executing an instruction, but also it can be used to finish any validation
45/// on the `Accounts` struct, particularly when instruction arguments are
46/// needed. Here, we use the given `bump_seed` to verify it creates a valid
47/// program-derived address.
48#[proc_macro_attribute]
49pub fn access_control(
50    args: proc_macro::TokenStream,
51    input: proc_macro::TokenStream,
52) -> proc_macro::TokenStream {
53    let mut args = args.to_string();
54    args.retain(|c| !c.is_whitespace());
55    let access_control: Vec<proc_macro2::TokenStream> = args
56        .split(')')
57        .filter(|ac| !ac.is_empty())
58        .map(|ac| format!("{ac})")) // Put back on the split char.
59        .map(|ac| format!("{ac}?;")) // Add `?;` syntax.
60        .map(|ac| ac.parse().unwrap())
61        .collect();
62
63    let item_fn = parse_macro_input!(input as syn::ItemFn);
64
65    let fn_attrs = item_fn.attrs;
66    let fn_vis = item_fn.vis;
67    let fn_sig = item_fn.sig;
68    let fn_block = item_fn.block;
69
70    let fn_stmts = fn_block.stmts;
71
72    proc_macro::TokenStream::from(quote! {
73        #(#fn_attrs)*
74        #fn_vis #fn_sig {
75
76            #(#access_control)*
77
78            #(#fn_stmts)*
79        }
80    })
81}
mainstay_attribute_access_control/lib.rs

mainstay_attribute_access_control/
lib.rs
