feat: pda upgrade poc

programs/gateway/src/contexts.rs

@@ -1,3 +1,4 @@
+use crate::errors::Errors;
 use crate::state::*;
 use anchor_lang::prelude::*;
 use anchor_spl::associated_token::AssociatedToken;
@@ -289,6 +290,26 @@ pub struct Whitelist<'info> {
     pub system_program: Program<'info, System>,
 }
 
+/// Instruction context for extending the PDA with new fields using realloc.
+#[derive(Accounts)]
+pub struct ExtendPda<'info> {
+    /// The account of the authority performing the operation.
+    #[account(mut)]
+    pub authority: Signer<'info>,
+
+    /// Gateway PDA to be extended.
+    #[account(
+        mut,
+        seeds = [b"meta"],
+        bump,
+        constraint = authority.key() == pda.authority @ Errors::SignerIsNotAuthority
+    )]
+    pub pda: Account<'info, Pda>,
+
+    /// The system program for realloc operations.
+    pub system_program: Program<'info, System>,
+}
+
 /// Instruction context for unwhitelisting SPL tokens.
 #[derive(Accounts)]
 pub struct Unwhitelist<'info> {

programs/gateway/src/instructions/admin.rs

@@ -1,6 +1,7 @@
 use crate::{
     contexts::{
-        Initialize, ResetNonce, Unwhitelist, UpdateAuthority, UpdatePaused, UpdateTss, Whitelist,
+        ExtendPda, Initialize, ResetNonce, Unwhitelist, UpdateAuthority, UpdatePaused, UpdateTss,
+        Whitelist,
     },
     state::InstructionId,
     utils::{
@@ -185,3 +186,55 @@ pub fn reset_nonce(ctx: Context<ResetNonce>, new_nonce: u64) -> Result<()> {
 
     Ok(())
 }
+
+// Extends the PDA with new fields using realloc. Caller is authority stored in PDA.
+pub fn extend_pda(ctx: Context<ExtendPda>, bump: u8) -> Result<()> {
+    let pda_account = &mut ctx.accounts.pda;
+    let authority = &ctx.accounts.authority;
+    let system_program = &ctx.accounts.system_program;
+
+    // Calculate new size - adding 2 bytes for bump (u8) and version (u8)
+    let current_size = pda_account.to_account_info().data_len();
+    let additional_space = 2; // 2 bytes for bump and version fields
+    let new_size = current_size + additional_space;
+
+    // Ensure rent exemption
+    let rent = Rent::get()?;
+    let new_minimum_balance = rent.minimum_balance(new_size);
+    let current_balance = pda_account.to_account_info().lamports();
+    let lamports_diff = new_minimum_balance.saturating_sub(current_balance);
+
+    if lamports_diff > 0 {
+        // Transfer additional lamports to maintain rent exemption
+        let transfer_instruction = anchor_lang::solana_program::system_instruction::transfer(
+            &authority.key(),
+            &pda_account.to_account_info().key(),
+            lamports_diff,
+        );
+
+        anchor_lang::solana_program::program::invoke(
+            &transfer_instruction,
+            &[
+                authority.to_account_info(),
+                pda_account.to_account_info(),
+                system_program.to_account_info(),
+            ],
+        )?;
+    }
+
+    // Reallocate the account to the new size
+    pda_account.to_account_info().realloc(new_size, false)?;
+
+    // Note: In a real implementation, you would need to handle the migration
+    // from the old Pda struct to the new ExtendedPda struct carefully.
+    // This is a simplified example showing the realloc pattern.
+
+    msg!(
+        "PDA extended: new size = {}, bump = {}, authority = {}",
+        new_size,
+        bump,
+        authority.key()
+    );
+
+    Ok(())
+}

programs/gateway/src/instructions/mod.rs

@@ -1,4 +1,5 @@
 pub mod admin;
 pub mod deposit;
 pub mod execute;
+pub mod migrate_pda;
 pub mod withdraw;

programs/gateway/src/lib.rs

@@ -269,6 +269,14 @@ pub mod gateway {
         instructions::admin::unwhitelist_spl_mint(ctx, signature, recovery_id, message_hash, nonce)
     }
 
+    /// Extends the PDA with new fields using realloc. Caller is authority stored in PDA.
+    /// # Arguments
+    /// * `ctx` - The instruction context.
+    /// * `bump` - The bump seed to store in the extended PDA.
+    pub fn extend_pda(ctx: Context<ExtendPda>, bump: u8) -> Result<()> {
+        instructions::admin::extend_pda(ctx, bump)
+    }
+
     /// Deposits SOL into the program and credits the `receiver` on ZetaChain zEVM.
     /// # Arguments
     /// * `ctx` - The instruction context.

programs/gateway/src/state.rs

@@ -15,6 +15,26 @@ pub struct Pda {
     pub deposit_paused: bool,
 }
 
+/// Extended PDA account with additional fields for future-proofing.
+/// This version includes a bump field and can be extended with more fields.
+#[account]
+pub struct ExtendedPda {
+    /// The nonce to ensure each signature can only be used once.
+    pub nonce: u64,
+    /// The Ethereum TSS address (20 bytes).
+    pub tss_address: [u8; 20],
+    /// The authority controlling the PDA.
+    pub authority: Pubkey,
+    /// The chain ID associated with the PDA.
+    pub chain_id: u64,
+    /// Flag to indicate whether deposits are paused.
+    pub deposit_paused: bool,
+    /// The bump seed used to derive the PDA.
+    pub bump: u8,
+    /// Version field to track PDA structure version.
+    pub version: u8,
+}
+
 /// Whitelist entry account for whitelisted SPL tokens.
 #[account]
 pub struct WhitelistEntry {}

tests/pda_extension.ts

@@ -0,0 +1,202 @@
+import * as anchor from "@coral-xyz/anchor";
+import { Program } from "@coral-xyz/anchor";
+import { Gateway } from "../target/types/gateway";
+import { expect } from "chai";
+import { PublicKey } from "@solana/web3.js";
+
+describe("PDA Extension Tests", () => {
+    const provider = anchor.AnchorProvider.env();
+    anchor.setProvider(provider);
+
+    const program = anchor.workspace.Gateway as Program<Gateway>;
+    const wallet = provider.wallet;
+
+    const chainId = 111111;
+    const tssAddress = new Array(20).fill(0).map(() => Math.floor(Math.random() * 256));
+    const bump = 255;
+
+    let pda: PublicKey;
+    let pdaBump: number;
+
+    before(async () => {
+        [pda, pdaBump] = PublicKey.findProgramAddressSync(
+            [Buffer.from("meta")],
+            program.programId
+        );
+    });
+
+    it("Initializes the gateway PDA", async () => {
+        try {
+            const tx = await program.methods
+                .initialize(tssAddress, new anchor.BN(chainId))
+                .accounts({
+                    signer: wallet.publicKey,
+                    pda: pda,
+                    systemProgram: anchor.web3.SystemProgram.programId,
+                })
+                .rpc();
+
+            console.log("Initialize transaction signature:", tx);
+
+            const pdaAccount = await program.account.pda.fetch(pda);
+            expect(pdaAccount.nonce.toString()).to.equal("0");
+            expect(pdaAccount.chainId.toString()).to.equal(chainId.toString());
+            expect(pdaAccount.authority.toString()).to.equal(wallet.publicKey.toString());
+            expect(pdaAccount.depositPaused).to.be.false;
+            expect(Array.from(pdaAccount.tssAddress)).to.deep.equal(tssAddress);
+
+            console.log("PDA initialized successfully");
+        } catch (error) {
+            console.error("Error initializing PDA:", error);
+            throw error;
+        }
+    });
+
+    it("Extends the PDA with new fields using realloc", async () => {
+        try {
+            const pdaAccountInfo = await program.provider.connection.getAccountInfo(pda);
+            const currentSize = pdaAccountInfo?.data.length || 0;
+            console.log("Current PDA size:", currentSize);
+
+            const tx = await program.methods
+                .extendPda(bump)
+                .accounts({
+                    authority: wallet.publicKey,
+                    pda: pda,
+                    systemProgram: anchor.web3.SystemProgram.programId,
+                })
+                .rpc();
+
+            console.log("Extend PDA transaction signature:", tx);
+            const pdaAccountInfoAfter = await program.provider.connection.getAccountInfo(pda);
+            const newSize = pdaAccountInfoAfter?.data.length || 0;
+
+            console.log("New PDA size:", newSize);
+            expect(newSize).to.be.greaterThan(currentSize);
+            expect(newSize - currentSize).to.equal(2);
+
+            const pdaAccount = await program.account.pda.fetch(pda);
+            expect(pdaAccount.nonce.toString()).to.equal("0");
+            expect(pdaAccount.chainId.toString()).to.equal(chainId.toString());
+            expect(pdaAccount.authority.toString()).to.equal(wallet.publicKey.toString());
+            expect(pdaAccount.depositPaused).to.be.false;
+            expect(Array.from(pdaAccount.tssAddress)).to.deep.equal(tssAddress);
+
+            console.log("PDA extended successfully");
+        } catch (error) {
+            console.error("Error extending PDA:", error);
+            throw error;
+        }
+    });
+
+    it("Verifies PDA functionality after extension", async () => {
+        try {
+            const newTssAddress = new Array(20).fill(1).map(() => Math.floor(Math.random() * 256));
+
+            const tx = await program.methods
+                .updateTss(newTssAddress)
+                .accounts({
+                    signer: wallet.publicKey,
+                    pda: pda,
+                })
+                .rpc();
+
+            console.log("Update TSS transaction signature:", tx);
+
+            const pdaAccount = await program.account.pda.fetch(pda);
+            expect(Array.from(pdaAccount.tssAddress)).to.deep.equal(newTssAddress);
+            expect(pdaAccount.nonce.toString()).to.equal("0");
+
+            console.log("PDA functionality verified after extension");
+        } catch (error) {
+            console.error("Error verifying PDA functionality:", error);
+            throw error;
+        }
+    });
+
+    it("Tests rent exemption after extension", async () => {
+        try {
+            const pdaAccountInfo = await program.provider.connection.getAccountInfo(pda);
+            const lamports = pdaAccountInfo?.lamports || 0;
+            const dataLength = pdaAccountInfo?.data.length || 0;
+            const rent = await program.provider.connection.getMinimumBalanceForRentExemption(dataLength);
+
+            console.log("PDA lamports:", lamports);
+            console.log("Minimum rent-exempt balance:", rent);
+            console.log("PDA data length:", dataLength);
+
+            expect(lamports).to.be.greaterThanOrEqual(rent);
+
+            console.log("PDA is rent-exempt after extension");
+        } catch (error) {
+            console.error("Error checking rent exemption:", error);
+            throw error;
+        }
+    });
+
+    it("Tests multiple extensions", async () => {
+        try {
+            const initialSize = (await program.provider.connection.getAccountInfo(pda))?.data.length || 0;
+
+            await program.methods
+                .extendPda(200)
+                .accounts({
+                    authority: wallet.publicKey,
+                    pda: pda,
+                    systemProgram: anchor.web3.SystemProgram.programId,
+                })
+                .rpc();
+
+            const sizeAfterFirst = (await program.provider.connection.getAccountInfo(pda))?.data.length || 0;
+            expect(sizeAfterFirst).to.be.greaterThan(initialSize);
+
+            await program.methods
+                .extendPda(150)
+                .accounts({
+                    authority: wallet.publicKey,
+                    pda: pda,
+                    systemProgram: anchor.web3.SystemProgram.programId,
+                })
+                .rpc();
+
+            const sizeAfterSecond = (await program.provider.connection.getAccountInfo(pda))?.data.length || 0;
+            expect(sizeAfterSecond).to.be.greaterThan(sizeAfterFirst);
+
+            console.log("Multiple extensions successful");
+        } catch (error) {
+            console.error("Error with multiple extensions:", error);
+            throw error;
+        }
+    });
+
+    it("Tests unauthorized access prevention", async () => {
+        try {
+            const unauthorizedWallet = anchor.web3.Keypair.generate();
+            const signature = await program.provider.connection.requestAirdrop(
+                unauthorizedWallet.publicKey,
+                2 * anchor.web3.LAMPORTS_PER_SOL
+            );
+            await program.provider.connection.confirmTransaction(signature);
+
+            try {
+                await program.methods
+                    .extendPda(100)
+                    .accounts({
+                        authority: unauthorizedWallet.publicKey,
+                        pda: pda,
+                        systemProgram: anchor.web3.SystemProgram.programId,
+                    })
+                    .signers([unauthorizedWallet])
+                    .rpc();
+
+                expect.fail("Unauthorized access should have been prevented");
+            } catch (error) {
+                expect(error.message).to.include("SignerIsNotAuthority");
+                console.log("Unauthorized access correctly prevented");
+            }
+        } catch (error) {
+            console.error("Error testing unauthorized access:", error);
+            throw error;
+        }
+    });
+});