token-upgrade: Define new conversion program (#3436)

* token-upgrade: Define interface and state for token upgrade program

* Update with feedback

* Address more feedback, make the whole thing stateless

* Address interface feedback

* Add implementation

* Update language

* Update language

* Add tests

* Add CI

* Change program id to tkup

* Rebase fix

* Address feedback

* Enforce decimals match

.github/workflows/pull-request-token-upgrade.yml

@@ -0,0 +1,62 @@
+name: Token Upgrade Pull Request
+
+on:
+  pull_request:
+    paths:
+    - 'token-upgrade/**'
+    - 'token/**'
+    - 'ci/*-version.sh'
+    - '.github/workflows/pull-request-token-upgrade.yml'
+  push:
+    branches: [master]
+    paths:
+    - 'token-upgrade/**'
+    - 'token/**'
+    - 'ci/*-version.sh'
+    - '.github/workflows/pull-request-token-upgrade.yml'
+
+jobs:
+  cargo-test-sbf:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Set env vars
+        run: |
+          source ci/rust-version.sh
+          echo "RUST_STABLE=$rust_stable" >> $GITHUB_ENV
+          source ci/solana-version.sh
+          echo "SOLANA_VERSION=$solana_version" >> $GITHUB_ENV
+
+      - uses: actions-rs/toolchain@v1
+        with:
+          toolchain: ${{ env.RUST_STABLE }}
+          override: true
+          profile: minimal
+
+      - uses: actions/cache@v2
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+          key: cargo-build-${{ hashFiles('**/Cargo.lock') }}-${{ env.RUST_STABLE}}
+
+      - uses: actions/cache@v2
+        with:
+          path: |
+            ~/.cargo/bin/rustfilt
+          key: cargo-sbf-bins-${{ runner.os }}
+
+      - uses: actions/cache@v2
+        with:
+          path: ~/.cache/solana
+          key: solana-${{ env.SOLANA_VERSION }}
+
+      - name: Install dependencies
+        run: |
+          ./ci/install-build-deps.sh
+          ./ci/install-program-deps.sh
+          echo "$HOME/.local/share/solana/install/active_release/bin" >> $GITHUB_PATH
+
+      - name: Build and test
+        run: ./ci/cargo-test-sbf.sh token-upgrade

Cargo.lock

@@ -6271,6 +6271,23 @@ dependencies = [
  "spl-token-swap",
 ]
 
+[[package]]
+name = "spl-token-upgrade"
+version = "0.1.0"
+dependencies = [
+ "num-derive",
+ "num-traits",
+ "num_enum",
+ "solana-program",
+ "solana-program-test",
+ "solana-sdk",
+ "spl-token 3.5.0",
+ "spl-token-2022 0.4.3",
+ "spl-token-client",
+ "test-case",
+ "thiserror",
+]
+
 [[package]]
 name = "stateless-asks"
 version = "0.1.0"

Cargo.toml

@@ -31,6 +31,7 @@ members = [
   "token-lending/program",
   "token-swap/program",
   "token-swap/program/fuzz",
+  "token-upgrade/program",
   "token/cli",
   "token/program",
   "token/program-2022",

token-upgrade/program/Cargo.toml

@@ -0,0 +1,33 @@
+[package]
+name = "spl-token-upgrade"
+version = "0.1.0"
+description = "Solana Program Library Token Upgrade"
+authors = ["Solana Maintainers <maintainers@solana.foundation>"]
+repository = "https://github.com/solana-labs/solana-program-library"
+license = "Apache-2.0"
+edition = "2018"
+
+[features]
+no-entrypoint = []
+test-sbf = []
+
+[dependencies]
+num-derive = "0.3"
+num-traits = "0.2"
+num_enum = "0.5.4"
+solana-program = "1.11.6"
+spl-token-2022 = { version = "0.4", path = "../../token/program-2022", features = ["no-entrypoint"] }
+thiserror = "1.0"
+
+[dev-dependencies]
+solana-program-test = "1.11.6"
+solana-sdk = "1.11.6"
+spl-token = { version = "3.5", path = "../../token/program", features = ["no-entrypoint"] }
+spl-token-client = { version = "0.1", path = "../../token/client" }
+test-case = "2.2"
+
+[lib]
+crate-type = ["cdylib", "lib"]
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]

token-upgrade/program/program-id.md

@@ -0,0 +1 @@
+TkupDoNseygccBCjSsrSpMccjwHfTYwcrjpnDSrFDhC

token-upgrade/program/src/entrypoint.rs

@@ -0,0 +1,16 @@
+//! Program entrypoint
+
+#![cfg(not(feature = "no-entrypoint"))]
+
+use solana_program::{
+    account_info::AccountInfo, entrypoint, entrypoint::ProgramResult, pubkey::Pubkey,
+};
+
+entrypoint!(process_instruction);
+fn process_instruction(
+    program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    instruction_data: &[u8],
+) -> ProgramResult {
+    crate::processor::process(program_id, accounts, instruction_data)
+}

token-upgrade/program/src/error.rs

@@ -0,0 +1,29 @@
+//! Error types
+
+use {
+    num_derive::FromPrimitive,
+    solana_program::{decode_error::DecodeError, program_error::ProgramError},
+    thiserror::Error,
+};
+
+/// Errors that may be returned by the program.
+#[derive(Clone, Debug, Eq, Error, FromPrimitive, PartialEq)]
+pub enum TokenUpgradeError {
+    // 0
+    /// Account does not match address derivation
+    #[error("Account does not match address derivation")]
+    InvalidOwner,
+    /// Decimals of original and new token mint do not match
+    #[error("Decimals of original and new token mint do not match")]
+    DecimalsMismatch,
+}
+impl From<TokenUpgradeError> for ProgramError {
+    fn from(e: TokenUpgradeError) -> Self {
+        ProgramError::Custom(e as u32)
+    }
+}
+impl<T> DecodeError<T> for TokenUpgradeError {
+    fn type_of() -> &'static str {
+        "TokenUpgradeError"
+    }
+}

