Skip to content

Instantly share code, notes, and snippets.

@5hanth
Created January 30, 2026 13:44
Show Gist options
  • Select an option

  • Save 5hanth/d8487215c71138a1b03b4b32920501da to your computer and use it in GitHub Desktop.

Select an option

Save 5hanth/d8487215c71138a1b03b4b32920501da to your computer and use it in GitHub Desktop.
PR #120 audit fixes - changes from 0075e5c to 8390806 (11 commits)
diff --git a/apps/contracts/stake/Anchor.toml b/apps/contracts/stake/Anchor.toml
index 9ef53fe..63a350f 100644
--- a/apps/contracts/stake/Anchor.toml
+++ b/apps/contracts/stake/Anchor.toml
@@ -18,6 +18,11 @@ url = "https://api.apr.dev"
cluster = "localnet"
wallet = "~/.config/solana/id.json"
+# Test configuration for H-04 (upgrade authority check)
+# We need proper upgradeable deployment, not genesis loading
+[test]
+startup_wait = 5000
+
[scripts]
create-test-mints = "bun run migrations/create-test-mints.ts"
deploy-and-initialize = "bun run migrations/deploy-and-initialize.ts"
diff --git a/apps/contracts/stake/Cargo.lock b/apps/contracts/stake/Cargo.lock
index 1984448..3a04a1a 100644
--- a/apps/contracts/stake/Cargo.lock
+++ b/apps/contracts/stake/Cargo.lock
@@ -2647,6 +2647,7 @@ version = "0.1.0"
dependencies = [
"anchor-lang",
"anchor-spl",
+ "bincode",
]
[[package]]
diff --git a/apps/contracts/stake/migrations/setup-test-authority.ts b/apps/contracts/stake/migrations/setup-test-authority.ts
new file mode 100644
index 0000000..67df6df
--- /dev/null
+++ b/apps/contracts/stake/migrations/setup-test-authority.ts
@@ -0,0 +1,95 @@
+/**
+ * Setup script to ensure the stake program has the correct upgrade authority
+ * for H-04 tests to pass.
+ *
+ * This script should be run before tests to verify/fix the program's upgrade authority.
+ * Run with: npx ts-node migrations/setup-test-authority.ts
+ */
+
+import * as anchor from "@coral-xyz/anchor";
+import { PublicKey, Connection } from "@solana/web3.js";
+import { execSync } from "child_process";
+
+const BPF_LOADER_UPGRADEABLE = new PublicKey("BPFLoaderUpgradeab1e11111111111111111111111");
+
+async function main() {
+ // Get the provider
+ const provider = anchor.AnchorProvider.env();
+ anchor.setProvider(provider);
+
+ const program = anchor.workspace.StakeProgram;
+ const programId = program.programId;
+ const walletPubkey = provider.wallet.publicKey;
+
+ console.log("=== Setup Test Authority ===");
+ console.log("Program ID:", programId.toBase58());
+ console.log("Wallet:", walletPubkey.toBase58());
+
+ // Derive program data PDA
+ const [programDataPda] = PublicKey.findProgramAddressSync(
+ [programId.toBuffer()],
+ BPF_LOADER_UPGRADEABLE
+ );
+ console.log("Program Data PDA:", programDataPda.toBase58());
+
+ // Check current authority
+ const programDataInfo = await provider.connection.getAccountInfo(programDataPda);
+ if (!programDataInfo) {
+ console.error("ERROR: Program data account not found!");
+ console.error("Make sure the program is deployed first.");
+ process.exit(1);
+ }
+
+ // Parse current authority
+ const variant = programDataInfo.data.readUInt32LE(0);
+ if (variant !== 3) {
+ console.error("ERROR: Invalid program data variant:", variant);
+ process.exit(1);
+ }
+
+ const hasAuthority = programDataInfo.data[12] === 1;
+ let currentAuthority: PublicKey | null = null;
+
+ if (hasAuthority) {
+ const authorityBytes = programDataInfo.data.slice(13, 45);
+ currentAuthority = new PublicKey(authorityBytes);
+ console.log("Current upgrade authority:", currentAuthority.toBase58());
+ } else {
+ console.log("Current upgrade authority: None (program is immutable)");
+ }
+
+ // Check if authority matches wallet
+ if (currentAuthority?.equals(walletPubkey)) {
+ console.log("βœ… Authority already matches wallet - no action needed");
+ return;
+ }
+
+ // If authority is default (System Program), we need to fix it
+ if (currentAuthority?.equals(PublicKey.default) || !hasAuthority) {
+ console.log("⚠️ Authority is not set correctly");
+ console.log("Attempting to set authority via solana CLI...");
+
+ try {
+ // Use solana CLI to set upgrade authority
+ // Note: This only works if we have authority to do so
+ const cmd = `solana program set-upgrade-authority ${programId.toBase58()} --new-upgrade-authority ${walletPubkey.toBase58()} --skip-new-upgrade-authority-signer-check`;
+ console.log("Running:", cmd);
+
+ const output = execSync(cmd, { encoding: 'utf-8' });
+ console.log(output);
+ console.log("βœ… Authority updated successfully");
+ } catch (err) {
+ console.error("Failed to set authority via CLI:", err);
+ console.log("\nπŸ“‹ Manual fix required:");
+ console.log("The program needs to be redeployed with proper upgrade authority.");
+ console.log("Run: anchor deploy");
+ process.exit(1);
+ }
+ } else {
+ console.error("❌ Authority is set to a different key:", currentAuthority?.toBase58());
+ console.error("Cannot change authority - need to redeploy program.");
+ process.exit(1);
+ }
+}
+
+main().catch(console.error);
diff --git a/apps/contracts/stake/programs/stake_program/Cargo.toml b/apps/contracts/stake/programs/stake_program/Cargo.toml
index 9997694..4a01bdb 100644
--- a/apps/contracts/stake/programs/stake_program/Cargo.toml
+++ b/apps/contracts/stake/programs/stake_program/Cargo.toml
@@ -23,7 +23,6 @@ custom-panic = []
[dependencies]
anchor-lang = { version = "0.32.1", features = ["init-if-needed"] }
anchor-spl = "0.32.1"
-bincode = "1.3"
[lints.rust]
diff --git a/apps/contracts/stake/programs/stake_program/src/lib.rs b/apps/contracts/stake/programs/stake_program/src/lib.rs
index 76fba09..02b273b 100644
--- a/apps/contracts/stake/programs/stake_program/src/lib.rs
+++ b/apps/contracts/stake/programs/stake_program/src/lib.rs
@@ -18,7 +18,6 @@
// =============================================================================
use anchor_lang::prelude::*;
-use anchor_lang::solana_program::bpf_loader_upgradeable::UpgradeableLoaderState;
use anchor_spl::associated_token::get_associated_token_address;
use anchor_spl::token::{self, Mint, Token, TokenAccount, Transfer};
@@ -140,20 +139,36 @@ pub mod stake_program {
// H-04: Verify admin is the program upgrade authority
// This restricts pool creation to whoever can upgrade the program
+ //
+ // ProgramData account layout (UpgradeableLoaderState::ProgramData):
+ // - Bytes 0-3: variant discriminator (u32 LE, value = 3 for ProgramData)
+ // - Bytes 4-11: slot when deployed (u64 LE)
+ // - Byte 12: Option tag (0 = None, 1 = Some upgrade authority)
+ // - Bytes 13-44: upgrade authority Pubkey (if Option tag is 1)
{
let program_data = &ctx.accounts.program_data;
let data = program_data.try_borrow_data()?;
- let state: UpgradeableLoaderState = bincode::deserialize(&data)
- .map_err(|_| error!(CustomError::InvalidProgramData))?;
- match state {
- UpgradeableLoaderState::ProgramData { upgrade_authority_address, .. } => {
- require!(
- upgrade_authority_address == Some(ctx.accounts.admin.key()),
- CustomError::UnauthorizedPoolCreator
- );
- }
- _ => return Err(error!(CustomError::InvalidProgramData)),
- }
+
+ // Verify minimum length for ProgramData header
+ require!(data.len() >= 45, CustomError::InvalidProgramData);
+
+ // Check variant discriminator (must be 3 for ProgramData)
+ let variant = u32::from_le_bytes(data[0..4].try_into().unwrap());
+ require!(variant == 3, CustomError::InvalidProgramData);
+
+ // Require an upgrade authority to be set (byte 12 = 1 means Some)
+ let has_authority = data[12] == 1;
+ require!(has_authority, CustomError::UnauthorizedPoolCreator);
+
+ // Extract and verify upgrade authority pubkey
+ let authority_bytes: [u8; 32] = data[13..45].try_into().unwrap();
+ let upgrade_authority = Pubkey::new_from_array(authority_bytes);
+
+ // Strictly require admin to match the upgrade authority
+ require!(
+ upgrade_authority == ctx.accounts.admin.key(),
+ CustomError::UnauthorizedPoolCreator
+ );
}
// Validate reward percentage to prevent accidental extreme values
@@ -819,7 +834,13 @@ pub struct CreatePool<'info> {
/// H-04: Program data account to verify admin is the upgrade authority
/// This restricts pool creation to whoever can upgrade the program
- /// CHECK: Validated in instruction handler that checks upgrade_authority matches admin
+ /// The PDA is derived from the program ID and owned by the BPF Loader Upgradeable
+ /// CHECK: Seeds validated via constraint, upgrade_authority checked in instruction handler
+ #[account(
+ seeds = [crate::ID.as_ref()],
+ bump,
+ seeds::program = anchor_lang::solana_program::bpf_loader_upgradeable::id()
+ )]
pub program_data: AccountInfo<'info>,
pub system_program: Program<'info, System>,
diff --git a/apps/contracts/stake/tests/account-reuse.test.ts b/apps/contracts/stake/tests/account-reuse.test.ts
index 862d4c6..357b5c9 100644
--- a/apps/contracts/stake/tests/account-reuse.test.ts
+++ b/apps/contracts/stake/tests/account-reuse.test.ts
@@ -8,7 +8,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ”’ Stake Program - Account Reuse Prevention Tests", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -56,6 +57,7 @@ describe("πŸ”’ Stake Program - Account Reuse Prevention Tests", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/atomic-deployment.test.ts b/apps/contracts/stake/tests/atomic-deployment.test.ts
index 3a38250..5fb75e5 100644
--- a/apps/contracts/stake/tests/atomic-deployment.test.ts
+++ b/apps/contracts/stake/tests/atomic-deployment.test.ts
@@ -1,7 +1,8 @@
-import anchor from "@coral-xyz/anchor";
+import * as anchor from "@coral-xyz/anchor";
const BN = anchor.BN;
import { expect } from "chai";
import { createMint, getAssociatedTokenAddressSync } from "@solana/spl-token";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ”’ Stake Program - Atomic Deployment Security", () => {
const provider = anchor.AnchorProvider.env();
@@ -93,6 +94,7 @@ describe("πŸ”’ Stake Program - Atomic Deployment Security", () => {
rewardVault: rewardVaultPda,
poolVault: poolVaultPda,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
systemProgram: anchor.web3.SystemProgram.programId,
tokenProgram: anchor.utils.token.TOKEN_PROGRAM_ID,
rent: anchor.web3.SYSVAR_RENT_PUBKEY,
@@ -125,6 +127,7 @@ describe("πŸ”’ Stake Program - Atomic Deployment Security", () => {
rewardVault: rewardVaultPda,
poolVault: poolVaultPda,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
systemProgram: anchor.web3.SystemProgram.programId,
tokenProgram: anchor.utils.token.TOKEN_PROGRAM_ID,
rent: anchor.web3.SYSVAR_RENT_PUBKEY,
diff --git a/apps/contracts/stake/tests/edge-cases.test.ts b/apps/contracts/stake/tests/edge-cases.test.ts
index 17dd4f6..0a77f7e 100644
--- a/apps/contracts/stake/tests/edge-cases.test.ts
+++ b/apps/contracts/stake/tests/edge-cases.test.ts
@@ -11,6 +11,7 @@ import {
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
import { getTestEnvironment, warpSlots, TEST_SLOTS_PER_PERIOD } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
// Use small slot counts for fast testing - reward logic works the same
const SLOTS_PER_DAY = TEST_SLOTS_PER_PERIOD;
@@ -62,6 +63,7 @@ describe("πŸ§ͺ Stake Program - Edge Cases", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/events.test.ts b/apps/contracts/stake/tests/events.test.ts
index 867aec0..be929ed 100644
--- a/apps/contracts/stake/tests/events.test.ts
+++ b/apps/contracts/stake/tests/events.test.ts
@@ -8,7 +8,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸŽ‰ Stake Program - Events", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -80,6 +81,7 @@ describe("πŸŽ‰ Stake Program - Events", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/multi-pool.test.ts b/apps/contracts/stake/tests/multi-pool.test.ts
index 7a02482..0b0636f 100644
--- a/apps/contracts/stake/tests/multi-pool.test.ts
+++ b/apps/contracts/stake/tests/multi-pool.test.ts
@@ -9,7 +9,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment, getPoolPDA, poolIdToBytes } from "./test-utils";
+import { getTestEnvironment, getPoolPDA, poolIdToBytes, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ”’ Stake Program - Multiple Pools per Token", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -74,6 +75,7 @@ describe("πŸ”’ Stake Program - Multiple Pools per Token", () => {
tokenMint: tokenMint,
rewardMint: rewardMint1,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
@@ -114,6 +116,7 @@ describe("πŸ”’ Stake Program - Multiple Pools per Token", () => {
tokenMint: tokenMint,
rewardMint: rewardMint2,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
@@ -148,6 +151,7 @@ describe("πŸ”’ Stake Program - Multiple Pools per Token", () => {
tokenMint: tokenMint,
rewardMint: rewardMint1, // Reuse first reward mint
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/pda-validation.test.ts b/apps/contracts/stake/tests/pda-validation.test.ts
index ff68b55..241879a 100644
--- a/apps/contracts/stake/tests/pda-validation.test.ts
+++ b/apps/contracts/stake/tests/pda-validation.test.ts
@@ -8,7 +8,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ” Stake Program - PDA Seed Validation", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -62,6 +63,7 @@ describe("πŸ” Stake Program - PDA Seed Validation", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/pool-association.test.ts b/apps/contracts/stake/tests/pool-association.test.ts
index 2717865..7ff97be 100644
--- a/apps/contracts/stake/tests/pool-association.test.ts
+++ b/apps/contracts/stake/tests/pool-association.test.ts
@@ -8,7 +8,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ”’ Stake Program - Pool Association Security Tests", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -68,6 +69,7 @@ describe("πŸ”’ Stake Program - Pool Association Security Tests", () => {
tokenMint: tokenMintA,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
@@ -89,6 +91,7 @@ describe("πŸ”’ Stake Program - Pool Association Security Tests", () => {
tokenMint: tokenMintB,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/pool-configuration.test.ts b/apps/contracts/stake/tests/pool-configuration.test.ts
index 70e6e34..5d93fa7 100644
--- a/apps/contracts/stake/tests/pool-configuration.test.ts
+++ b/apps/contracts/stake/tests/pool-configuration.test.ts
@@ -3,7 +3,8 @@ import { Program } from "@coral-xyz/anchor";
import { createMint, getMint } from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ”§ Stake Program - Pool Configuration", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -52,6 +53,7 @@ describe("πŸ”§ Stake Program - Pool Configuration", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/pool-creation.test.ts b/apps/contracts/stake/tests/pool-creation.test.ts
index e8e6d29..ac66ddd 100644
--- a/apps/contracts/stake/tests/pool-creation.test.ts
+++ b/apps/contracts/stake/tests/pool-creation.test.ts
@@ -3,7 +3,7 @@ import { Program } from "@coral-xyz/anchor";
import { createMint, getMint, TOKEN_PROGRAM_ID } from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment, getPoolPDA } from "./test-utils";
+import { getTestEnvironment, getPoolPDA, getProgramDataPDA } from "./test-utils";
describe("πŸͺ™ Stake Program - Create Pool", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -57,6 +57,7 @@ describe("πŸͺ™ Stake Program - Create Pool", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
@@ -104,6 +105,7 @@ describe("πŸͺ™ Stake Program - Create Pool", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
@@ -146,6 +148,7 @@ describe("πŸͺ™ Stake Program - Create Pool", () => {
tokenMint: testTokenMint,
rewardMint: testRewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
@@ -201,6 +204,7 @@ describe("πŸͺ™ Stake Program - Create Pool", () => {
tokenMint: testTokenMint,
rewardMint: testRewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
@@ -243,6 +247,7 @@ describe("πŸͺ™ Stake Program - Create Pool", () => {
tokenMint: testTokenMint,
rewardMint: testRewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/reward-epochs.test.ts b/apps/contracts/stake/tests/reward-epochs.test.ts
index ec97d81..f22ab40 100644
--- a/apps/contracts/stake/tests/reward-epochs.test.ts
+++ b/apps/contracts/stake/tests/reward-epochs.test.ts
@@ -8,6 +8,7 @@ import {
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
import { getTestEnvironment, warpSlots, TEST_SLOTS_PER_PERIOD } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
// Use small slot counts for fast testing - reward logic works the same
const SLOTS_PER_DAY = TEST_SLOTS_PER_PERIOD;
@@ -55,6 +56,7 @@ describe("πŸ• Stake Program - Reward Epochs", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/reward-scenario.test.ts b/apps/contracts/stake/tests/reward-scenario.test.ts
index b48af15..4514ec7 100644
--- a/apps/contracts/stake/tests/reward-scenario.test.ts
+++ b/apps/contracts/stake/tests/reward-scenario.test.ts
@@ -9,6 +9,7 @@ import {
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
import { getTestEnvironment, advanceToSlot } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
// Test constants for slot-based timing
// Using the same SLOTS_PER_YEAR constant as defined in the stake program (lib.rs)
@@ -77,6 +78,7 @@ describe("πŸ’° Stake Program - Reward Scenario", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/reward-vault.test.ts b/apps/contracts/stake/tests/reward-vault.test.ts
index 3ef3c62..dfc0c9e 100644
--- a/apps/contracts/stake/tests/reward-vault.test.ts
+++ b/apps/contracts/stake/tests/reward-vault.test.ts
@@ -10,7 +10,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("🏦 Stake Program - Reward Vault Management", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -59,6 +60,7 @@ describe("🏦 Stake Program - Reward Vault Management", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/safety-features.test.ts b/apps/contracts/stake/tests/safety-features.test.ts
index 07aa5d9..fb96126 100644
--- a/apps/contracts/stake/tests/safety-features.test.ts
+++ b/apps/contracts/stake/tests/safety-features.test.ts
@@ -8,7 +8,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ”’ Stake Program - Safety Features", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -50,6 +51,7 @@ describe("πŸ”’ Stake Program - Safety Features", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/security.test.ts b/apps/contracts/stake/tests/security.test.ts
index 0fbf5ba..77d65b4 100644
--- a/apps/contracts/stake/tests/security.test.ts
+++ b/apps/contracts/stake/tests/security.test.ts
@@ -10,7 +10,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ”’ Stake Program - Security Tests", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -60,6 +61,7 @@ describe("πŸ”’ Stake Program - Security Tests", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/test-utils.ts b/apps/contracts/stake/tests/test-utils.ts
index 3d35934..36f9772 100644
--- a/apps/contracts/stake/tests/test-utils.ts
+++ b/apps/contracts/stake/tests/test-utils.ts
@@ -116,3 +116,99 @@ export function getPoolPDA(
);
}
+/**
+ * BPF Loader Upgradeable Program ID
+ */
+export const BPF_LOADER_UPGRADEABLE_PROGRAM_ID = new anchor.web3.PublicKey(
+ "BPFLoaderUpgradeab1e11111111111111111111111"
+);
+
+/**
+ * Creates the program data account bytes with the specified upgrade authority.
+ * This is used to fix the program data in test environments where bankrun
+ * doesn't properly set the upgrade authority.
+ *
+ * ProgramData layout:
+ * - Bytes 0-3: variant discriminator (u32 LE, value = 3 for ProgramData)
+ * - Bytes 4-11: slot when deployed (u64 LE)
+ * - Byte 12: Option tag (0 = None, 1 = Some upgrade authority)
+ * - Bytes 13-44: upgrade authority Pubkey (if Option tag is 1)
+ * - Bytes 45+: program data (BPF bytecode)
+ *
+ * @param upgradeAuthority The pubkey to set as upgrade authority
+ * @param existingData Optional existing program data to preserve (bytecode)
+ * @returns Buffer with properly formatted program data header
+ */
+export function createProgramDataHeader(
+ upgradeAuthority: anchor.web3.PublicKey,
+ slot: bigint = BigInt(0)
+): Buffer {
+ const header = Buffer.alloc(45);
+
+ // Variant discriminator: 3 for ProgramData (little-endian u32)
+ header.writeUInt32LE(3, 0);
+
+ // Slot (little-endian u64)
+ header.writeBigUInt64LE(slot, 4);
+
+ // Option tag: 1 for Some (has upgrade authority)
+ header[12] = 1;
+
+ // Upgrade authority pubkey (32 bytes)
+ upgradeAuthority.toBuffer().copy(header, 13);
+
+ return header;
+}
+
+/**
+ * Fixes the program data account to have the correct upgrade authority.
+ * Call this in test setup when using bankrun to ensure H-04 checks pass.
+ *
+ * @param connection The Solana connection
+ * @param programId The program ID
+ * @param upgradeAuthority The pubkey to set as upgrade authority
+ */
+export async function fixProgramDataAuthority(
+ connection: anchor.web3.Connection,
+ programId: anchor.web3.PublicKey,
+ upgradeAuthority: anchor.web3.PublicKey
+): Promise<void> {
+ const programDataPda = getProgramDataPDA(programId);
+
+ // Get existing program data account
+ const existingAccount = await connection.getAccountInfo(programDataPda);
+ if (!existingAccount) {
+ throw new Error(`Program data account not found: ${programDataPda.toBase58()}`);
+ }
+
+ // Create new header with correct upgrade authority
+ const newHeader = createProgramDataHeader(upgradeAuthority);
+
+ // Preserve the existing program bytecode (everything after the 45-byte header)
+ const newData = Buffer.concat([
+ newHeader,
+ existingAccount.data.slice(45)
+ ]);
+
+ // Note: In bankrun, you would use context.setAccount() to update this.
+ // This function documents the expected format. The actual update must be
+ // done through the test framework's account manipulation API.
+ console.log(`Program data should be updated with upgrade authority: ${upgradeAuthority.toBase58()}`);
+}
+
+/**
+ * Derives the program data PDA for a given program ID
+ * Required for H-04: Restrict create_pool to program upgrade authority
+ * @param programId The program ID
+ * @returns The program data PDA
+ */
+export function getProgramDataPDA(
+ programId: anchor.web3.PublicKey
+): anchor.web3.PublicKey {
+ const [programData] = anchor.web3.PublicKey.findProgramAddressSync(
+ [programId.toBuffer()],
+ BPF_LOADER_UPGRADEABLE_PROGRAM_ID
+ );
+ return programData;
+}
+
diff --git a/apps/contracts/stake/tests/user-staking.test.ts b/apps/contracts/stake/tests/user-staking.test.ts
index 40c9bb9..33e2c7a 100644
--- a/apps/contracts/stake/tests/user-staking.test.ts
+++ b/apps/contracts/stake/tests/user-staking.test.ts
@@ -9,7 +9,8 @@ import {
} from "@solana/spl-token";
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
-import { getTestEnvironment } from "./test-utils";
+import { getTestEnvironment, getProgramDataPDA } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
describe("πŸ§‘β€πŸ’Ό Stake Program - User Staking", () => {
const { provider, program, admin } = getTestEnvironment();
@@ -60,6 +61,7 @@ describe("πŸ§‘β€πŸ’Ό Stake Program - User Staking", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
diff --git a/apps/contracts/stake/tests/user-withdrawal.test.ts b/apps/contracts/stake/tests/user-withdrawal.test.ts
index e6b8214..b06f2f6 100644
--- a/apps/contracts/stake/tests/user-withdrawal.test.ts
+++ b/apps/contracts/stake/tests/user-withdrawal.test.ts
@@ -11,6 +11,7 @@ import {
import { expect } from "chai";
import { StakeProgram } from "../target/types/stake_program";
import { getTestEnvironment, warpSlots, TEST_SLOTS_PER_PERIOD } from "./test-utils";
+import { getProgramDataPDA } from "./test-utils";
// Use small slot counts for fast testing - reward logic works the same
const SLOTS_PER_DAY = TEST_SLOTS_PER_PERIOD;
@@ -65,6 +66,7 @@ describe("πŸ’Έ Stake Program - User Withdrawal", () => {
tokenMint: tokenMint,
rewardMint: rewardMint,
admin: admin.publicKey,
+ programData: getProgramDataPDA(program.programId),
})
.rpc();
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment