Program Upgrades

┌─────────────────────────────────────────────────────────────────┐
│                  Safe Upgrade Process                           │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  1. Develop & Test                                              │
│     ┌─────────────────────────────────────────────────────┐    │
│     │ • Write new program code                             │    │
│     │ • Comprehensive unit tests                           │    │
│     │ • Integration tests on localnet                      │    │
│     │ • Security audit (for major changes)                 │    │
│     └─────────────────────────────────────────────────────┘    │
│                           │                                     │
│                           ▼                                     │
│  2. Deploy to Devnet                                            │
│     ┌─────────────────────────────────────────────────────┐    │
│     │ • Deploy new version to devnet                       │    │
│     │ • Run migration scripts                              │    │
│     │ • Verify state consistency                           │    │
│     │ • Test all user flows                                │    │
│     └─────────────────────────────────────────────────────┘    │
│                           │                                     │
│                           ▼                                     │
│  3. Governance Proposal                                         │
│     ┌─────────────────────────────────────────────────────┐    │
│     │ • Create upgrade proposal                            │    │
│     │ • Include: changelog, audit report, buffer address   │    │
│     │ • Community discussion period                        │    │
│     │ • Voting period                                      │    │
│     └─────────────────────────────────────────────────────┘    │
│                           │                                     │
│                           ▼                                     │
│  4. Execute Upgrade                                             │
│     ┌─────────────────────────────────────────────────────┐    │
│     │ • Timelock delay (e.g., 48 hours)                    │    │
│     │ • Execute upgrade via governance                     │    │
│     │ • Run post-upgrade migrations                        │    │
│     │ • Verify deployment                                  │    │
│     └─────────────────────────────────────────────────────┘    │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

use anchor_lang::prelude::*;

declare_id!("Prog1111111111111111111111111111111111111");

// Store version in a global config account
#[account]
pub struct ProgramConfig {
    pub authority: Pubkey,
    pub version: ProgramVersion,
    pub paused: bool,
    pub upgrade_buffer: Option<Pubkey>,
    pub scheduled_upgrade_time: Option<i64>,
}

#[derive(AnchorSerialize, AnchorDeserialize, Clone, PartialEq, Eq)]
pub struct ProgramVersion {
    pub major: u8,
    pub minor: u8,
    pub patch: u8,
}

impl ProgramVersion {
    pub fn new(major: u8, minor: u8, patch: u8) -> Self {
        Self { major, minor, patch }
    }

    pub fn is_compatible(&self, other: &ProgramVersion) -> bool {
        // Same major version is compatible
        self.major == other.major
    }
}

// Account versioning for state migrations
#[account]
pub struct UserAccountV1 {
    pub owner: Pubkey,
    pub balance: u64,
    pub created_at: i64,
}

#[account]
pub struct UserAccountV2 {
    pub owner: Pubkey,
    pub balance: u64,
    pub created_at: i64,
    // New fields
    pub last_activity: i64,
    pub total_transactions: u64,
    pub version: u8,
}

impl UserAccountV2 {
    pub const VERSION: u8 = 2;

    pub fn from_v1(v1: &UserAccountV1, current_time: i64) -> Self {
        Self {
            owner: v1.owner,
            balance: v1.balance,
            created_at: v1.created_at,
            last_activity: current_time,
            total_transactions: 0,
            version: Self::VERSION,
        }
    }
}

#[program]
pub mod versioned_program {
    use super::*;

    // Migration instruction
    pub fn migrate_user_account(ctx: Context<MigrateUserAccount>) -> Result<()> {
        let old_account = &ctx.accounts.user_account_v1;
        let new_account = &mut ctx.accounts.user_account_v2;

        let clock = Clock::get()?;

        // Copy and upgrade data
        new_account.owner = old_account.owner;
        new_account.balance = old_account.balance;
        new_account.created_at = old_account.created_at;
        new_account.last_activity = clock.unix_timestamp;
        new_account.total_transactions = 0;
        new_account.version = UserAccountV2::VERSION;

        // Close old account (return rent to user)
        // Done via close constraint

        Ok(())
    }