token-upgrade/program/src/instruction.rs

@@ -0,0 +1,77 @@
+//! Program instructions
+
+use {
+    crate::get_token_upgrade_authority_address,
+    num_enum::{IntoPrimitive, TryFromPrimitive},
+    solana_program::{
+        instruction::{AccountMeta, Instruction},
+        pubkey::Pubkey,
+    },
+};
+
+/// Instructions supported by the TokenUpgrade program
+#[derive(Clone, Copy, Debug, TryFromPrimitive, IntoPrimitive)]
+#[repr(u8)]
+pub enum TokenUpgradeInstruction {
+    /// Burns all of the original tokens in the user's account, and transfers the same
+    /// amount of tokens from an account owned by a PDA into another account.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[writeable]` Original token account to burn from
+    ///   1. `[writeable]` Original token mint
+    ///   2. `[writeable]` Escrow of new tokens held by or delegated to PDA at address:
+    ///       `get_token_upgrade_authority_address(original_mint, new_mint, program_id)`
+    ///   3. `[writeable]` New token account to transfer into
+    ///   4. `[]` New token mint
+    ///   5. `[]` Transfer authority (owner or delegate) of new token escrow held by PDA, must be:
+    ///       `get_token_upgrade_authority_address(original_mint, new_mint, program_id)`
+    ///   6. `[]` SPL Token program for original mint
+    ///   7. `[]` SPL Token program for new mint
+    ///   8. `[]` Original token account transfer authority (owner or delegate)
+    ///   9. ..9+M `[signer]` M multisig signer accounts
+    ///
+    /// Data expected by this instruction:
+    ///   None
+    ///
+    Exchange,
+}
+
+/// Create an `Exchange` instruction
+#[allow(clippy::too_many_arguments)]
+pub fn exchange(
+    program_id: &Pubkey,
+    original_account: &Pubkey,
+    original_mint: &Pubkey,
+    new_escrow: &Pubkey,
+    new_account: &Pubkey,
+    new_mint: &Pubkey,
+    original_token_program_id: &Pubkey,
+    new_token_program_id: &Pubkey,
+    original_transfer_authority: &Pubkey,
+    original_multisig_signers: &[&Pubkey],
+) -> Instruction {
+    let escrow_authority = get_token_upgrade_authority_address(original_mint, new_mint, program_id);
+    let mut accounts = Vec::with_capacity(9 + original_multisig_signers.len());
+    accounts.push(AccountMeta::new(*original_account, false));
+    accounts.push(AccountMeta::new(*original_mint, false));
+    accounts.push(AccountMeta::new(*new_escrow, false));
+    accounts.push(AccountMeta::new(*new_account, false));
+    accounts.push(AccountMeta::new(*new_mint, false));
+    accounts.push(AccountMeta::new_readonly(escrow_authority, false));
+    accounts.push(AccountMeta::new_readonly(*original_token_program_id, false));
+    accounts.push(AccountMeta::new_readonly(*new_token_program_id, false));
+    accounts.push(AccountMeta::new_readonly(
+        *original_transfer_authority,
+        original_multisig_signers.is_empty(),
+    ));
+    for signer_pubkey in original_multisig_signers.iter() {
+        accounts.push(AccountMeta::new_readonly(**signer_pubkey, true));
+    }
+
+    Instruction {
+        program_id: *program_id,
+        accounts,
+        data: vec![TokenUpgradeInstruction::Exchange.into()],
+    }
+}

