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    TokenStream::from(quote! {
38        #additional_macro
39        #modified
40    })
41}
42
43/// Modifies the component module and adds the necessary functions and structs.
44fn modify_component_module(mut module: ItemMod, component_type: &Type) -> ItemMod {
45    let (initialize_fn, initialize_struct) = generate_initialize(component_type);
46    let (destroy_fn, destroy_struct) = generate_destroy(component_type);
47    //let (apply_fn, apply_struct, apply_impl, update_fn, update_struct) = generate_instructions(component_type);
48    let (update_fn, update_with_session_fn, update_struct, update_with_session_struct) =
49        generate_update(component_type);
50
51    module.content = module.content.map(|(brace, mut items)| {
52        items.extend(
53            vec![
54                initialize_fn,
55                initialize_struct,
56                update_fn,
57                update_struct,
58                update_with_session_fn,
59                update_with_session_struct,
60                destroy_fn,
61                destroy_struct,
62            ]
63            .into_iter()
64            .map(|item| syn::parse2(item).unwrap())
65            .collect::<Vec<_>>(),
66        );
67
68        let modified_items = items
69            .into_iter()
70            .map(|item| match item {
71                syn::Item::Struct(mut struct_item)
72                    if struct_item.ident == "Apply" || struct_item.ident == "ApplyWithSession" =>
73                {
74                    modify_apply_struct(&mut struct_item);
75                    syn::Item::Struct(struct_item)
76                }
77                _ => item,
78            })
79            .collect();
80        (brace, modified_items)
81    });
82
83    module
84}
85
86/// Extracts the type name from attribute arguments.
87fn extract_type_name(args: &AttributeArgs) -> Option<Type> {
88    args.iter().find_map(|arg| {
89        if let NestedMeta::Meta(syn::Meta::Path(path)) = arg {
90            Some(Type::Path(syn::TypePath {
91                qself: None,
92                path: path.clone(),
93            }))
94        } else {
95            None
96        }
97    })
98}
99
100/// Modifies the Apply struct, change the bolt system to accept any compatible system.
101fn modify_apply_struct(struct_item: &mut ItemStruct) {
102    if let Fields::Named(fields_named) = &mut struct_item.fields {
103        fields_named
104            .named
105            .iter_mut()
106            .filter(|field| is_expecting_program(field))
107            .for_each(|field| {
108                field.ty = syn::parse_str("UncheckedAccount<'info>").expect("Failed to parse type");
109                field.attrs.push(create_check_attribute());
110            });
111    }
112}
113
114/// Creates the check attribute.
115fn create_check_attribute() -> Attribute {
116    parse_quote! {
117        #[doc = "CHECK: This program can modify the data of the component"]
118    }
119}
120
121/// Generates the destroy function and struct.
122fn generate_destroy(component_type: &Type) -> (TokenStream2, TokenStream2) {
123    (
124        quote! {
125            #[automatically_derived]
126            pub fn destroy(ctx: Context<Destroy>) -> Result<()> {
127                let program_data_address =
128                    Pubkey::find_program_address(&[crate::id().as_ref()], &bolt_lang::prelude::solana_program::bpf_loader_upgradeable::id()).0;
129
130                if !program_data_address.eq(ctx.accounts.component_program_data.key) {
131                    return Err(BoltError::InvalidAuthority.into());
132                }
133
134                let program_account_data = ctx.accounts.component_program_data.try_borrow_data()?;
135                let upgrade_authority = if let bolt_lang::prelude::solana_program::bpf_loader_upgradeable::UpgradeableLoaderState::ProgramData {
136                    upgrade_authority_address,
137                    ..
138                } =
139                    bolt_lang::prelude::bincode::deserialize(&program_account_data).map_err(|_| BoltError::InvalidAuthority)?
140                {
141                    Ok(upgrade_authority_address)
142                } else {
143                    Err(anchor_lang::error::Error::from(BoltError::InvalidAuthority))
144                }?.ok_or_else(|| BoltError::InvalidAuthority)?;
145
146                if ctx.accounts.authority.key != &ctx.accounts.component.bolt_metadata.authority && ctx.accounts.authority.key != &upgrade_authority {
147                    return Err(BoltError::InvalidAuthority.into());
148                }
149
150                let instruction = anchor_lang::solana_program::sysvar::instructions::get_instruction_relative(
151                    0, &ctx.accounts.instruction_sysvar_account.to_account_info()
152                ).map_err(|_| BoltError::InvalidCaller)?;
153                if instruction.program_id != World::id() {
154                    return Err(BoltError::InvalidCaller.into());
155                }
156                Ok(())
157            }
158        },
159        quote! {
160            #[automatically_derived]
161            #[derive(Accounts)]
162            pub struct Destroy<'info> {
163                #[account()]
164                pub authority: Signer<'info>,
165                #[account(mut)]
166                pub receiver: AccountInfo<'info>,
167                #[account()]
168                pub entity: Account<'info, Entity>,
169                #[account(mut, close = receiver, seeds = [<#component_type>::seed(), entity.key().as_ref()], bump)]
170                pub component: Account<'info, #component_type>,
171                #[account()]
172                pub component_program_data: AccountInfo<'info>,
173                #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
174                pub instruction_sysvar_account: AccountInfo<'info>,
175                pub system_program: Program<'info, System>,
176            }
177        },
178    )
179}
180
181/// Generates the initialize function and struct.
182fn generate_initialize(component_type: &Type) -> (TokenStream2, TokenStream2) {
183    (
184        quote! {
185            #[automatically_derived]
186            pub fn initialize(ctx: Context<Initialize>) -> Result<()> {
187                let instruction = anchor_lang::solana_program::sysvar::instructions::get_instruction_relative(
188                    0, &ctx.accounts.instruction_sysvar_account.to_account_info()
189                ).map_err(|_| BoltError::InvalidCaller)?;
190                if instruction.program_id != World::id() {
191                    return Err(BoltError::InvalidCaller.into());
192                }
193                ctx.accounts.data.set_inner(<#component_type>::default());
194                ctx.accounts.data.bolt_metadata.authority = *ctx.accounts.authority.key;
195                Ok(())
196            }
197        },
198        quote! {
199            #[automatically_derived]
200            #[derive(Accounts)]
201            pub struct Initialize<'info>  {
202                #[account(mut)]
203                pub payer: Signer<'info>,
204                #[account(init_if_needed, payer = payer, space = <#component_type>::size(), seeds = [<#component_type>::seed(), entity.key().as_ref()], bump)]
205                pub data: Account<'info, #component_type>,
206                #[account()]
207                pub entity: Account<'info, Entity>,
208                #[account()]
209                pub authority: AccountInfo<'info>,
210                #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
211                pub instruction_sysvar_account: UncheckedAccount<'info>,
212                pub system_program: Program<'info, System>,
213            }
214        },
215    )
216}
217
218/// Generates the instructions and related structs to inject in the component.
219fn generate_update(
220    component_type: &Type,
221) -> (TokenStream2, TokenStream2, TokenStream2, TokenStream2) {
222    (
223        quote! {
224            #[automatically_derived]
225            pub fn update(ctx: Context<Update>, data: Vec<u8>) -> Result<()> {
226                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);
227
228                // Check if the instruction is called from the world program
229                let instruction = anchor_lang::solana_program::sysvar::instructions::get_instruction_relative(
230                    0, &ctx.accounts.instruction_sysvar_account.to_account_info()
231                ).map_err(|_| BoltError::InvalidCaller)?;
232                require_eq!(instruction.program_id, World::id(), BoltError::InvalidCaller);
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                // Check if the instruction is called from the world program
255                let instruction = anchor_lang::solana_program::sysvar::instructions::get_instruction_relative(
256                    0, &ctx.accounts.instruction_sysvar_account.to_account_info()
257                ).map_err(|_| BoltError::InvalidCaller)?;
258                require_eq!(instruction.program_id, World::id(), BoltError::InvalidCaller);
259
260                ctx.accounts.bolt_component.set_inner(<#component_type>::try_from_slice(&data)?);
261                Ok(())
262            }
263        },
264        quote! {
265            #[automatically_derived]
266            #[derive(Accounts)]
267            pub struct Update<'info> {
268                #[account(mut)]
269                pub bolt_component: Account<'info, #component_type>,
270                #[account()]
271                pub authority: Signer<'info>,
272                #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
273                pub instruction_sysvar_account: UncheckedAccount<'info>
274            }
275        },
276        quote! {
277            #[automatically_derived]
278            #[derive(Accounts)]
279            pub struct UpdateWithSession<'info> {
280                #[account(mut)]
281                pub bolt_component: Account<'info, #component_type>,
282                #[account()]
283                pub authority: Signer<'info>,
284                #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
285                pub instruction_sysvar_account: UncheckedAccount<'info>,
286                #[account(constraint = session_token.to_account_info().owner == &bolt_lang::session_keys::ID)]
287                pub session_token: Account<'info, bolt_lang::session_keys::SessionToken>,
288            }
289        },
290    )
291}
292
293/// Checks if the field is expecting a program.
294fn is_expecting_program(field: &Field) -> bool {
295    field.ty.to_token_stream().to_string().contains("Program")
296}