    // Version-aware instruction
    pub fn transfer(ctx: Context<Transfer>, amount: u64) -> Result<()> {
        let config = &ctx.accounts.program_config;

        // Check program is not paused
        require!(!config.paused, ErrorCode::ProgramPaused);

        // Version-specific logic
        let sender = &mut ctx.accounts.sender;
        let receiver = &mut ctx.accounts.receiver;

        require!(sender.balance >= amount, ErrorCode::InsufficientBalance);

        sender.balance -= amount;
        sender.last_activity = Clock::get()?.unix_timestamp;
        sender.total_transactions += 1;

        receiver.balance += amount;
        receiver.last_activity = Clock::get()?.unix_timestamp;
        receiver.total_transactions += 1;

        Ok(())
    }
}

#[derive(Accounts)]
pub struct MigrateUserAccount<'info> {
    #[account(mut)]
    pub authority: Signer<'info>,

    #[account(
        mut,
        close = authority,
        constraint = user_account_v1.owner == authority.key()
    )]
    pub user_account_v1: Account<'info, UserAccountV1>,

    #[account(
        init,
        payer = authority,
        space = 8 + std::mem::size_of::<UserAccountV2>(),
        seeds = [b"user_v2", authority.key().as_ref()],
        bump
    )]
    pub user_account_v2: Account<'info, UserAccountV2>,

    pub system_program: Program<'info, System>,
}

#[error_code]
pub enum ErrorCode {
    #[msg("Program is paused")]
    ProgramPaused,
    #[msg("Insufficient balance")]
    InsufficientBalance,
    #[msg("Version mismatch")]
    VersionMismatch,
}

import {
  Connection,
  Keypair,
  PublicKey,
  Transaction,
  TransactionInstruction,
} from "@solana/web3.js";
import * as fs from "fs";

const BPF_LOADER_UPGRADEABLE = new PublicKey(
  "BPFLoaderUpgradeab1e11111111111111111111111",
);

interface BufferInfo {
  address: PublicKey;
  authority: PublicKey;
  dataLength: number;
}

// Create upgrade buffer
export async function createUpgradeBuffer(
  connection: Connection,
  payer: Keypair,
  programPath: string,
  authority: PublicKey,
): Promise<BufferInfo> {
  const programData = fs.readFileSync(programPath);
  const bufferKeypair = Keypair.generate();

  const bufferSize = programData.length + 37; // Buffer header
  const rent = await connection.getMinimumBalanceForRentExemption(bufferSize);

  // Create buffer account
  const createAccountIx = SystemProgram.createAccount({
    fromPubkey: payer.publicKey,
    newAccountPubkey: bufferKeypair.publicKey,
    lamports: rent,
    space: bufferSize,
    programId: BPF_LOADER_UPGRADEABLE,
  });

  // Initialize buffer with authority
  const initBufferIx: TransactionInstruction = {
    programId: BPF_LOADER_UPGRADEABLE,
    keys: [
      { pubkey: bufferKeypair.publicKey, isSigner: false, isWritable: true },
      { pubkey: authority, isSigner: false, isWritable: false },
    ],
    data: Buffer.from([0]), // InitializeBuffer
  };

  const tx = new Transaction().add(createAccountIx, initBufferIx);
  await connection.sendTransaction(tx, [payer, bufferKeypair]);

  // Write program data in chunks
  const chunkSize = 900; // Safe chunk size
  for (let offset = 0; offset < programData.length; offset += chunkSize) {
    const chunk = programData.slice(offset, offset + chunkSize);

    const writeIx: TransactionInstruction = {
      programId: BPF_LOADER_UPGRADEABLE,
      keys: [
        { pubkey: bufferKeypair.publicKey, isSigner: false, isWritable: true },
        { pubkey: authority, isSigner: true, isWritable: false },
      ],
      data: Buffer.concat([
        Buffer.from([1]), // Write instruction
        Buffer.from(new Uint32Array([offset]).buffer),
        chunk,
      ]),
    };

    const writeTx = new Transaction().add(writeIx);
    await connection.sendTransaction(writeTx, [payer]);

    console.log(`Written ${offset + chunk.length}/${programData.length} bytes`);
  }

  return {
    address: bufferKeypair.publicKey,
    authority,
    dataLength: programData.length,
  };
}

// Verify buffer matches expected program
export async function verifyBuffer(
  connection: Connection,
  bufferAddress: PublicKey,
  expectedHash: string,
): Promise<boolean> {
  const accountInfo = await connection.getAccountInfo(bufferAddress);
  if (!accountInfo) {
    throw new Error("Buffer account not found");
  }

  // Skip header (37 bytes) to get program data
  const programData = accountInfo.data.slice(37);

  const crypto = await import("crypto");
  const hash = crypto.createHash("sha256").update(programData).digest("hex");

  return hash === expectedHash;
}

// Execute upgrade via governance
export async function createUpgradeProposal(
  connection: Connection,
  realmAddress: PublicKey,
  governanceAddress: PublicKey,
  programId: PublicKey,
  bufferAddress: PublicKey,
  proposer: Keypair,
): Promise<PublicKey> {
  const instructions: TransactionInstruction[] = [];

  // Get program data address
  const [programDataAddress] = PublicKey.findProgramAddressSync(
    [programId.toBuffer()],
    BPF_LOADER_UPGRADEABLE,
  );

  // The upgrade instruction that governance will execute
  const upgradeInstruction: TransactionInstruction = {
    programId: BPF_LOADER_UPGRADEABLE,
    keys: [
      { pubkey: programDataAddress, isSigner: false, isWritable: true },
      { pubkey: programId, isSigner: false, isWritable: true },
      { pubkey: bufferAddress, isSigner: false, isWritable: true },
      { pubkey: proposer.publicKey, isSigner: false, isWritable: true }, // Spill
      { pubkey: SystemProgram.programId, isSigner: false, isWritable: false },
      { pubkey: governanceAddress, isSigner: true, isWritable: false }, // Authority
    ],
    data: Buffer.from([3]), // Upgrade
  };

  // Create proposal with upgrade instruction
  // (Using SPL Governance)
  const proposalAddress = await withCreateProposal(
    instructions,
    GOVERNANCE_PROGRAM_ID,
    3, // Version
    realmAddress,
    governanceAddress,
    tokenOwnerRecord,
    "Upgrade Program to v2.0.0",
    "This upgrade includes: ...",
    communityMint,
    proposer.publicKey,
    0, // Proposal index
    VoteType.SingleChoice,
    ["Approve"],
    true, // Use deny option
    proposer.publicKey,
  );

  // Add the upgrade instruction to the proposal
  await withInsertTransaction(
    instructions,
    GOVERNANCE_PROGRAM_ID,
    3,
    governanceAddress,
    proposalAddress,
    tokenOwnerRecord,
    proposer.publicKey,
    0, // Option index
    0, // Instruction index
    0, // Hold up time
    [upgradeInstruction],
    proposer.publicKey,
  );

  const tx = new Transaction().add(...instructions);
  await connection.sendTransaction(tx, [proposer]);

  return proposalAddress;
}

use anchor_lang::prelude::*;

// Strategy 1: Lazy Migration
// Migrate accounts when they're used
#[account]
pub struct AccountWithVersion {
    pub version: u8,
    pub data: [u8; 100],
}

pub fn process_with_migration(ctx: Context<ProcessAccount>) -> Result<()> {
    let account = &mut ctx.accounts.account;

    // Check version and migrate if needed
    if account.version < CURRENT_VERSION {
        migrate_account(account)?;
    }

    // Continue with normal processing
    Ok(())
}

fn migrate_account(account: &mut AccountWithVersion) -> Result<()> {
    match account.version {
        0 => {
            // V0 -> V1 migration
            // ... transform data ...
            account.version = 1;
        }
        1 => {
            // V1 -> V2 migration
            // ... transform data ...
            account.version = 2;
        }
        _ => {}
    }
    Ok(())
}

// Strategy 2: Batch Migration
// Migrate multiple accounts in one transaction
#[derive(Accounts)]
pub struct BatchMigrate<'info> {
    #[account(mut)]
    pub authority: Signer<'info>,

    #[account(
        seeds = [b"config"],
        bump,
        constraint = config.authority == authority.key()
    )]
    pub config: Account<'info, ProgramConfig>,

    // Remaining accounts are accounts to migrate
}

