:sparkles: Separating apply and apply_with_session (#141)

clients/bolt-sdk/src/generated/idl/world.json

@@ -2,7 +2,7 @@
   "address": "WorLD15A7CrDwLcLy4fRqtaTb9fbd8o8iqiEMUDse2n",
   "metadata": {
     "name": "world",
-    "version": "0.2.0",
+    "version": "0.2.1",
     "spec": "0.1.0",
     "description": "Bolt World program",
     "repository": "https://github.com/magicblock-labs/bolt"
@@ -97,6 +97,41 @@
         162,
         225
       ],
+      "accounts": [
+        {
+          "name": "bolt_system"
+        },
+        {
+          "name": "authority",
+          "signer": true
+        },
+        {
+          "name": "instruction_sysvar_account",
+          "address": "Sysvar1nstructions1111111111111111111111111"
+        },
+        {
+          "name": "world"
+        }
+      ],
+      "args": [
+        {
+          "name": "args",
+          "type": "bytes"
+        }
+      ]
+    },
+    {
+      "name": "apply_with_session",
+      "discriminator": [
+        213,
+        69,
+        29,
+        230,
+        142,
+        107,
+        134,
+        103
+      ],
       "accounts": [
         {
           "name": "bolt_system"
@@ -113,8 +148,7 @@
           "name": "world"
         },
         {
-          "name": "session_token",
-          "optional": true
+          "name": "session_token"
         }
       ],
       "args": [

clients/bolt-sdk/src/generated/types/world.ts

@@ -8,7 +8,7 @@ export type World = {
   address: "WorLD15A7CrDwLcLy4fRqtaTb9fbd8o8iqiEMUDse2n";
   metadata: {
     name: "world";
-    version: "0.2.0";
+    version: "0.2.1";
     spec: "0.1.0";
     description: "Bolt World program";
     repository: "https://github.com/magicblock-labs/bolt";
@@ -76,6 +76,32 @@ export type World = {
     {
       name: "apply";
       discriminator: [248, 243, 145, 24, 105, 50, 162, 225];
+      accounts: [
+        {
+          name: "boltSystem";
+        },
+        {
+          name: "authority";
+          signer: true;
+        },
+        {
+          name: "instructionSysvarAccount";
+          address: "Sysvar1nstructions1111111111111111111111111";
+        },
+        {
+          name: "world";
+        },
+      ];
+      args: [
+        {
+          name: "args";
+          type: "bytes";
+        },
+      ];
+    },
+    {
+      name: "applyWithSession";
+      discriminator: [213, 69, 29, 230, 142, 107, 134, 103];
       accounts: [
         {
           name: "boltSystem";
@@ -93,7 +119,6 @@ export type World = {
         },
         {
           name: "sessionToken";
-          optional: true;
         },
       ];
       args: [

clients/bolt-sdk/src/world/transactions.ts

@@ -405,15 +405,6 @@ async function createApplySystemInstruction({
     throw new Error("No components provided");
   }
 
-  let sessionToken = session ? session.token : null;
-
-  const applyAccounts = {
-    authority: authority ?? PROGRAM_ID,
-    boltSystem: systemId,
-    sessionToken,
-    world,
-  };
-
   let remainingAccounts: web3.AccountMeta[] = [];
   let components: { id: PublicKey; pda: PublicKey }[] = [];
   for (const entity of entities) {
@@ -452,11 +443,28 @@ async function createApplySystemInstruction({
       remainingAccounts.push(account);
     }
   }
-  return program.methods
-    .apply(SerializeArgs(args))
-    .accounts(applyAccounts)
-    .remainingAccounts(remainingAccounts)
-    .instruction();
+
+  if (session)
+    return program.methods
+      .applyWithSession(SerializeArgs(args))
+      .accounts({
+        authority: authority ?? PROGRAM_ID,
+        boltSystem: systemId,
+        sessionToken: session.token,
+        world,
+      })
+      .remainingAccounts(remainingAccounts)
+      .instruction();
+  else
+    return program.methods
+      .apply(SerializeArgs(args))
+      .accounts({
+        authority: authority ?? PROGRAM_ID,
+        boltSystem: systemId,
+        world,
+      })
+      .remainingAccounts(remainingAccounts)
+      .instruction();
 }
 
 interface ApplySystemEntity {

crates/bolt-lang/attribute/bolt-program/src/lib.rs

@@ -44,20 +44,30 @@ pub fn bolt_program(args: TokenStream, input: TokenStream) -> TokenStream {
 fn modify_component_module(mut module: ItemMod, component_type: &Type) -> ItemMod {
     let (initialize_fn, initialize_struct) = generate_initialize(component_type);
     //let (apply_fn, apply_struct, apply_impl, update_fn, update_struct) = generate_instructions(component_type);
-    let (update_fn, update_struct) = generate_update(component_type);
+    let (update_fn, update_with_session_fn, update_struct, update_with_session_struct) =
+        generate_update(component_type);
 
     module.content = module.content.map(|(brace, mut items)| {
         items.extend(
-            vec![initialize_fn, initialize_struct, update_fn, update_struct]
-                .into_iter()
-                .map(|item| syn::parse2(item).unwrap())
-                .collect::<Vec<_>>(),
+            vec![
+                initialize_fn,
+                initialize_struct,
+                update_fn,
+                update_struct,
+                update_with_session_fn,
+                update_with_session_struct,
+            ]
+            .into_iter()
+            .map(|item| syn::parse2(item).unwrap())
+            .collect::<Vec<_>>(),
         );
 
         let modified_items = items
             .into_iter()
             .map(|item| match item {
-                syn::Item::Struct(mut struct_item) if struct_item.ident == "Apply" => {
+                syn::Item::Struct(mut struct_item)
+                    if struct_item.ident == "Apply" || struct_item.ident == "ApplyWithSession" =>
+                {
                     modify_apply_struct(&mut struct_item);
                     syn::Item::Struct(struct_item)
                 }
@@ -113,7 +123,7 @@ fn generate_initialize(component_type: &Type) -> (TokenStream2, TokenStream2) {
             pub fn initialize(ctx: Context<Initialize>) -> Result<()> {
                 let instruction = anchor_lang::solana_program::sysvar::instructions::get_instruction_relative(
                     0, &ctx.accounts.instruction_sysvar_account.to_account_info()
-                ).unwrap();
+                ).map_err(|_| BoltError::InvalidCaller)?;
                 if instruction.program_id != World::id() {
                     return Err(BoltError::InvalidCaller.into());
                 }
@@ -143,32 +153,45 @@ fn generate_initialize(component_type: &Type) -> (TokenStream2, TokenStream2) {
 }
 
 /// Generates the instructions and related structs to inject in the component.
-fn generate_update(component_type: &Type) -> (TokenStream2, TokenStream2) {
+fn generate_update(
+    component_type: &Type,
+) -> (TokenStream2, TokenStream2, TokenStream2, TokenStream2) {
     (
         quote! {
             #[automatically_derived]
             pub fn update(ctx: Context<Update>, data: Vec<u8>) -> Result<()> {
-                if let Some(session_token) = &ctx.accounts.session_token {
-                    if ctx.accounts.bolt_component.bolt_metadata.authority == World::id() {
-                        require!(Clock::get()?.unix_timestamp < session_token.valid_until, bolt_lang::session_keys::SessionError::InvalidToken);
-                    } else {
-                        let validity_ctx = bolt_lang::session_keys::ValidityChecker {
-                            session_token: session_token.clone(),
-                            session_signer: ctx.accounts.authority.clone(),
-                            authority: ctx.accounts.bolt_component.bolt_metadata.authority.clone(),
-                            target_program: World::id(),
-                        };
-                        require!(session_token.validate(validity_ctx)?, bolt_lang::session_keys::SessionError::InvalidToken);
-                        require_eq!(ctx.accounts.bolt_component.bolt_metadata.authority, session_token.authority, bolt_lang::session_keys::SessionError::InvalidToken);
-                    }
+                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);
+
+                // Check if the instruction is called from the world program
+                let instruction = anchor_lang::solana_program::sysvar::instructions::get_instruction_relative(
+                    0, &ctx.accounts.instruction_sysvar_account.to_account_info()
+                ).map_err(|_| BoltError::InvalidCaller)?;
+                require_eq!(instruction.program_id, World::id(), BoltError::InvalidCaller);
+
+                ctx.accounts.bolt_component.set_inner(<#component_type>::try_from_slice(&data)?);
+                Ok(())
+            }
+        },
+        quote! {
+            #[automatically_derived]
+            pub fn update_with_session(ctx: Context<UpdateWithSession>, data: Vec<u8>) -> Result<()> {
+                if ctx.accounts.bolt_component.bolt_metadata.authority == World::id() {
+                    require!(Clock::get()?.unix_timestamp < ctx.accounts.session_token.valid_until, bolt_lang::session_keys::SessionError::InvalidToken);
                 } else {
-                    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);
+                    let validity_ctx = bolt_lang::session_keys::ValidityChecker {
+                        session_token: ctx.accounts.session_token.clone(),
+                        session_signer: ctx.accounts.authority.clone(),
+                        authority: ctx.accounts.bolt_component.bolt_metadata.authority.clone(),
+                        target_program: World::id(),
+                    };
+                    require!(ctx.accounts.session_token.validate(validity_ctx)?, bolt_lang::session_keys::SessionError::InvalidToken);
+                    require_eq!(ctx.accounts.bolt_component.bolt_metadata.authority, ctx.accounts.session_token.authority, bolt_lang::session_keys::SessionError::InvalidToken);
                 }
 
                 // Check if the instruction is called from the world program
                 let instruction = anchor_lang::solana_program::sysvar::instructions::get_instruction_relative(
                     0, &ctx.accounts.instruction_sysvar_account.to_account_info()
-                ).unwrap();
+                ).map_err(|_| BoltError::InvalidCaller)?;
                 require_eq!(instruction.program_id, World::id(), BoltError::InvalidCaller);
 
                 ctx.accounts.bolt_component.set_inner(<#component_type>::try_from_slice(&data)?);
@@ -179,6 +202,18 @@ fn generate_update(component_type: &Type) -> (TokenStream2, TokenStream2) {
             #[automatically_derived]
             #[derive(Accounts)]
             pub struct Update<'info> {
+                #[account(mut)]
+                pub bolt_component: Account<'info, #component_type>,
+                #[account()]
+                pub authority: Signer<'info>,
+                #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
+                pub instruction_sysvar_account: UncheckedAccount<'info>
+            }
+        },
+        quote! {
+            #[automatically_derived]
+            #[derive(Accounts)]
+            pub struct UpdateWithSession<'info> {
                 #[account(mut)]
                 pub bolt_component: Account<'info, #component_type>,
                 #[account()]
@@ -186,7 +221,7 @@ fn generate_update(component_type: &Type) -> (TokenStream2, TokenStream2) {
                 #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
                 pub instruction_sysvar_account: UncheckedAccount<'info>,
                 #[account(constraint = session_token.to_account_info().owner == &bolt_lang::session_keys::ID)]
-                pub session_token: Option<Account<'info, bolt_lang::session_keys::SessionToken>>,
+                pub session_token: Account<'info, bolt_lang::session_keys::SessionToken>,
             }
         },
     )

crates/programs/bolt-component/src/lib.rs

@@ -14,6 +14,10 @@ pub mod bolt_component {
         Ok(())
     }
 
+    pub fn update_with_session(_ctx: Context<UpdateWithSession>, _data: Vec<u8>) -> Result<()> {
+        Ok(())
+    }
+
     #[derive(Accounts)]
     pub struct Update<'info> {
         #[account(mut)]
@@ -25,8 +29,21 @@ pub mod bolt_component {
         #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
         /// CHECK: The instruction sysvar
         pub instruction_sysvar_account: AccountInfo<'info>,
+    }
+
+    #[derive(Accounts)]
+    pub struct UpdateWithSession<'info> {
+        #[account(mut)]
+        /// CHECK: The component to update
+        pub bolt_component: UncheckedAccount<'info>,
+        #[account()]
+        /// CHECK: The authority of the component
+        pub authority: Signer<'info>,
+        #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
+        /// CHECK: The instruction sysvar
+        pub instruction_sysvar_account: AccountInfo<'info>,
         #[account()]
-        pub session_token: Option<UncheckedAccount<'info>>,
+        pub session_token: UncheckedAccount<'info>,
     }
 }
 
@@ -65,11 +82,21 @@ pub trait CpiContextBuilder<'info>: ToAccountMetas + ToAccountInfos<'info> + Siz
 #[cfg(feature = "cpi")]
 impl<'info> CpiContextBuilder<'info> for cpi::accounts::Update<'info> {
     fn build_cpi_context(
-        mut self,
+        self,
+        program: AccountInfo<'info>,
+    ) -> CpiContext<'info, 'info, 'info, 'info, Self> {
+        let cpi_program = program.to_account_info();
+        CpiContext::new(cpi_program, self)
+    }
+}
+
+#[cfg(feature = "cpi")]
+impl<'info> CpiContextBuilder<'info> for cpi::accounts::UpdateWithSession<'info> {
+    fn build_cpi_context(
+        self,
         program: AccountInfo<'info>,
     ) -> CpiContext<'info, 'info, 'info, 'info, Self> {
         let cpi_program = program.to_account_info();
-        self.session_token = Some(self.session_token.unwrap_or(program.to_account_info()));
         CpiContext::new(cpi_program, self)
     }
 }

crates/programs/world/src/lib.rs

@@ -1,6 +1,7 @@
 #![allow(clippy::manual_unwrap_or_default)]
 use anchor_lang::prelude::*;
 use bolt_component::CpiContextBuilder;
+use error::WorldError;
 use std::collections::BTreeSet;
 
 #[cfg(not(feature = "no-entrypoint"))]
@@ -268,58 +269,70 @@ pub mod world {
         ctx: Context<'_, '_, '_, 'info, Apply<'info>>,
         args: Vec<u8>,
     ) -> Result<()> {
-        if !ctx.accounts.authority.is_signer
-            && ctx.accounts.authority.key != &ID
-            && ctx.accounts.session_token.is_none()
-        {
-            return Err(WorldError::InvalidAuthority.into());
-        }
-        if !ctx.accounts.world.permissionless
-            && !ctx
-                .accounts
-                .world
-                .systems()
-                .approved_systems
-                .contains(&ctx.accounts.bolt_system.key())
-        {
-            return Err(WorldError::SystemNotApproved.into());
+        let (pairs, results) = apply_impl(
+            &ctx.accounts.authority,
+            &ctx.accounts.world,
+            &ctx.accounts.bolt_system,
+            ctx.accounts.build(),
+            args,
+            ctx.remaining_accounts.to_vec(),
+        )?;
+        for ((program, component), result) in pairs.into_iter().zip(results.into_iter()) {
+            bolt_component::cpi::update(
+                build_update_context(
+                    program,
+                    component,
+                    ctx.accounts.authority.clone(),
+                    ctx.accounts.instruction_sysvar_account.clone(),
+                ),
+                result,
+            )?;
         }
+        Ok(())
+    }
 
-        let mut remaining_accounts: Vec<AccountInfo<'info>> = ctx.remaining_accounts.to_vec();
+    #[derive(Accounts)]
+    pub struct Apply<'info> {
+        /// CHECK: bolt system program check
+        #[account()]
+        pub bolt_system: UncheckedAccount<'info>,
+        /// CHECK: authority check
+        #[account()]
+        pub authority: Signer<'info>,
+        /// CHECK: instruction sysvar check
+        #[account(address = anchor_lang::solana_program::sysvar::instructions::id())]
+        pub instruction_sysvar_account: UncheckedAccount<'info>,
+        #[account()]
+        pub world: Account<'info, World>,
+    }
 
-        let mut pairs = Vec::new();
-        while remaining_accounts.len() >= 2 {
-            let program = remaining_accounts.remove(0);
-            if program.key() == ID {
-                break;
-            }
-            let component = remaining_accounts.remove(0);
-            pairs.push((program, component));
+    impl<'info> Apply<'info> {
+        pub fn build(
+            &self,
+        ) -> CpiContext<'_, '_, '_, 'info, bolt_system::cpi::accounts::BoltExecute<'info>> {
+            let cpi_program = self.bolt_system.to_account_info();
+            let cpi_accounts = bolt_system::cpi::accounts::BoltExecute {
+                authority: self.authority.to_account_info(),
+            };
+            CpiContext::new(cpi_program, cpi_accounts)
         }
+    }
 
-        let mut components_accounts = pairs
-            .iter()
-            .map(|(_, component)| component)
-            .cloned()
-            .collect::<Vec<_>>();
-        components_accounts.append(&mut remaining_accounts);
-        let remaining_accounts = components_accounts;
-
-        let results = bolt_system::cpi::bolt_execute(
-            ctx.accounts
-                .build()
-                .with_remaining_accounts(remaining_accounts),
+    pub fn apply_with_session<'info>(
+        ctx: Context<'_, '_, '_, 'info, ApplyWithSession<'info>>,
+        args: Vec<u8>,
+    ) -> Result<()> {
+        let (pairs, results) = apply_impl(
+            &ctx.accounts.authority,
+            &ctx.accounts.world,
+            &ctx.accounts.bolt_system,
+            ctx.accounts.build(),
             args,
-        )?
-        .get();
-
-        if results.len() != pairs.len() {
-            return Err(WorldError::InvalidSystemOutput.into());
-        }
-
+            ctx.remaining_accounts.to_vec(),
+        )?;
         for ((program, component), result) in pairs.into_iter().zip(results.into_iter()) {
-            bolt_component::cpi::update(
-                build_update_context(
+            bolt_component::cpi::update_with_session(
+                build_update_context_with_session(
                     program,
                     component,
                     ctx.accounts.authority.clone(),
@@ -333,7 +346,7 @@ pub mod world {
     }
 
     #[derive(Accounts)]
-    pub struct Apply<'info> {
+    pub struct ApplyWithSession<'info> {
         /// CHECK: bolt system program check
         #[account()]
         pub bolt_system: UncheckedAccount<'info>,
@@ -346,10 +359,10 @@ pub mod world {
         #[account()]
         pub world: Account<'info, World>,
         #[account()]
-        pub session_token: Option<UncheckedAccount<'info>>,
+        pub session_token: UncheckedAccount<'info>,
     }
 
-    impl<'info> Apply<'info> {
+    impl<'info> ApplyWithSession<'info> {
         pub fn build(
             &self,
         ) -> CpiContext<'_, '_, '_, 'info, bolt_system::cpi::accounts::BoltExecute<'info>> {
@@ -362,6 +375,57 @@ pub mod world {
     }
 }
 
+#[allow(clippy::type_complexity)]
+fn apply_impl<'info>(
+    authority: &Signer<'info>,
+    world: &Account<'info, World>,
+    bolt_system: &UncheckedAccount<'info>,
+    cpi_context: CpiContext<'_, '_, '_, 'info, bolt_system::cpi::accounts::BoltExecute<'info>>,
+    args: Vec<u8>,
+    mut remaining_accounts: Vec<AccountInfo<'info>>,
+) -> Result<(Vec<(AccountInfo<'info>, AccountInfo<'info>)>, Vec<Vec<u8>>)> {
+    if !authority.is_signer && authority.key != &ID {
+        return Err(WorldError::InvalidAuthority.into());
+    }
+    if !world.permissionless
+        && !world
+            .systems()
+            .approved_systems
+            .contains(&bolt_system.key())
+    {
+        return Err(WorldError::SystemNotApproved.into());
+    }
+
+    let mut pairs = Vec::new();
+    while remaining_accounts.len() >= 2 {
+        let program = remaining_accounts.remove(0);
+        if program.key() == ID {
+            break;
+        }
+        let component = remaining_accounts.remove(0);
+        pairs.push((program, component));
+    }
+
+    let mut components_accounts = pairs
+        .iter()
+        .map(|(_, component)| component)
+        .cloned()
+        .collect::<Vec<_>>();
+    components_accounts.append(&mut remaining_accounts);
+    let remaining_accounts = components_accounts;
+
+    let results = bolt_system::cpi::bolt_execute(
+        cpi_context.with_remaining_accounts(remaining_accounts),
+        args,
+    )?
+    .get();
+
+    if results.len() != pairs.len() {
+        return Err(WorldError::InvalidSystemOutput.into());
+    }
+    Ok((pairs, results))
+}
+
 #[derive(Accounts)]
 pub struct InitializeRegistry<'info> {
     #[account(init, payer = payer, space = Registry::size(), seeds = [Registry::seed()], bump)]
@@ -598,16 +662,35 @@ pub fn build_update_context<'info>(
     bolt_component: AccountInfo<'info>,
     authority: Signer<'info>,
     instruction_sysvar_account: UncheckedAccount<'info>,
-    session_token: Option<UncheckedAccount<'info>>,
 ) -> CpiContext<'info, 'info, 'info, 'info, bolt_component::cpi::accounts::Update<'info>> {
     let authority = authority.to_account_info();
     let instruction_sysvar_account = instruction_sysvar_account.to_account_info();
     let cpi_program = component_program;
-    let session_token = session_token.map(|x| x.to_account_info());
     bolt_component::cpi::accounts::Update {
         bolt_component,
         authority,
         instruction_sysvar_account,
+    }
+    .build_cpi_context(cpi_program)
+}
+
+/// Builds the context for updating a component.
+pub fn build_update_context_with_session<'info>(
+    component_program: AccountInfo<'info>,
+    bolt_component: AccountInfo<'info>,
+    authority: Signer<'info>,
+    instruction_sysvar_account: UncheckedAccount<'info>,
+    session_token: UncheckedAccount<'info>,
+) -> CpiContext<'info, 'info, 'info, 'info, bolt_component::cpi::accounts::UpdateWithSession<'info>>
+{
+    let authority = authority.to_account_info();
+    let instruction_sysvar_account = instruction_sysvar_account.to_account_info();
+    let cpi_program = component_program;
+    let session_token = session_token.to_account_info();
+    bolt_component::cpi::accounts::UpdateWithSession {
+        bolt_component,
+        authority,
+        instruction_sysvar_account,
         session_token,
     }
     .build_cpi_context(cpi_program)

tests/low-level/session.ts

@@ -84,7 +84,7 @@ export function session(framework) {
         );
 
       const instruction = await framework.worldProgram.methods
-        .apply(SerializeArgs())
+        .applyWithSession(SerializeArgs())
         .accounts({
           authority: sessionSigner.publicKey,
           boltSystem: framework.systemFly.programId,
@@ -170,7 +170,7 @@ export function session(framework) {
         );
 
       const instruction = await framework.worldProgram.methods
-        .apply(SerializeArgs())
+        .applyWithSession(SerializeArgs())
         .accounts({
           authority: sessionSigner.publicKey,
           boltSystem: framework.systemFly.programId,