token-upgrade/program/src/lib.rs

@@ -0,0 +1,60 @@
+//! Convention for upgrading tokens from one program to another
+#![deny(missing_docs)]
+#![forbid(unsafe_code)]
+
+mod entrypoint;
+pub mod error;
+pub mod instruction;
+pub mod processor;
+
+// Export current SDK types for downstream users building with a different SDK version
+pub use solana_program;
+use solana_program::pubkey::Pubkey;
+
+solana_program::declare_id!("TkupDoNseygccBCjSsrSpMccjwHfTYwcrjpnDSrFDhC");
+
+const TOKEN_ESCROW_AUTHORITY_SEED: &[u8] = b"token-escrow-authority";
+
+/// Get the upgrade token account authority
+pub fn get_token_upgrade_authority_address(
+    original_mint: &Pubkey,
+    new_mint: &Pubkey,
+    program_id: &Pubkey,
+) -> Pubkey {
+    get_token_upgrade_authority_address_and_bump_seed(original_mint, new_mint, program_id).0
+}
+
+pub(crate) fn get_token_upgrade_authority_address_and_bump_seed(
+    original_mint: &Pubkey,
+    new_mint: &Pubkey,
+    program_id: &Pubkey,
+) -> (Pubkey, u8) {
+    Pubkey::find_program_address(
+        &collect_token_upgrade_authority_seeds(original_mint, new_mint),
+        program_id,
+    )
+}
+
+pub(crate) fn collect_token_upgrade_authority_seeds<'a>(
+    original_mint: &'a Pubkey,
+    new_mint: &'a Pubkey,
+) -> [&'a [u8]; 3] {
+    [
+        TOKEN_ESCROW_AUTHORITY_SEED,
+        original_mint.as_ref(),
+        new_mint.as_ref(),
+    ]
+}
+
+pub(crate) fn collect_token_upgrade_authority_signer_seeds<'a>(
+    original_mint: &'a Pubkey,
+    new_mint: &'a Pubkey,
+    bump_seed: &'a [u8],
+) -> [&'a [u8]; 4] {
+    [
+        TOKEN_ESCROW_AUTHORITY_SEED,
+        original_mint.as_ref(),
+        new_mint.as_ref(),
+        bump_seed,
+    ]
+}

