Merge pull request #89 from regolith-labs/upgrade-ix

Upgrade ORE tokens from v1 to v2

src/args.rs

@@ -103,3 +103,13 @@ pub struct StakeArgs {
 pub struct UpdateAdminArgs {
     pub new_admin: String,
 }
+
+#[derive(Parser, Debug)]
+pub struct UpgradeArgs {
+    #[arg(
+        long,
+        value_name = "AMOUNT",
+        help = "The amount of Ore to upgrade from v1 to v2. Defaults to max."
+    )]
+    pub amount: Option<f64>,
+}

src/cu_limits.rs

@@ -1,3 +1,4 @@
+pub const CU_LIMIT_UPGRADE: u32 = 20_000;
 pub const CU_LIMIT_CLAIM: u32 = 32_000;
 pub const _CU_LIMIT_RESET: u32 = 12_200;
 pub const _CU_LIMIT_MINE: u32 = 3200;

src/main.rs

@@ -13,6 +13,7 @@ mod open;
 mod rewards;
 mod send_and_confirm;
 mod stake;
+mod upgrade;
 mod utils;
 
 use std::sync::Arc;
@@ -60,6 +61,9 @@ enum Commands {
     #[command(about = "Stake to earn a rewards multiplier")]
     Stake(StakeArgs),
 
+    #[command(about = "Upgrade your ORE tokens from v1 to v2")]
+    Upgrade(UpgradeArgs),
+
     #[cfg(feature = "admin")]
     #[command(about = "Initialize the program")]
     Initialize(InitializeArgs),
@@ -162,6 +166,9 @@ async fn main() {
         Commands::Stake(args) => {
             miner.stake(args).await;
         }
+        Commands::Upgrade(args) => {
+            miner.upgrade(args).await;
+        }
         #[cfg(feature = "admin")]
         Commands::Initialize(_) => {
             miner.initialize().await;

src/upgrade.rs

@@ -0,0 +1,116 @@
+use colored::*;
+use solana_sdk::{pubkey::Pubkey, signer::Signer};
+use spl_token::amount_to_ui_amount;
+
+use crate::{
+    cu_limits::CU_LIMIT_UPGRADE,
+    send_and_confirm::ComputeBudget,
+    utils::{amount_f64_to_u64_v1, ask_confirm},
+    Miner, UpgradeArgs,
+};
+
+impl Miner {
+    pub async fn upgrade(&self, args: UpgradeArgs) {
+        let signer = &self.signer();
+        let beneficiary = self.get_or_initialize_ata().await;
+        let (sender, sender_balance) = self.get_ata_v1().await;
+
+        let amount_f64 = match args.amount {
+            Some(f64) => f64,
+            None => {
+                println!(
+                    "Defaulting to max amount of v1 Ore token in wallet: {}",
+                    sender_balance
+                );
+                sender_balance
+            }
+        };
+        let amount = amount_f64_to_u64_v1(amount_f64);
+        let amount_ui = amount_to_ui_amount(amount, ore_api::consts::TOKEN_DECIMALS_V1);
+
+        if !ask_confirm(
+            format!(
+                "\n You are about to upgrade {}. \n\nAre you sure you want to continue? [Y/n]",
+                format!("{} ORE", amount_ui).bold(),
+            )
+            .as_str(),
+        ) {
+            return;
+        }
+
+        let ix = ore_api::instruction::upgrade(signer.pubkey(), beneficiary, sender, amount);
+        match self
+            .send_and_confirm(&[ix], ComputeBudget::Fixed(CU_LIMIT_UPGRADE), false)
+            .await
+        {
+            Ok(_sig) => {}
+            Err(err) => {
+                println!("error: {}", err);
+            }
+        }
+    }
+
+    // asserts that token account exists and gets balance
+    async fn get_ata_v1(&self) -> (Pubkey, f64) {
+        // Initialize client.
+        let signer = self.signer();
+        let client = self.rpc_client.clone();
+
+        // Derive assoicated token address (for v1 account)
+        let token_account_pubkey_v1 = spl_associated_token_account::get_associated_token_address(
+            &signer.pubkey(),
+            &ore_api::consts::MINT_V1_ADDRESS,
+        );
+
+        // Get token account balance
+        let balance = match client.get_token_account(&token_account_pubkey_v1).await {
+            Ok(None) => {
+                panic!("v1 token account doesn't exist")
+            }
+            Ok(Some(token_account)) => match token_account.token_amount.ui_amount {
+                Some(ui_amount) => ui_amount,
+                None => {
+                    panic!(
+                        "Error parsing token account UI amount: {}",
+                        token_account.token_amount.amount
+                    )
+                }
+            },
+            Err(err) => {
+                panic!("Error fetching token account: {}", err)
+            }
+        };
+
+        // Return v1 token account address
+        (token_account_pubkey_v1, balance)
+    }
+
+    async fn get_or_initialize_ata(&self) -> Pubkey {
+        // Initialize client
+        let signer = self.signer();
+        let client = self.rpc_client.clone();
+
+        // Derive assoicated token address (ata)
+        let token_account_pubkey = spl_associated_token_account::get_associated_token_address(
+            &signer.pubkey(),
+            &ore_api::consts::MINT_ADDRESS,
+        );
+
+        // Check if ata already exists or init
+        if let Err(_err) = client.get_token_account(&token_account_pubkey).await {
+            println!("Initializing v2 token account...");
+            let ix = spl_associated_token_account::instruction::create_associated_token_account(
+                &signer.pubkey(),
+                &signer.pubkey(),
+                &ore_api::consts::MINT_ADDRESS,
+                &spl_token::id(),
+            );
+            self.send_and_confirm(&[ix], ComputeBudget::Dynamic, false)
+                .await
+                .ok();
+        }
+
+        // Return token account address
+        token_account_pubkey
+    }
+}

src/utils.rs

@@ -2,7 +2,9 @@ use std::io::Read;
 
 use cached::proc_macro::cached;
 use ore_api::{
-    consts::{CONFIG_ADDRESS, MINT_ADDRESS, PROOF, TOKEN_DECIMALS, TREASURY_ADDRESS},
+    consts::{
+        CONFIG_ADDRESS, MINT_ADDRESS, PROOF, TOKEN_DECIMALS, TOKEN_DECIMALS_V1, TREASURY_ADDRESS,
+    },
     state::{Config, Proof, Treasury},
 };
 use ore_utils::AccountDeserialize;
@@ -60,6 +62,10 @@ pub fn amount_f64_to_u64(amount: f64) -> u64 {
     (amount * 10f64.powf(TOKEN_DECIMALS as f64)) as u64
 }
 
+pub fn amount_f64_to_u64_v1(amount: f64) -> u64 {
+    (amount * 10f64.powf(TOKEN_DECIMALS_V1 as f64)) as u64
+}
+
 pub fn ask_confirm(question: &str) -> bool {
     println!("{}", question);
     loop {