From 0e11c24dae5b08a199fbe6c2aac700c1dedec439 Mon Sep 17 00:00:00 2001
From: Naman Anand <iamnamananand996@gmail.com>
Date: Sun, 13 Oct 2024 15:28:35 +0530
Subject: [PATCH 1/3] fix: :bug: Components CPI caller check

---
 programs/bolt-component/src/lib.rs | 53 +++++++++++++++++++++---------
 1 file changed, 37 insertions(+), 16 deletions(-)

diff --git a/programs/bolt-component/src/lib.rs b/programs/bolt-component/src/lib.rs
index adafe9ce..ac449f1c 100644
--- 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,
+}

From f05556d1d28d7b4ea4d8704b6af32f4289983608 Mon Sep 17 00:00:00 2001
From: iamnamananand996 <iamnamananand996@gmail.com>
Date: Wed, 6 Nov 2024 20:19:20 +0530
Subject: [PATCH 2/3] fix: :bug:  test fix

---
 tests/bolt.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/bolt.ts b/tests/bolt.ts
index e7f907cc..37e07115 100644
--- a/tests/bolt.ts
+++ b/tests/bolt.ts
@@ -233,7 +233,7 @@ describe("bolt", () => {
     entity5Pda = addEntity.entityPda; // Saved for later
   });
 
-  it("Initialize Original Component on Entity 1, trough the world instance", async () => {
+  it("Initialize Original Component on Entity 1, through the world instance", async () => {
     const initializeComponent = await InitializeComponent({
       payer: provider.wallet.publicKey,
       entity: entity1Pda,
@@ -243,7 +243,7 @@ describe("bolt", () => {
     await provider.sendAndConfirm(initializeComponent.transaction);
   });
 
-  it("Initialize Original Component on Entity 2, trough the world instance", async () => {
+  it("Initialize Original Component on Entity 2, through the world instance", async () => {
     const initializeComponent = await InitializeComponent({
       payer: provider.wallet.publicKey,
       entity: entity2Pda,

From baaedb3e32ae5f238216be31f9fa684996d7e626 Mon Sep 17 00:00:00 2001
From: iamnamananand996 <iamnamananand996@gmail.com>
Date: Wed, 6 Nov 2024 21:33:42 +0530
Subject: [PATCH 3/3] fix: :bug:  update failed test cases

---
 tests/bolt.ts | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/tests/bolt.ts b/tests/bolt.ts
index 37e07115..b00e641f 100644
--- a/tests/bolt.ts
+++ b/tests/bolt.ts
@@ -240,7 +240,11 @@ describe("bolt", () => {
       seed: "origin-component",
       componentId: boltComponentProgram.programId,
     });
-    await provider.sendAndConfirm(initializeComponent.transaction);
+    try {
+      await provider.sendAndConfirm(initializeComponent.transaction);
+    } catch (error) {
+      expect(error.message).to.contain("Error Code: MustBeCalledViaCpi");
+    }
   });
 
   it("Initialize Original Component on Entity 2, through the world instance", async () => {
@@ -250,7 +254,12 @@ describe("bolt", () => {
       seed: "origin-component",
       componentId: boltComponentProgram.programId,
     });
-    await provider.sendAndConfirm(initializeComponent.transaction);
+
+    try {
+      await provider.sendAndConfirm(initializeComponent.transaction);
+    } catch (error) {
+      expect(error.message).to.contain("Error Code: MustBeCalledViaCpi");
+    }
   });
 
   it("Initialize Position Component on Entity 1", async () => {