token-upgrade/program/src/processor.rs

@@ -0,0 +1,190 @@
+//! Program state processor
+
+use {
+    crate::{
+        collect_token_upgrade_authority_signer_seeds, error::TokenUpgradeError,
+        get_token_upgrade_authority_address_and_bump_seed, instruction::TokenUpgradeInstruction,
+    },
+    solana_program::{
+        account_info::{next_account_info, AccountInfo},
+        entrypoint::ProgramResult,
+        msg,
+        program::{invoke, invoke_signed},
+        program_error::ProgramError,
+        pubkey::Pubkey,
+    },
+    spl_token_2022::{
+        extension::StateWithExtensions,
+        instruction::decode_instruction_type,
+        state::{Account, Mint},
+    },
+};
+
+fn check_owner(account_info: &AccountInfo, expected_owner: &Pubkey) -> ProgramResult {
+    if account_info.owner != expected_owner {
+        Err(ProgramError::IllegalOwner)
+    } else {
+        Ok(())
+    }
+}
+
+fn burn_original_tokens<'a>(
+    original_token_program: AccountInfo<'a>,
+    source: AccountInfo<'a>,
+    mint: AccountInfo<'a>,
+    authority: AccountInfo<'a>,
+    multisig_signers: &[AccountInfo<'a>],
+    amount: u64,
+    decimals: u8,
+) -> ProgramResult {
+    let multisig_pubkeys = multisig_signers.iter().map(|s| s.key).collect::<Vec<_>>();
+    let ix = spl_token_2022::instruction::burn_checked(
+        original_token_program.key,
+        source.key,
+        mint.key,
+        authority.key,
+        &multisig_pubkeys,
+        amount,
+        decimals,
+    )?;
+    let mut account_infos = vec![source, mint, authority];
+    account_infos.extend_from_slice(multisig_signers);
+    invoke(&ix, &account_infos)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn transfer_new_tokens<'a>(
+    new_token_program: AccountInfo<'a>,
+    source: AccountInfo<'a>,
+    mint: AccountInfo<'a>,
+    destination: AccountInfo<'a>,
+    authority: AccountInfo<'a>,
+    authority_seeds: &[&[u8]],
+    amount: u64,
+    decimals: u8,
+) -> Result<(), ProgramError> {
+    let ix = spl_token_2022::instruction::transfer_checked(
+        new_token_program.key,
+        source.key,
+        mint.key,
+        destination.key,
+        authority.key,
+        &[],
+        amount,
+        decimals,
+    )?;
+    invoke_signed(
+        &ix,
+        &[source, mint, destination, authority],
+        &[authority_seeds],
+    )
+}
+
+fn process_exchange(program_id: &Pubkey, accounts: &[AccountInfo]) -> ProgramResult {
+    let account_info_iter = &mut accounts.iter();
+
+    let original_account_info = next_account_info(account_info_iter)?;
+    let original_mint_info = next_account_info(account_info_iter)?;
+    let new_escrow_info = next_account_info(account_info_iter)?;
+    let new_account_info = next_account_info(account_info_iter)?;
+    let new_mint_info = next_account_info(account_info_iter)?;
+    let new_transfer_authority_info = next_account_info(account_info_iter)?;
+    let original_token_program = next_account_info(account_info_iter)?;
+    let new_token_program = next_account_info(account_info_iter)?;
+    let original_transfer_authority_info = next_account_info(account_info_iter)?;
+
+    // owner checks
+    check_owner(original_account_info, original_token_program.key)?;
+    check_owner(original_mint_info, original_token_program.key)?;
+    check_owner(new_escrow_info, new_token_program.key)?;
+    check_owner(new_account_info, new_token_program.key)?;
+    check_owner(new_mint_info, new_token_program.key)?;
+
+    // PDA derivation check
+    let (expected_escrow_authority, bump_seed) = get_token_upgrade_authority_address_and_bump_seed(
+        original_mint_info.key,
+        new_mint_info.key,
+        program_id,
+    );
+    let bump_seed = [bump_seed];
+    let authority_seeds = collect_token_upgrade_authority_signer_seeds(
+        original_mint_info.key,
+        new_mint_info.key,
+        &bump_seed,
+    );
+    if expected_escrow_authority != *new_transfer_authority_info.key {
+        msg!(
+            "Expected escrow authority {}, received {}",
+            &expected_escrow_authority,
+            new_transfer_authority_info.key
+        );
+        return Err(TokenUpgradeError::InvalidOwner.into());
+    }
+
+    // pull out these values in a block to drop all data before performing CPIs
+    let (token_amount, decimals) = {
+        // check mints are actually mints
+        let original_mint_data = original_mint_info.try_borrow_data()?;
+        let original_mint = StateWithExtensions::<Mint>::unpack(&original_mint_data)?;
+        let new_mint_data = new_mint_info.try_borrow_data()?;
+        let new_mint = StateWithExtensions::<Mint>::unpack(&new_mint_data)?;
+
+        // check accounts are actually accounts
+        let original_account_data = original_account_info.try_borrow_data()?;
+        let original_account = StateWithExtensions::<Account>::unpack(&original_account_data)?;
+        let new_escrow_data = new_escrow_info.try_borrow_data()?;
+        let new_escrow = StateWithExtensions::<Account>::unpack(&new_escrow_data)?;
+        let new_account_data = new_account_info.try_borrow_data()?;
+        let _ = StateWithExtensions::<Account>::unpack(&new_account_data)?;
+
+        let token_amount = original_account.base.amount;
+        if new_escrow.base.amount < token_amount {
+            msg!(
+                "Escrow only has {} tokens, needs at least {}",
+                new_escrow.base.amount,
+                token_amount
+            );
+            return Err(ProgramError::InsufficientFunds);
+        }
+        if original_mint.base.decimals != new_mint.base.decimals {
+            msg!(
+                "Original and new token mint decimals mismatch: original has {} decimals, and new has {}",
+                original_mint.base.decimals,
+                new_mint.base.decimals,
+            );
+            return Err(TokenUpgradeError::DecimalsMismatch.into());
+        }
+
+        (original_account.base.amount, original_mint.base.decimals)
+    };
+
+    burn_original_tokens(
+        original_token_program.clone(),
+        original_account_info.clone(),
+        original_mint_info.clone(),
+        original_transfer_authority_info.clone(),
+        account_info_iter.as_slice(),
+        token_amount,
+        decimals,
+    )?;
+
+    transfer_new_tokens(
+        new_token_program.clone(),
+        new_escrow_info.clone(),
+        new_mint_info.clone(),
+        new_account_info.clone(),
+        new_transfer_authority_info.clone(),
+        &authority_seeds,
+        token_amount,
+        decimals,
+    )?;
+
+    Ok(())
+}
+
+/// Instruction processor
+pub fn process(program_id: &Pubkey, accounts: &[AccountInfo], input: &[u8]) -> ProgramResult {
+    match decode_instruction_type(input)? {
+        TokenUpgradeInstruction::Exchange => process_exchange(program_id, accounts),
+    }
+}

