✨ Session Keys

magicblock-labs · Jan 24, 2025 · 2c71613 · 2c71613
1 parent ede6485
commit 2c71613
Show file tree

Hide file tree

Showing 16 changed files with 1,036 additions and 685 deletions.
diff --git a/clients/bolt-sdk/src/index.ts b/clients/bolt-sdk/src/index.ts
@@ -76,6 +76,26 @@ export function FindEntityPda({
   )[0];
 }
 
+export function FindSessionTokenPda({
+  targetProgram,
+  sessionSigner,
+  authority,
+}: {
+  targetProgram: PublicKey;
+  sessionSigner: PublicKey;
+  authority: PublicKey;
+}) {
+  return PublicKey.findProgramAddressSync(
+    [
+      Buffer.from("session_token"),
+      targetProgram.toBytes(),
+      sessionSigner.toBytes(),
+      authority.toBytes(),
+    ],
+    SessionProgram.programId,
+  )[0];
+}
+
 // TODO: seed must be Uint8Array like the other FindPda functions
 export function FindComponentPda({
   componentId,

diff --git a/clients/bolt-sdk/src/world/transactions.ts b/clients/bolt-sdk/src/world/transactions.ts
@@ -11,11 +11,14 @@ import {
   SerializeArgs,
   SYSVAR_INSTRUCTIONS_PUBKEY,
   World,
+  SessionProgram,
+  FindSessionTokenPda,
 } from "../index";
 import BN from "bn.js";
 import type web3 from "@solana/web3.js";
 import {
   type Connection,
+  Keypair,
   type PublicKey,
   Transaction,
   type TransactionInstruction,
@@ -47,6 +50,49 @@ export async function InitializeRegistry({
   };
 }
 
+export async function CreateSession({
+  sessionSigner,
+  authority,
+  targetProgram,
+  topUp,
+  validity,
+}: {
+  sessionSigner?: Keypair;
+  authority: PublicKey;
+  targetProgram: PublicKey;
+  topUp?: boolean;
+  validity?: BN;
+}): Promise<{
+  instruction: TransactionInstruction;
+  transaction: Transaction;
+  sessionToken: PublicKey;
+  sessionSigner: Keypair;
+}> {
+  sessionSigner = sessionSigner ?? Keypair.generate();
+  const sessionToken = FindSessionTokenPda({
+    targetProgram,
+    sessionSigner: sessionSigner.publicKey,
+    authority,
+  });
+  topUp = topUp ?? false;
+  let instruction = await SessionProgram.methods
+    .createSession(topUp, validity ?? null)
+    .accounts({
+      sessionSigner: sessionSigner.publicKey,
+      authority,
+      targetProgram,
+      sessionToken,
+    })
+    .instruction();
+  const transaction = new Transaction().add(instruction);
+  return {
+    instruction,
+    transaction,
+    sessionToken,
+    sessionSigner,
+  };
+}
+
 /**
  * Create the transaction to Initialize a new world
  * @param payer
@@ -438,13 +484,15 @@ export async function ApplySystem({
   world,
   extraAccounts,
   args,
+  sessionToken,
 }: {
   authority: PublicKey;
   systemId: PublicKey;
   entities: ApplySystemEntity[];
   world: PublicKey;
   extraAccounts?: web3.AccountMeta[];
   args?: object;
+  sessionToken?: PublicKey;
 }): Promise<{ instruction: TransactionInstruction; transaction: Transaction }> {
   const instruction = await createApplySystemInstruction({
     authority,
@@ -453,6 +501,7 @@ export async function ApplySystem({
     world,
     extraAccounts,
     args,
+    sessionToken,
   });
   const transaction = new Transaction().add(instruction);
   return {

diff --git a/crates/bolt-lang/attribute/bolt-program/src/lib.rs b/crates/bolt-lang/attribute/bolt-program/src/lib.rs
@@ -150,14 +150,18 @@ fn generate_update(component_type: &Type) -> (TokenStream2, TokenStream2) {
             pub fn update(ctx: Context<Update>, data: Vec<u8>) -> Result<()> {
                 // TODO: Should we also check if the session token authority can be the World::id?
                 if let Some(session_token) = &ctx.accounts.session_token {
-                    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: crate::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);
+                    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: crate::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);
+                    }
                 } 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);
                 }

diff --git a/crates/programs/bolt-component/src/lib.rs b/crates/programs/bolt-component/src/lib.rs
@@ -17,6 +17,7 @@ pub mod bolt_component {
     #[derive(Accounts)]
     pub struct Update<'info> {
         #[account(mut)]
+        /// CHECK: The component to update
         pub bolt_component: UncheckedAccount<'info>,
         #[account()]
         /// CHECK: The authority of the component
@@ -34,6 +35,7 @@ pub struct Initialize<'info> {
     #[account(mut)]
     pub payer: Signer<'info>,
     #[account(mut)]
+    /// CHECK: The component to initialize
     pub data: UncheckedAccount<'info>,
     #[account()]
     /// CHECK: A generic entity account

diff --git a/scripts/lint.sh b/scripts/lint.sh
@@ -3,11 +3,11 @@ SCRIPT_DIR=$(dirname "$0")
 PROJECT_DIR="$SCRIPT_DIR/.."
 pushd "$PROJECT_DIR"
 echo "### Checking formatting..."
-cargo fmt -- --check --verbose
+cargo fmt -- --verbose
 
 echo "### Checking clippy..."
-cargo clippy -- --deny=warnings
+cargo clippy --fix -- --deny=warnings
 
 echo "### Checking yarn lint..."
-yarn lint
+yarn lint --write
 popd
diff --git a/tests/framework.ts b/tests/framework.ts
@@ -71,20 +71,14 @@ export class Framework {
   entity1Pda: PublicKey;
   entity2Pda: PublicKey;
   entity4Pda: PublicKey;
-  entity5Pda: PublicKey;
 
   componentPositionEntity1Pda: PublicKey;
   componentVelocityEntity1Pda: PublicKey;
 
   componentPositionEntity4Pda: PublicKey;
-  componentPositionEntity5Pda: PublicKey;
-
-  sessionSigner: Keypair;
-  sessionToken: PublicKey;
 
   constructor() {
     this.secondAuthority = Keypair.generate().publicKey;
-    this.sessionSigner = Keypair.generate();
     this.worldProgram = anchor.workspace.World;
     this.exampleComponentPosition = anchor.workspace.Position;
     this.exampleComponentVelocity = anchor.workspace.Velocity;

diff --git a/tests/intermediate-level/acceleration.ts b/tests/intermediate-level/acceleration.ts
@@ -13,10 +13,14 @@ export function acceleration(framework) {
         componentId: framework.exampleComponentPosition.programId,
       });
 
-      await framework.provider.sendAndConfirm(delegateComponent.transaction, [], {
-        skipPreflight: true,
-        commitment: "confirmed",
-      });
+      await framework.provider.sendAndConfirm(
+        delegateComponent.transaction,
+        [],
+        {
+          skipPreflight: true,
+          commitment: "confirmed",
+        },
+      );
       const acc = await framework.provider.connection.getAccountInfo(
         delegateComponent.componentPda,
       );