bolt_attribute_bolt_program/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::TokenStream as TokenStream2;
3use quote::{quote, ToTokens};
4use syn::{
5    parse_macro_input, parse_quote, Attribute, AttributeArgs, Field, Fields, ItemMod, ItemStruct,
6    NestedMeta, Type,
7};
8
9/// This macro attribute is used to define a BOLT component.
10///
11/// Bolt components are themselves programs that can be called by other programs.
12///
13/// # Example
14/// ```ignore
15/// #[bolt_program(Position)]
16/// #[program]
17/// pub mod component_position {
18///     use super::*;
19/// }
20///
21///
22/// #[component]
23/// pub struct Position {
24///     pub x: i64,
25///     pub y: i64,
26///     pub z: i64,
27/// }
28/// ```
29#[proc_macro_attribute]
30pub fn bolt_program(args: TokenStream, input: TokenStream) -> TokenStream {
31    let ast = parse_macro_input!(input as syn::ItemMod);
32    let args = parse_macro_input!(args as syn::AttributeArgs);
33    let component_type =
34        extract_type_name(&args).expect("Expected a component type in macro arguments");
35    let modified = modify_component_module(ast, &component_type);
36    let additional_macro: Attribute = parse_quote! { #[program] };
37    let cpi_checker = generate_cpi_checker();
38    TokenStream::from(quote! {
39        #cpi_checker
40        #additional_macro
41        #modified
42    })
43}
44
45/// Modifies the component module and adds the necessary functions and structs.
46fn modify_component_module(mut module: ItemMod, component_type: &Type) -> ItemMod {
47    let (initialize_fn, initialize_struct) = generate_initialize(component_type);
48    let (destroy_fn, destroy_struct) = generate_destroy(component_type);
49    let (update_fn, update_with_session_fn, update_struct, update_with_session_struct) =
50        generate_update(component_type);
51
52    module.content = module.content.map(|(brace, mut items)| {
53        items.extend(
54            vec![
55                initialize_fn,
56                initialize_struct,
57                update_fn,
58                update_struct,
59                update_with_session_fn,
60                update_with_session_struct,
61                destroy_fn,
62                destroy_struct,
63            ]
64            .into_iter()
65            .map(|item| syn::parse2(item).unwrap())
66            .collect::<Vec<_>>(),
67        );
68
69        let modified_items = items
70            .into_iter()
71            .map(|item| match item {
72                syn::Item::Struct(mut struct_item)
73                    if struct_item.ident == "Apply" || struct_item.ident == "ApplyWithSession" =>
74                {
75                    modify_apply_struct(&mut struct_item);
76                    syn::Item::Struct(struct_item)
77                }
78                _ => item,
79            })
80            .collect();
81        (brace, modified_items)
82    });
83
84    module
85}
86
87/// Extracts the type name from attribute arguments.
88fn extract_type_name(args: &AttributeArgs) -> Option<Type> {
89    args.iter().find_map(|arg| {
90        if let NestedMeta::Meta(syn::Meta::Path(path)) = arg {
91            Some(Type::Path(syn::TypePath {
92                qself: None,
93                path: path.clone(),
94            }))
95        } else {
96            None
97        }
98    })
99}
100
101/// Modifies the Apply struct, change the bolt system to accept any compatible system.
102fn modify_apply_struct(struct_item: &mut ItemStruct) {
103    if let Fields::Named(fields_named) = &mut struct_item.fields {
104        fields_named
105            .named
106            .iter_mut()
107            .filter(|field| is_expecting_program(field))
108            .for_each(|field| {
109                field.ty = syn::parse_str("UncheckedAccount<'info>").expect("Failed to parse type");
110                field.attrs.push(create_check_attribute());
111            });
112    }
113}
114
115/// Creates the check attribute.
116fn create_check_attribute() -> Attribute {
117    parse_quote! {
118        #[doc = "CHECK: This program can modify the data of the component"]
119    }
120}
121
122/// Generates the CPI checker function.
123fn generate_cpi_checker() -> TokenStream2 {
124    quote! {
125        fn cpi_checker<'info>(cpi_auth: &AccountInfo<'info>) -> Result<()> {
126            if !cpi_auth.is_signer || cpi_auth.key != &bolt_lang::world::World::cpi_auth_address() {
127                return Err(BoltError::InvalidCaller.into());
128            }
129            Ok(())
130        }
131    }
132}
133
134/// Generates the destroy function and struct.
135fn generate_destroy(component_type: &Type) -> (TokenStream2, TokenStream2) {
136    (
137        quote! {
138            #[automatically_derived]
139            pub fn destroy(ctx: Context<Destroy>) -> Result<()> {
140                let program_data_address =
141                    Pubkey::find_program_address(&[crate::id().as_ref()], &bolt_lang::prelude::solana_program::bpf_loader_upgradeable::id()).0;
142
143                if !program_data_address.eq(ctx.accounts.component_program_data.key) {
144                    return Err(BoltError::InvalidAuthority.into());
145                }
146
147                let program_account_data = ctx.accounts.component_program_data.try_borrow_data()?;
148                let upgrade_authority = if let bolt_lang::prelude::solana_program::bpf_loader_upgradeable::UpgradeableLoaderState::ProgramData {
149                    upgrade_authority_address,
150                    ..
151                } =
152                    bolt_lang::prelude::bincode::deserialize(&program_account_data).map_err(|_| BoltError::InvalidAuthority)?
153                {
154                    Ok(upgrade_authority_address)
155                } else {
156                    Err(anchor_lang::error::Error::from(BoltError::InvalidAuthority))
157                }?.ok_or_else(|| BoltError::InvalidAuthority)?;
158
159                if ctx.accounts.authority.key != &ctx.accounts.component.bolt_metadata.authority && ctx.accounts.authority.key != &upgrade_authority {
160                    return Err(BoltError::InvalidAuthority.into());
161                }
162
163                cpi_checker(&ctx.accounts.cpi_auth.to_account_info())?;
164
165                Ok(())
166            }
167        },
168        quote! {
169            #[automatically_derived]
170            #[derive(Accounts)]
171            pub struct Destroy<'info> {
172                #[account()]
173                pub cpi_auth: Signer<'info>,
174                #[account()]
175                pub authority: Signer<'info>,
176                #[account(mut)]
177                pub receiver: AccountInfo<'info>,
178                #[account()]
179                pub entity: Account<'info, Entity>,
180                #[account(mut, close = receiver, seeds = [<#component_type>::seed(), entity.key().as_ref()], bump)]
181                pub component: Account<'info, #component_type>,
182                #[account()]
183                pub component_program_data: AccountInfo<'info>,
184                pub system_program: Program<'info, System>,
185            }
186        },
187    )
188}
189
190/// Generates the initialize function and struct.
191fn generate_initialize(component_type: &Type) -> (TokenStream2, TokenStream2) {
192    (
193        quote! {
194            #[automatically_derived]
195            pub fn initialize(ctx: Context<Initialize>) -> Result<()> {
196                cpi_checker(&ctx.accounts.cpi_auth.to_account_info())?;
197                ctx.accounts.data.set_inner(<#component_type>::default());
198                ctx.accounts.data.bolt_metadata.authority = *ctx.accounts.authority.key;
199                Ok(())
200            }
201        },
202        quote! {
203            #[automatically_derived]
204            #[derive(Accounts)]
205            pub struct Initialize<'info>  {
206                #[account()]
207                pub cpi_auth: Signer<'info>,
208                #[account(mut)]
209                pub payer: Signer<'info>,
210                #[account(init_if_needed, payer = payer, space = <#component_type>::size(), seeds = [<#component_type>::seed(), entity.key().as_ref()], bump)]
211                pub data: Account<'info, #component_type>,
212                #[account()]
213                pub entity: Account<'info, Entity>,
214                #[account()]
215                pub authority: AccountInfo<'info>,
216                pub system_program: Program<'info, System>,
217            }
218        },
219    )
220}
221
222/// Generates the instructions and related structs to inject in the component.
223fn generate_update(
224    component_type: &Type,
225) -> (TokenStream2, TokenStream2, TokenStream2, TokenStream2) {
226    (
227        quote! {
228            #[automatically_derived]
229            pub fn update(ctx: Context<Update>, data: Vec<u8>) -> Result<()> {
230                require!(ctx.accounts.bolt_component.bolt_metadata.authority == World::id() || (ctx.accounts.bolt_component.bolt_metadata.authority == *ctx.accounts.authority.key && ctx.accounts.authority.is_signer), BoltError::InvalidAuthority);
231
232                cpi_checker(&ctx.accounts.cpi_auth.to_account_info())?;
233
234                ctx.accounts.bolt_component.set_inner(<#component_type>::try_from_slice(&data)?);
235                Ok(())
236            }
237        },
238        quote! {
239            #[automatically_derived]
240            pub fn update_with_session(ctx: Context<UpdateWithSession>, data: Vec<u8>) -> Result<()> {
241                if ctx.accounts.bolt_component.bolt_metadata.authority == World::id() {
242                    require!(Clock::get()?.unix_timestamp < ctx.accounts.session_token.valid_until, bolt_lang::session_keys::SessionError::InvalidToken);
243                } else {
244                    let validity_ctx = bolt_lang::session_keys::ValidityChecker {
245                        session_token: ctx.accounts.session_token.clone(),
246                        session_signer: ctx.accounts.authority.clone(),
247                        authority: ctx.accounts.bolt_component.bolt_metadata.authority.clone(),
248                        target_program: World::id(),
249                    };
250                    require!(ctx.accounts.session_token.validate(validity_ctx)?, bolt_lang::session_keys::SessionError::InvalidToken);
251                    require_eq!(ctx.accounts.bolt_component.bolt_metadata.authority, ctx.accounts.session_token.authority, bolt_lang::session_keys::SessionError::InvalidToken);
252                }
253
254                cpi_checker(&ctx.accounts.cpi_auth.to_account_info())?;
255
256                ctx.accounts.bolt_component.set_inner(<#component_type>::try_from_slice(&data)?);
257                Ok(())
258            }
259        },
260        quote! {
261            #[automatically_derived]
262            #[derive(Accounts)]
263            pub struct Update<'info> {
264                #[account()]
265                pub cpi_auth: Signer<'info>,
266                #[account(mut)]
267                pub bolt_component: Account<'info, #component_type>,
268                #[account()]
269                pub authority: Signer<'info>,
270            }
271        },
272        quote! {
273            #[automatically_derived]
274            #[derive(Accounts)]
275            pub struct UpdateWithSession<'info> {
276                #[account()]
277                pub cpi_auth: Signer<'info>,
278                #[account(mut)]
279                pub bolt_component: Account<'info, #component_type>,
280                #[account()]
281                pub authority: Signer<'info>,
282                #[account(constraint = session_token.to_account_info().owner == &bolt_lang::session_keys::ID)]
283                pub session_token: Account<'info, bolt_lang::session_keys::SessionToken>,
284            }
285        },
286    )
287}
288
289/// Checks if the field is expecting a program.
290fn is_expecting_program(field: &Field) -> bool {
291    field.ty.to_token_stream().to_string().contains("Program")
292}
bolt_attribute_bolt_program/lib.rs

bolt_attribute_bolt_program/
lib.rs
