From fffba72723b19cba80f2c6f29f624aede5b349f0 Mon Sep 17 00:00:00 2001
From: Hardhat Chad <hardhatchad@gmail.com>
Date: Mon, 3 Jun 2024 16:46:17 +0000
Subject: [PATCH] cap max supply on upgrade

---
 src/processor/upgrade.rs | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/processor/upgrade.rs b/src/processor/upgrade.rs
index f1f2bf5..53930d1 100644
--- a/src/processor/upgrade.rs
+++ b/src/processor/upgrade.rs
@@ -1,10 +1,12 @@
 use solana_program::{
     account_info::AccountInfo, entrypoint::ProgramResult, program_error::ProgramError,
-    pubkey::Pubkey,
+    program_pack::Pack, pubkey::Pubkey,
 };
+use spl_token::state::Mint;
 
 use crate::{
-    instruction::StakeArgs, loaders::*, MINT_ADDRESS, MINT_V1_ADDRESS, TREASURY, TREASURY_BUMP,
+    error::OreError, instruction::StakeArgs, loaders::*, MAX_SUPPLY, MINT_ADDRESS, MINT_V1_ADDRESS,
+    TREASURY, TREASURY_BUMP,
 };
 
 /// Upgrade allows a user to migrate a v1 token to a v2 token one-for-one. Its responsibilies include:
@@ -58,6 +60,13 @@ pub fn process_upgrade<'a, 'info>(
     // v1 token has 9 decimals. v2 token has 11.
     let amount_to_mint = amount.saturating_mul(100);
 
+    // Cap at max supply.
+    let mint_data = mint_info.data.borrow();
+    let mint = Mint::unpack(&mint_data)?;
+    if mint.supply.saturating_add(amount_to_mint).gt(&MAX_SUPPLY) {
+        return Err(OreError::MaxSupply.into());
+    }
+
     // Mint to the beneficiary account
     solana_program::program::invoke_signed(
         &spl_token::instruction::mint_to(