fix: 🐛 Components CPI caller check

magicblock-labs · Nov 3, 2024 · 0e11c24 · 0e11c24
1 parent 20f7ffa
commit 0e11c24
Showing 1 changed file with 37 additions and 16 deletions.
diff --git a/programs/bolt-component/src/lib.rs b/programs/bolt-component/src/lib.rs
@@ -1,4 +1,5 @@
 use anchor_lang::prelude::*;
+use anchor_lang::solana_program::sysvar::instructions::get_instruction_relative;
 
 declare_id!("CmP2djJgABZ4cRokm4ndxuq6LerqpNHLBsaUv2XKEJua");
 
@@ -7,15 +8,22 @@ pub mod bolt_component {
     use super::*;
 
     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();
-        if instruction.program_id == id() {
-            panic!("The instruction must be called from a CPI");
+        // Check if the program is called via CPI
+        match get_instruction_relative(
+            -1,
+            &ctx.accounts.instruction_sysvar_account.to_account_info(),
+        ) {
+            Ok(instruction) => {
+                // Verify the caller's program ID, if necessary
+                if instruction.program_id != ctx.accounts.authority.key() {
+                    return Err(ErrorCode::UnauthorizedCaller.into());
+                }
+            }
+            Err(_) => {
+                return Err(ErrorCode::MustBeCalledViaCpi.into());
+            }
         }
+
         ctx.accounts.data.bolt_metadata.authority = *ctx.accounts.authority.key;
         Ok(())
     }
@@ -47,14 +55,19 @@ pub mod bolt_component {
     }
 
     pub fn update(ctx: Context<Update>, _data: Vec<u8>) -> Result<()> {
-        let instruction =
-            anchor_lang::solana_program::sysvar::instructions::get_instruction_relative(
-                0,
-                &ctx.accounts.instruction_sysvar_account.to_account_info(),
-            )
-            .unwrap();
-        if instruction.program_id == id() {
-            panic!("The instruction must be called from a CPI");
+        // Same CPI check as in initialize
+        match get_instruction_relative(
+            -1,
+            &ctx.accounts.instruction_sysvar_account.to_account_info(),
+        ) {
+            Ok(instruction) => {
+                if instruction.program_id != ctx.accounts.authority.key() {
+                    return Err(ErrorCode::UnauthorizedCaller.into());
+                }
+            }
+            Err(_) => {
+                return Err(ErrorCode::MustBeCalledViaCpi.into());
+            }
         }
         Ok(())
     }
@@ -118,3 +131,11 @@ pub struct Position {
 pub struct BoltMetadata {
     pub authority: Pubkey,
 }
+
+#[error_code]
+pub enum ErrorCode {
+    #[msg("The instruction must be called from a CPI")]
+    MustBeCalledViaCpi,
+    #[msg("Unauthorized caller program")]
+    UnauthorizedCaller,
+}