🚧 Getting session keys on the new system execution framework

magicblock-labs · Jan 17, 2025 · df85aac · df85aac
1 parent 65c1eae
commit df85aac
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 98 deletions.
diff --git a/Anchor.toml b/Anchor.toml
@@ -41,8 +41,8 @@ echo "Waiting for 2 seconds..."
 sleep 2
 echo "Running low level API tests..."
 yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/bolt.low-level.api.ts
-# echo "Running intermediate level API tests..."
-# yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/bolt.intermediate-level.api.ts
+echo "Running intermediate level API tests..."
+yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/bolt.intermediate-level.api.ts
 """
 
 [test]

diff --git a/crates/bolt-lang/attribute/system-input/src/lib.rs b/crates/bolt-lang/attribute/system-input/src/lib.rs
@@ -17,23 +17,10 @@ use syn::{parse_macro_input, Fields, ItemStruct, Lit, Meta, NestedMeta};
 ///
 /// ```
 #[proc_macro_attribute]
-pub fn system_input(attr: TokenStream, item: TokenStream) -> TokenStream {
+pub fn system_input(_attr: TokenStream, item: TokenStream) -> TokenStream {
     // Parse the input TokenStream (the struct) into a Rust data structure
     let input = parse_macro_input!(item as ItemStruct);
 
-    let attr = parse_macro_input!(attr as syn::AttributeArgs);
-
-    let has_session_key = attr.iter().any(|meta| {
-        if let syn::NestedMeta::Meta(syn::Meta::Path(path)) = meta {
-            path.segments
-                .first()
-                .map(|segment| segment.ident == "session_key")
-                .unwrap_or(false)
-        } else {
-            false
-        }
-    });
-
     // Ensure the struct has named fields
     let fields = match &input.fields {
         Fields::Named(fields) => &fields.named,
@@ -84,32 +71,15 @@ pub fn system_input(attr: TokenStream, item: TokenStream) -> TokenStream {
         }
     });
 
-    let (derives, optional_fields) = if has_session_key {
-        (
-            quote! {
-                #[derive(Accounts, Session)]
-            },
-            quote! {
-                #[session(signer = authority, authority = position.bolt_metadata.authority.key())] // FIXME: position is hardcoded.
-                pub session_token: Option<Account<'info, SessionToken>>
-            },
-        )
-    } else {
-        (
-            quote! {
-                #[derive(Accounts)]
-            },
-            Default::default(),
-        )
-    };
     // Generate the new struct with the Accounts derive and transformed fields
     let output_struct = quote! {
-        #derives
+        #[derive(Accounts, Session)]
         pub struct #name<'info> {
             #(#transformed_fields)*
             #[account()]
             pub authority: Signer<'info>,
-            #optional_fields
+            #[session(signer = authority, authority = authority.key())] // FIXME: authority = authority.key() is hardcoded.
+            pub session_token: Option<Account<'info, SessionToken>>
         }
     };
 
@@ -152,6 +122,7 @@ pub fn system_input(attr: TokenStream, item: TokenStream) -> TokenStream {
             fn try_from<'a, 'b>(context: &Context<'a, 'b, 'info, 'info, VariadicBoltComponents<'info>>) -> Result<Self> {
                 Ok(Self {
                     authority: context.accounts.authority.clone(),
+                    session_token: context.accounts.session_token.clone(),
                     #(#try_from_fields)*
                 })
             }
@@ -170,6 +141,8 @@ pub fn system_input(attr: TokenStream, item: TokenStream) -> TokenStream {
         pub struct VariadicBoltComponents<'info> {
             #[account()]
             pub authority: Signer<'info>,
+            #[account()]
+            pub session_token: Option<Account<'info, SessionToken>>
         }
     };
 

diff --git a/crates/bolt-lang/attribute/system/src/lib.rs b/crates/bolt-lang/attribute/system/src/lib.rs
@@ -7,9 +7,7 @@ use syn::{
 };
 
 #[derive(Default)]
-struct SystemTransform {
-    has_session_key: bool,
-}
+struct SystemTransform;
 
 #[derive(Default)]
 struct Extractor {
@@ -36,27 +34,13 @@ struct Extractor {
 /// }
 /// ```
 #[proc_macro_attribute]
-pub fn system(attr: TokenStream, item: TokenStream) -> TokenStream {
+pub fn system(_attr: TokenStream, item: TokenStream) -> TokenStream {
     let mut ast = parse_macro_input!(item as ItemMod);
-    let attr = parse_macro_input!(attr as syn::AttributeArgs);
-
-    let has_session_key = attr.iter().any(|meta| {
-        if let syn::NestedMeta::Meta(syn::Meta::Path(path)) = meta {
-            path.segments
-                .first()
-                .map(|segment| segment.ident == "session_key")
-                .unwrap_or(false)
-        } else {
-            false
-        }
-    });
 
-    if has_session_key {
-        if let Some(content) = &mut ast.content {
-            content.1.push(syn::Item::Use(syn::parse_quote! {
-                use bolt_lang::session_keys::{Session, SessionToken};
-            }));
-        }
+    if let Some(content) = &mut ast.content {
+        content.1.push(syn::Item::Use(syn::parse_quote! {
+            use bolt_lang::session_keys::{Session, SessionToken};
+        }));
     }
 
     // Extract the number of components from the module
@@ -70,9 +54,7 @@ pub fn system(attr: TokenStream, item: TokenStream) -> TokenStream {
             SystemTransform::add_variadic_execute_function(items);
         }
 
-        let mut transform = SystemTransform {
-            has_session_key,
-        };
+        let mut transform = SystemTransform;
         transform.visit_item_mod_mut(&mut ast);
 
         // Add `#[program]` macro and try_to_vec implementation
@@ -147,11 +129,6 @@ impl VisitMut for SystemTransform {
     // Modify the return type of the system function to Result<Vec<u8>,*>
     fn visit_item_fn_mut(&mut self, item_fn: &mut ItemFn) {
         if item_fn.sig.ident == "execute" {
-            if self.has_session_key {
-                // item_fn.attrs.push(parse_quote! {
-                //     #[session_auth_or(ctx.accounts.position.bolt_metadata.authority.key() == ctx.accounts.authority.key(), SessionError::InvalidToken)] // FIXME: Position is hardcoded here
-                //  });
-            }
             // Modify the return type to Result<Vec<u8>> if necessary
             if let ReturnType::Type(_, type_box) = &item_fn.sig.output {
                 if let Type::Path(type_path) = &**type_box {
@@ -187,9 +164,7 @@ impl VisitMut for SystemTransform {
                         .iter_mut()
                         .find(|attr| attr.path.is_ident("system_input"))
                     {
-                        if self.has_session_key {
-                            attr.tokens.append_all(quote! { (session_key) });
-                        }
+                        attr.tokens.append_all(quote! { (session_key) });
                     }
                     if item_struct
                         .attrs

diff --git a/crates/programs/world/src/lib.rs b/crates/programs/world/src/lib.rs
@@ -350,6 +350,7 @@ pub mod world {
         pub fn build(
             &self,
         ) -> CpiContext<'_, '_, '_, 'info, bolt_system::cpi::accounts::SetData<'info>> {
+            msg!("session_token: {:?}", self.session_token);
             bolt_system::cpi::accounts::SetData {
                 authority: self.authority.to_account_info(),
                 session_token: self.session_token.as_ref().map(|x| x.to_account_info()),

diff --git a/tests/bolt.low-level.api.ts b/tests/bolt.low-level.api.ts
@@ -6,7 +6,6 @@ import { type SystemSimpleMovement } from "../target/types/system_simple_movemen
 import { type SystemFly } from "../target/types/system_fly";
 import { type SystemApplyVelocity } from "../target/types/system_apply_velocity";
 import { type World } from "../target/types/world";
-import { type World } from "../target/types/world";
 import { expect } from "chai";
 import BN from "bn.js";
 import {
@@ -512,34 +511,6 @@ describe("bolt", () => {
     expect(position.z.toNumber()).to.equal(0);
   });
 
-  it("Apply Simple Movement System (Up) on Entity 1", async () => {
-    const instruction = await worldProgram.methods
-      .apply(SerializeArgs({ direction: Direction.Up }))
-      .accounts({
-        authority: provider.wallet.publicKey,
-        boltSystem: exampleSystemSimpleMovement,
-        boltComponent: componentPositionEntity1Pda,
-        componentProgram: exampleComponentPosition.programId,
-        world: worldPda,
-      })
-      .instruction();
-
-    const transaction = new anchor.web3.Transaction().add(instruction);
-    const signature = await provider.sendAndConfirm(transaction);
-    console.log(
-      "Apply Simple Movement System (Up) on Entity 1 signature: ",
-      signature,
-    );
-
-    const position = await exampleComponentPosition.account.position.fetch(
-      componentPositionEntity1Pda,
-    );
-    logPosition("Movement System: Entity 1", position);
-    expect(position.x.toNumber()).to.equal(0);
-    expect(position.y.toNumber()).to.equal(2);
-    expect(position.z.toNumber()).to.equal(0);
-  });
-
   it("Apply Simple Movement System (Right) on Entity 1", async () => {
     const instruction = await worldProgram.methods
       .apply(SerializeArgs({ direction: Direction.Right }))
@@ -620,6 +591,7 @@ describe("bolt", () => {
         authority: provider.wallet.publicKey,
         boltSystem: exampleSystemApplyVelocity,
         world: worldPda,
+        sessionToken: null
       })
       .remainingAccounts([
         {
@@ -673,6 +645,7 @@ describe("bolt", () => {
         authority: provider.wallet.publicKey,
         boltSystem: exampleSystemApplyVelocity,
         world: worldPda,
+        sessionToken: null
       })
       .remainingAccounts([
         {