pub fn batch_migrate(ctx: Context<BatchMigrate>) -> Result<()> {
    let remaining = &mut ctx.remaining_accounts.iter();

    for account_info in remaining {
        // Deserialize, migrate, serialize
        let mut data = account_info.try_borrow_mut_data()?;

        // Check discriminator and version
        let version = data[8]; // After 8-byte discriminator

        if version < CURRENT_VERSION {
            // Perform in-place migration
            migrate_data(&mut data)?;
        }
    }

    Ok(())
}

// Strategy 3: Dual-Account Period
// Support both old and new accounts during transition
pub mod dual_account {
    use super::*;

    #[account]
    pub struct LegacyAccount {
        pub owner: Pubkey,
        pub balance: u64,
    }

    #[account]
    pub struct NewAccount {
        pub owner: Pubkey,
        pub balance: u64,
        pub metadata: [u8; 32],
    }

    // Accept either account type
    pub fn withdraw(ctx: Context<WithdrawDual>, amount: u64) -> Result<()> {
        // Try new account first
        if let Some(new_account) = &mut ctx.accounts.new_account {
            require!(new_account.balance >= amount, ErrorCode::Insufficient);
            new_account.balance -= amount;
            return Ok(());
        }

        // Fall back to legacy account
        if let Some(legacy_account) = &mut ctx.accounts.legacy_account {
            require!(legacy_account.balance >= amount, ErrorCode::Insufficient);
            legacy_account.balance -= amount;
            return Ok(());
        }

        Err(ErrorCode::NoAccountProvided.into())
    }

    #[derive(Accounts)]
    pub struct WithdrawDual<'info> {
        #[account(mut)]
        pub authority: Signer<'info>,

        #[account(
            mut,
            constraint = legacy_account.as_ref().map_or(true, |a| a.owner == authority.key())
        )]
        pub legacy_account: Option<Account<'info, LegacyAccount>>,

        #[account(
            mut,
            constraint = new_account.as_ref().map_or(true, |a| a.owner == authority.key())
        )]
        pub new_account: Option<Account<'info, NewAccount>>,
    }
}

// Store previous program versions for rollback
interface ProgramVersionRecord {
  version: string;
  bufferAddress: string;
  deployedAt: Date;
  deployer: string;
  hash: string;
}

class ProgramVersionManager {
  private versions: Map<string, ProgramVersionRecord[]> = new Map();

  async recordDeployment(
    programId: string,
    bufferAddress: string,
    version: string,
    deployer: string,
  ): Promise<void> {
    const hash = await this.computeProgramHash(bufferAddress);

    const record: ProgramVersionRecord = {
      version,
      bufferAddress,
      deployedAt: new Date(),
      deployer,
      hash,
    };

    const existing = this.versions.get(programId) || [];
    existing.push(record);
    this.versions.set(programId, existing);

    // Store in database
    await db.query(
      `INSERT INTO program_versions (program_id, version, buffer_address, deployed_at, deployer, hash)
       VALUES ($1, $2, $3, $4, $5, $6)`,
      [programId, version, bufferAddress, record.deployedAt, deployer, hash],
    );
  }

  async rollback(programId: string, targetVersion: string): Promise<string> {
    const versions = this.versions.get(programId);
    if (!versions) throw new Error("No version history");

    const targetRecord = versions.find((v) => v.version === targetVersion);
    if (!targetRecord) throw new Error("Target version not found");

    // Verify buffer still exists and matches
    const bufferValid = await this.verifyBuffer(
      targetRecord.bufferAddress,
      targetRecord.hash,
    );
    if (!bufferValid) {
      // Need to recreate buffer from stored bytecode
      const bytecode = await this.fetchStoredBytecode(programId, targetVersion);
      targetRecord.bufferAddress = await this.createNewBuffer(bytecode);
    }

    // Create rollback proposal
    const proposalId = await this.createRollbackProposal(
      programId,
      targetRecord,
    );

    return proposalId;
  }
}