token-upgrade/program/tests/functional.rs

@@ -0,0 +1,412 @@
+// Mark this test as SBF-only due to current `ProgramTest` limitations when CPIing into the system program
+#![cfg(feature = "test-sbf")]
+
+use {
+    solana_program_test::{
+        processor,
+        tokio::{self, sync::Mutex},
+        ProgramTest, ProgramTestContext,
+    },
+    solana_sdk::{
+        instruction::{AccountMeta, InstructionError},
+        pubkey::Pubkey,
+        signature::Signer,
+        signer::keypair::Keypair,
+        transaction::{Transaction, TransactionError},
+    },
+    spl_token_client::{
+        client::{
+            ProgramBanksClient, ProgramBanksClientProcessTransaction, ProgramClient,
+            SendTransaction,
+        },
+        token::Token,
+    },
+    spl_token_upgrade::{
+        error::TokenUpgradeError, get_token_upgrade_authority_address, instruction::exchange,
+    },
+    std::sync::Arc,
+    test_case::test_case,
+};
+
+fn keypair_clone(kp: &Keypair) -> Keypair {
+    Keypair::from_bytes(&kp.to_bytes()).expect("failed to copy keypair")
+}
+
+async fn setup() -> (
+    Arc<Mutex<ProgramTestContext>>,
+    Arc<dyn ProgramClient<ProgramBanksClientProcessTransaction>>,
+    Arc<Keypair>,
+) {
+    let mut program_test = ProgramTest::new(
+        "spl_token_upgrade",
+        spl_token_upgrade::id(),
+        processor!(spl_token_upgrade::processor::process),
+    );
+
+    program_test.prefer_bpf(false); // simplicity in the build
+
+    program_test.add_program(
+        "spl_token_2022",
+        spl_token_2022::id(),
+        processor!(spl_token_2022::processor::Processor::process),
+    );
+    program_test.add_program(
+        "spl_token",
+        spl_token::id(),
+        processor!(spl_token::processor::Processor::process),
+    );
+
+    let context = program_test.start_with_context().await;
+    let payer = Arc::new(keypair_clone(&context.payer));
+    let context = Arc::new(Mutex::new(context));
+
+    let client: Arc<dyn ProgramClient<ProgramBanksClientProcessTransaction>> =
+        Arc::new(ProgramBanksClient::new_from_context(
+            Arc::clone(&context),
+            ProgramBanksClientProcessTransaction,
+        ));
+    (context, client, payer)
+}
+
+async fn setup_mint<T: SendTransaction>(
+    program_id: &Pubkey,
+    mint_authority: &Pubkey,
+    decimals: u8,
+    payer: Arc<Keypair>,
+    client: Arc<dyn ProgramClient<T>>,
+) -> Token<T> {
+    let mint_account = Keypair::new();
+    let token = Token::new(client, program_id, &mint_account.pubkey(), payer);
+    token
+        .create_mint(mint_authority, None, decimals, vec![], &[&mint_account])
+        .await
+        .unwrap();
+    token
+}
+
+#[test_case(spl_token::id(), spl_token_2022::id() ; "upgrade to token-2022")]
+#[test_case(spl_token_2022::id(), spl_token::id() ; "downgrade to token")]
+#[test_case(spl_token::id(), spl_token::id() ; "token to token")]
+#[test_case(spl_token_2022::id(), spl_token_2022::id() ; "token-2022 to token-2022")]
+#[tokio::test]
+async fn success(original_program_id: Pubkey, new_program_id: Pubkey) {
+    let (context, client, payer) = setup().await;
+
+    let wallet = Keypair::new();
+    let mint_authority = Keypair::new();
+    let mint_authority_pubkey = mint_authority.pubkey();
+
+    let decimals = 2;
+    let original_token = setup_mint(
+        &original_program_id,
+        &mint_authority_pubkey,
+        decimals,
+        payer.clone(),
+        client.clone(),
+    )
+    .await;
+    let new_token = setup_mint(
+        &new_program_id,
+        &mint_authority_pubkey,
+        decimals,
+        payer.clone(),
+        client.clone(),
+    )
+    .await;
+
+    let program_escrow = get_token_upgrade_authority_address(
+        original_token.get_address(),
+        new_token.get_address(),
+        &spl_token_upgrade::id(),
+    );
+
+    original_token
+        .create_associated_token_account(&wallet.pubkey())
+        .await
+        .unwrap();
+    let original_account = original_token.get_associated_token_address(&wallet.pubkey());
+    let token_amount = 1_000_000_000_000;
+    original_token
+        .mint_to(
+            &original_account,
+            &mint_authority_pubkey,
+            token_amount,
+            Some(decimals),
+            &[&mint_authority],
+        )
+        .await
+        .unwrap();
+
+    new_token
+        .create_associated_token_account(&wallet.pubkey())
+        .await
+        .unwrap();
+    let new_account = new_token.get_associated_token_address(&wallet.pubkey());
+    new_token
+        .create_associated_token_account(&program_escrow)
+        .await
+        .unwrap();
+    let escrow_account = new_token.get_associated_token_address(&program_escrow);
+    new_token
+        .mint_to(
+            &escrow_account,
+            &mint_authority_pubkey,
+            token_amount,
+            Some(decimals),
+            &[&mint_authority],
+        )
+        .await
+        .unwrap();
+
+    {
+        let mut context = context.lock().await;
+        let transaction = Transaction::new_signed_with_payer(
+            &[exchange(
+                &spl_token_upgrade::id(),
+                &original_account,
+                original_token.get_address(),
+                &escrow_account,
+                &new_account,
+                new_token.get_address(),
+                &original_program_id,
+                &new_program_id,
+                &wallet.pubkey(),
+                &[],
+            )],
+            Some(&context.payer.pubkey()),
+            &[&context.payer, &wallet],
+            context.last_blockhash,
+        );
+        context
+            .banks_client
+            .process_transaction(transaction)
+            .await
+            .unwrap();
+    }
+
+    let original_mint = original_token.get_mint_info().await.unwrap();
+    assert_eq!(original_mint.base.supply, 0);
+    let original_account_info = original_token
+        .get_account_info(&original_account)
+        .await
+        .unwrap();
+    assert_eq!(original_account_info.base.amount, 0);
+    let new_account_info = new_token.get_account_info(&new_account).await.unwrap();
+    assert_eq!(new_account_info.base.amount, token_amount);
+    let escrow_info = new_token.get_account_info(&escrow_account).await.unwrap();
+    assert_eq!(escrow_info.base.amount, 0);
+}
+
+#[test_case(spl_token::id(), spl_token_2022::id() ; "fail upgrade to token-2022")]
+#[tokio::test]
+async fn fail_incorrect_escrow_derivation(original_program_id: Pubkey, new_program_id: Pubkey) {
+    let (context, client, payer) = setup().await;
+
+    let wallet = Keypair::new();
+    let mint_authority = Keypair::new();
+    let mint_authority_pubkey = mint_authority.pubkey();
+
+    let decimals = 2;
+    let original_token = setup_mint(
+        &original_program_id,
+        &mint_authority_pubkey,
+        decimals,
+        payer.clone(),
+        client.clone(),
+    )
+    .await;
+    let new_token = setup_mint(
+        &new_program_id,
+        &mint_authority_pubkey,
+        decimals,
+        payer.clone(),
+        client.clone(),
+    )
+    .await;
+
+    // backwards derivation
+    let program_escrow = get_token_upgrade_authority_address(
+        new_token.get_address(),
+        original_token.get_address(),
+        &spl_token_upgrade::id(),
+    );
+
+    original_token
+        .create_associated_token_account(&wallet.pubkey())
+        .await
+        .unwrap();
+    let original_account = original_token.get_associated_token_address(&wallet.pubkey());
+    let token_amount = 1_000_000_000_000;
+    original_token
+        .mint_to(
+            &original_account,
+            &mint_authority_pubkey,
+            token_amount,
+            Some(decimals),
+            &[&mint_authority],
+        )
+        .await
+        .unwrap();
+
+    new_token
+        .create_associated_token_account(&wallet.pubkey())
+        .await
+        .unwrap();
+    let new_account = new_token.get_associated_token_address(&wallet.pubkey());
+    new_token
+        .create_associated_token_account(&program_escrow)
+        .await
+        .unwrap();
+    let escrow_account = new_token.get_associated_token_address(&program_escrow);
+    new_token
+        .mint_to(
+            &escrow_account,
+            &mint_authority_pubkey,
+            token_amount,
+            Some(decimals),
+            &[&mint_authority],
+        )
+        .await
+        .unwrap();
+
+    let mut instruction = exchange(
+        &spl_token_upgrade::id(),
+        &original_account,
+        original_token.get_address(),
+        &escrow_account,
+        &new_account,
+        new_token.get_address(),
+        &original_program_id,
+        &new_program_id,
+        &wallet.pubkey(),
+        &[],
+    );
+    instruction.accounts[5] = AccountMeta::new_readonly(program_escrow, false);
+    let mut context = context.lock().await;
+    let transaction = Transaction::new_signed_with_payer(
+        &[instruction],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &wallet],
+        context.last_blockhash,
+    );
+    let error = context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .unwrap_err()
+        .unwrap();
+    assert_eq!(
+        error,
+        TransactionError::InstructionError(
+            0,
+            InstructionError::Custom(TokenUpgradeError::InvalidOwner as u32)
+        )
+    );
+}
+
+#[test_case(spl_token::id(), spl_token_2022::id() ; "fail upgrade to token-2022")]
+#[tokio::test]
+async fn fail_decimals_mismatch(original_program_id: Pubkey, new_program_id: Pubkey) {
+    let (context, client, payer) = setup().await;
+
+    let wallet = Keypair::new();
+    let mint_authority = Keypair::new();
+    let mint_authority_pubkey = mint_authority.pubkey();
+
+    // different decimals
+    let original_decimals = 2;
+    let new_decimals = 3;
+
+    let original_token = setup_mint(
+        &original_program_id,
+        &mint_authority_pubkey,
+        original_decimals,
+        payer.clone(),
+        client.clone(),
+    )
+    .await;
+    let new_token = setup_mint(
+        &new_program_id,
+        &mint_authority_pubkey,
+        new_decimals,
+        payer.clone(),
+        client.clone(),
+    )
+    .await;
+
+    let program_escrow = get_token_upgrade_authority_address(
+        original_token.get_address(),
+        new_token.get_address(),
+        &spl_token_upgrade::id(),
+    );
+
+    original_token
+        .create_associated_token_account(&wallet.pubkey())
+        .await
+        .unwrap();
+    let original_account = original_token.get_associated_token_address(&wallet.pubkey());
+    let token_amount = 1_000_000_000_000;
+    original_token
+        .mint_to(
+            &original_account,
+            &mint_authority_pubkey,
+            token_amount,
+            Some(original_decimals),
+            &[&mint_authority],
+        )
+        .await
+        .unwrap();
+
+    new_token
+        .create_associated_token_account(&wallet.pubkey())
+        .await
+        .unwrap();
+    let new_account = new_token.get_associated_token_address(&wallet.pubkey());
+    new_token
+        .create_associated_token_account(&program_escrow)
+        .await
+        .unwrap();
+    let escrow_account = new_token.get_associated_token_address(&program_escrow);
+    new_token
+        .mint_to(
+            &escrow_account,
+            &mint_authority_pubkey,
+            token_amount,
+            Some(new_decimals),
+            &[&mint_authority],
+        )
+        .await
+        .unwrap();
+
+    let mut context = context.lock().await;
+    let transaction = Transaction::new_signed_with_payer(
+        &[exchange(
+            &spl_token_upgrade::id(),
+            &original_account,
+            original_token.get_address(),
+            &escrow_account,
+            &new_account,
+            new_token.get_address(),
+            &original_program_id,
+            &new_program_id,
+            &wallet.pubkey(),
+            &[],
+        )],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &wallet],
+        context.last_blockhash,
+    );
+    let error = context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .unwrap_err()
+        .unwrap();
+    assert_eq!(
+        error,
+        TransactionError::InstructionError(
+            0,
+            InstructionError::Custom(TokenUpgradeError::DecimalsMismatch as u32)
+        )
+    );
+}