p-token: Allow withdraw excess lamports from mint without authority (#51)

* Add mint as authority test

* Add tests

* Fix test comment

* Update test names

p-token/src/processor/withdraw_excess_lamports.rs

@@ -18,7 +18,7 @@ pub fn process_withdraw_excess_lamports(accounts: &[AccountInfo]) -> ProgramResu
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
-    // SAFETY: single mutable borrow to `source_account_info` account data
+    // SAFETY: single immutable borrow to `source_account_info` account data
     let source_data = unsafe { source_account_info.borrow_data_unchecked() };
 
     match source_data.len() {
@@ -36,10 +36,25 @@ pub fn process_withdraw_excess_lamports(accounts: &[AccountInfo]) -> ProgramResu
             // SAFETY: `source_data` has the same length as `Mint`.
             let mint = unsafe { load::<Mint>(source_data)? };
 
-            if let Some(mint_authority) = mint.mint_authority() {
-                validate_owner(mint_authority, authority_info, remaining)?;
-            } else {
-                return Err(TokenError::AuthorityTypeNotSupported.into());
+            match mint.mint_authority() {
+                Some(mint_authority) => {
+                    validate_owner(mint_authority, authority_info, remaining)?;
+                }
+                None if source_account_info == authority_info => {
+                    // Comparing whether the AccountInfo's "point" to the same account or
+                    // not - this is a faster comparison since it just checks the internal
+                    // raw pointer.
+                    //
+                    // This is a special case where there is no mint authority set but the mint
+                    // account is the same as the authority account and, therefore, needs to be
+                    // a signer.
+                    if !authority_info.is_signer() {
+                        return Err(ProgramError::MissingRequiredSignature);
+                    }
+                }
+                _ => {
+                    return Err(TokenError::AuthorityTypeNotSupported.into());
+                }
             }
         }
         Multisig::LEN => {

p-token/tests/withdraw_excess_lamports.rs

@@ -8,6 +8,7 @@ use {
     solana_program_test::{tokio, BanksClientError, ProgramTest},
     solana_sdk::{
         instruction::InstructionError,
+        program_pack::Pack,
         pubkey::Pubkey,
         signature::{Keypair, Signer},
         system_instruction,
@@ -758,3 +759,381 @@ async fn fail_withdraw_excess_lamports_from_multisig_missing_signer(token_progra
         ))
     );
 }
+
+#[test_case::test_case(TOKEN_PROGRAM_ID ; "p-token")]
+#[tokio::test]
+async fn withdraw_excess_lamports_from_mint_with_no_authority(token_program: Pubkey) {
+    let context = ProgramTest::new("pinocchio_token_program", TOKEN_PROGRAM_ID, None)
+        .start_with_context()
+        .await;
+
+    let excess_lamports = 4_000_000_000_000;
+
+    // Given a mint authority, freeze authority and an account keypair.
+
+    let mint_authority = Keypair::new();
+    let freeze_authority = Pubkey::new_unique();
+    let mint_account = Keypair::new();
+    let account_pubkey = mint_account.pubkey();
+
+    let account_size = size_of::<Mint>();
+    let rent = context.banks_client.get_rent().await.unwrap();
+
+    let mut initialize_ix = spl_token::instruction::initialize_mint(
+        &spl_token::ID,
+        &mint_account.pubkey(),
+        &mint_authority.pubkey(),
+        Some(&freeze_authority),
+        0,
+    )
+    .unwrap();
+    // Switches the program id to the token program.
+    initialize_ix.program_id = token_program;
+
+    // And we initialize a mint account with excess lamports.
+
+    let instructions = vec![
+        system_instruction::create_account(
+            &context.payer.pubkey(),
+            &mint_account.pubkey(),
+            rent.minimum_balance(account_size) + excess_lamports,
+            account_size as u64,
+            &token_program,
+        ),
+        initialize_ix,
+    ];
+
+    let tx = Transaction::new_signed_with_payer(
+        &instructions,
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_account],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+
+    let account = context
+        .banks_client
+        .get_account(mint_account.pubkey())
+        .await
+        .unwrap();
+
+    assert!(account.is_some());
+
+    let account = account.unwrap();
+    assert_eq!(
+        account.lamports,
+        rent.minimum_balance(account_size) + excess_lamports
+    );
+
+    // And we remove the mint authority.
+
+    let mut set_authority_ix = spl_token::instruction::set_authority(
+        &spl_token::ID,
+        &mint_account.pubkey(),
+        None,
+        spl_token::instruction::AuthorityType::MintTokens,
+        &mint_authority.pubkey(),
+        &[&mint_authority.pubkey()],
+    )
+    .unwrap();
+    // Switches the program id to the token program.
+    set_authority_ix.program_id = token_program;
+
+    let tx = Transaction::new_signed_with_payer(
+        &[set_authority_ix],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_authority],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+
+    let account = context
+        .banks_client
+        .get_account(mint_account.pubkey())
+        .await
+        .unwrap();
+
+    assert!(account.is_some());
+
+    let account = account.unwrap();
+    let account = spl_token::state::Mint::unpack(&account.data).unwrap();
+
+    assert!(account.mint_authority.is_none());
+
+    // When we withdraw the excess lamports with no authority.
+
+    let destination = Pubkey::new_unique();
+
+    let mut withdraw_ix = spl_token_2022::instruction::withdraw_excess_lamports(
+        &spl_token_2022::ID,
+        &account_pubkey,
+        &destination,
+        &mint_account.pubkey(),
+        &[],
+    )
+    .unwrap();
+    // Switches the program id to the token program.
+    withdraw_ix.program_id = token_program;
+
+    let tx = Transaction::new_signed_with_payer(
+        &[withdraw_ix],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_account],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+
+    // Then the destination account has the excess lamports.
+
+    let destination = context.banks_client.get_account(destination).await.unwrap();
+
+    assert!(destination.is_some());
+
+    let destination = destination.unwrap();
+    assert_eq!(destination.lamports, excess_lamports);
+}
+
+#[test_case::test_case(TOKEN_PROGRAM_ID ; "p-token")]
+#[tokio::test]
+async fn fail_withdraw_excess_lamports_from_mint_with_authority_and_mint_as_signer(
+    token_program: Pubkey,
+) {
+    let context = ProgramTest::new("pinocchio_token_program", TOKEN_PROGRAM_ID, None)
+        .start_with_context()
+        .await;
+
+    let excess_lamports = 4_000_000_000_000;
+
+    // Given a mint authority, freeze authority and an account keypair.
+
+    let mint_authority = Keypair::new();
+    let freeze_authority = Pubkey::new_unique();
+    let mint_account = Keypair::new();
+    let account_pubkey = mint_account.pubkey();
+
+    let account_size = size_of::<Mint>();
+    let rent = context.banks_client.get_rent().await.unwrap();
+
+    let mut initialize_ix = spl_token::instruction::initialize_mint(
+        &spl_token::ID,
+        &mint_account.pubkey(),
+        &mint_authority.pubkey(),
+        Some(&freeze_authority),
+        0,
+    )
+    .unwrap();
+    // Switches the program id to the token program.
+    initialize_ix.program_id = token_program;
+
+    // And we initialize a mint account with excess lamports.
+
+    let instructions = vec![
+        system_instruction::create_account(
+            &context.payer.pubkey(),
+            &mint_account.pubkey(),
+            rent.minimum_balance(account_size) + excess_lamports,
+            account_size as u64,
+            &token_program,
+        ),
+        initialize_ix,
+    ];
+
+    let tx = Transaction::new_signed_with_payer(
+        &instructions,
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_account],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+
+    let account = context
+        .banks_client
+        .get_account(mint_account.pubkey())
+        .await
+        .unwrap();
+
+    assert!(account.is_some());
+
+    let account = account.unwrap();
+    assert_eq!(
+        account.lamports,
+        rent.minimum_balance(account_size) + excess_lamports
+    );
+
+    // When we try to withdraw the excess lamports with the mint as authority.
+
+    let destination = Pubkey::new_unique();
+
+    let mut withdraw_ix = spl_token_2022::instruction::withdraw_excess_lamports(
+        &spl_token_2022::ID,
+        &account_pubkey,
+        &destination,
+        &mint_account.pubkey(),
+        &[],
+    )
+    .unwrap();
+    // Switches the program id to the token program.
+    withdraw_ix.program_id = token_program;
+
+    let tx = Transaction::new_signed_with_payer(
+        &[withdraw_ix],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_account],
+        context.last_blockhash,
+    );
+    let error = context
+        .banks_client
+        .process_transaction(tx)
+        .await
+        .unwrap_err();
+
+    // Then we expect an error.
+
+    assert_matches!(
+        error,
+        BanksClientError::TransactionError(TransactionError::InstructionError(
+            _,
+            InstructionError::Custom(4) // TokenError::OwnerMismatch
+        ))
+    );
+}
+
+#[test_case::test_case(TOKEN_PROGRAM_ID ; "p-token")]
+#[tokio::test]
+async fn fail_withdraw_excess_lamports_from_mint_with_no_authority_and_authority_signer(
+    token_program: Pubkey,
+) {
+    let context = ProgramTest::new("pinocchio_token_program", TOKEN_PROGRAM_ID, None)
+        .start_with_context()
+        .await;
+
+    let excess_lamports = 4_000_000_000_000;
+
+    // Given a mint authority, freeze authority and an account keypair.
+
+    let mint_authority = Keypair::new();
+    let freeze_authority = Pubkey::new_unique();
+    let mint_account = Keypair::new();
+    let account_pubkey = mint_account.pubkey();
+
+    let account_size = size_of::<Mint>();
+    let rent = context.banks_client.get_rent().await.unwrap();
+
+    let mut initialize_ix = spl_token::instruction::initialize_mint(
+        &spl_token::ID,
+        &mint_account.pubkey(),
+        &mint_authority.pubkey(),
+        Some(&freeze_authority),
+        0,
+    )
+    .unwrap();
+    // Switches the program id to the token program.
+    initialize_ix.program_id = token_program;
+
+    // And we initialize a mint account with excess lamports.
+
+    let instructions = vec![
+        system_instruction::create_account(
+            &context.payer.pubkey(),
+            &mint_account.pubkey(),
+            rent.minimum_balance(account_size) + excess_lamports,
+            account_size as u64,
+            &token_program,
+        ),
+        initialize_ix,
+    ];
+
+    let tx = Transaction::new_signed_with_payer(
+        &instructions,
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_account],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+
+    let account = context
+        .banks_client
+        .get_account(mint_account.pubkey())
+        .await
+        .unwrap();
+
+    assert!(account.is_some());
+
+    let account = account.unwrap();
+    assert_eq!(
+        account.lamports,
+        rent.minimum_balance(account_size) + excess_lamports
+    );
+
+    // And we remove the mint authority.
+
+    let mut set_authority_ix = spl_token::instruction::set_authority(
+        &spl_token::ID,
+        &mint_account.pubkey(),
+        None,
+        spl_token::instruction::AuthorityType::MintTokens,
+        &mint_authority.pubkey(),
+        &[&mint_authority.pubkey()],
+    )
+    .unwrap();
+    // Switches the program id to the token program.
+    set_authority_ix.program_id = token_program;
+
+    let tx = Transaction::new_signed_with_payer(
+        &[set_authority_ix],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_authority],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+
+    let account = context
+        .banks_client
+        .get_account(mint_account.pubkey())
+        .await
+        .unwrap();
+
+    assert!(account.is_some());
+
+    let account = account.unwrap();
+    let account = spl_token::state::Mint::unpack(&account.data).unwrap();
+
+    assert!(account.mint_authority.is_none());
+
+    // When we try to withdraw the excess lamports with the "old" mint authority.
+
+    let destination = Pubkey::new_unique();
+
+    let mut withdraw_ix = spl_token_2022::instruction::withdraw_excess_lamports(
+        &spl_token_2022::ID,
+        &account_pubkey,
+        &destination,
+        &mint_authority.pubkey(),
+        &[],
+    )
+    .unwrap();
+    // Switches the program id to the token program.
+    withdraw_ix.program_id = token_program;
+
+    let tx = Transaction::new_signed_with_payer(
+        &[withdraw_ix],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_authority],
+        context.last_blockhash,
+    );
+    let error = context
+        .banks_client
+        .process_transaction(tx)
+        .await
+        .unwrap_err();
+
+    // Then we expect an error.
+
+    assert_matches!(
+        error,
+        BanksClientError::TransactionError(TransactionError::InstructionError(
+            _,
+            InstructionError::Custom(15) // TokenError::AuthorityTypeNotSupported
+        ))
+    );
+}