From 937993cff886986ad41ffbc4f7f07aabfd53387d Mon Sep 17 00:00:00 2001 From: hanako mumei <81144685+2501babe@users.noreply.github.com> Date: Thu, 13 Oct 2022 14:58:34 -0700 Subject: [PATCH] token-2022: implement CpiGuard (this is a reversion of a previous reversion of the original CpiGuard commit) --- Cargo.lock | 1 + token/client/src/token.rs | 48 +- token/program-2022-test/Cargo.toml | 1 + token/program-2022-test/build.rs | 35 +- token/program-2022-test/tests/cpi_guard.rs | 767 ++++++++++++++++++ token/program-2022/src/error.rs | 51 +- .../src/extension/cpi_guard/instruction.rs | 100 +++ .../src/extension/cpi_guard/mod.rs | 39 + .../src/extension/cpi_guard/processor.rs | 98 +++ .../extension/memo_transfer/instruction.rs | 2 +- token/program-2022/src/extension/mod.rs | 9 +- token/program-2022/src/instruction.rs | 9 + token/program-2022/src/processor.rs | 78 +- 13 files changed, 1208 insertions(+), 30 deletions(-) create mode 100644 token/program-2022-test/tests/cpi_guard.rs create mode 100644 token/program-2022/src/extension/cpi_guard/instruction.rs create mode 100644 token/program-2022/src/extension/cpi_guard/mod.rs create mode 100644 token/program-2022/src/extension/cpi_guard/processor.rs diff --git a/Cargo.lock b/Cargo.lock index 2b8f7a86244..b5de229a323 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6196,6 +6196,7 @@ dependencies = [ "solana-program-test", "solana-sdk", "spl-associated-token-account 1.1.1", + "spl-instruction-padding", "spl-memo 3.0.1", "spl-token-2022 0.4.3", "spl-token-client", diff --git a/token/client/src/token.rs b/token/client/src/token.rs index dd69be0dfe5..8fda770349e 100644 --- a/token/client/src/token.rs +++ b/token/client/src/token.rs @@ -20,8 +20,8 @@ use { }, spl_token_2022::{ extension::{ - confidential_transfer, default_account_state, interest_bearing_mint, memo_transfer, - transfer_fee, ExtensionType, StateWithExtensionsOwned, + confidential_transfer, cpi_guard, default_account_state, interest_bearing_mint, + memo_transfer, transfer_fee, ExtensionType, StateWithExtensionsOwned, }, instruction, solana_zk_token_sdk::{ @@ -1388,6 +1388,50 @@ where .await } + /// Prevent unsafe usage of token account through CPI + pub async fn enable_cpi_guard( + &self, + account: &Pubkey, + authority: &Pubkey, + signing_keypairs: &S, + ) -> TokenResult { + let signing_pubkeys = signing_keypairs.pubkeys(); + let multisig_signers = self.get_multisig_signers(authority, &signing_pubkeys); + + self.process_ixs( + &[cpi_guard::instruction::enable_cpi_guard( + &self.program_id, + account, + authority, + &multisig_signers, + )?], + signing_keypairs, + ) + .await + } + + /// Stop preventing unsafe usage of token account through CPI + pub async fn disable_cpi_guard( + &self, + account: &Pubkey, + authority: &Pubkey, + signing_keypairs: &S, + ) -> TokenResult { + let signing_pubkeys = signing_keypairs.pubkeys(); + let multisig_signers = self.get_multisig_signers(authority, &signing_pubkeys); + + self.process_ixs( + &[cpi_guard::instruction::disable_cpi_guard( + &self.program_id, + account, + authority, + &multisig_signers, + )?], + signing_keypairs, + ) + .await + } + /// Update interest rate pub async fn update_interest_rate( &self, diff --git a/token/program-2022-test/Cargo.toml b/token/program-2022-test/Cargo.toml index 6c115e3c31e..8da9f7717bf 100644 --- a/token/program-2022-test/Cargo.toml +++ b/token/program-2022-test/Cargo.toml @@ -21,4 +21,5 @@ solana-sdk = "=1.14.4" spl-associated-token-account = { version = "1.1", path = "../../associated-token-account/program" } spl-memo = { version = "3.0.1", path = "../../memo/program", features = ["no-entrypoint"] } spl-token-2022 = { version = "0.4", path="../program-2022", features = ["no-entrypoint"] } +spl-instruction-padding = { version = "0.1.0", path="../../instruction-padding/program", features = ["no-entrypoint"] } spl-token-client = { version = "0.2.1", path = "../client" } diff --git a/token/program-2022-test/build.rs b/token/program-2022-test/build.rs index a166f6661aa..12e0f75b217 100644 --- a/token/program-2022-test/build.rs +++ b/token/program-2022-test/build.rs @@ -31,24 +31,37 @@ fn rerun_if_changed(directory: &Path) { fn main() { let cwd = env::current_dir().expect("Unable to get current working directory"); + let spl_token_2022_dir = cwd .parent() .expect("Unable to get parent directory of current working dir") .join("program-2022"); rerun_if_changed(&spl_token_2022_dir); + + let instruction_padding_dir = cwd + .parent() + .expect("Unable to get parent directory of current working dir") + .parent() + .expect("Unable to get grandparent directory of current working dir") + .join("instruction-padding") + .join("program"); + rerun_if_changed(&instruction_padding_dir); + println!("cargo:rerun-if-changed=build.rs"); - let spl_token_2022_toml = spl_token_2022_dir.join("Cargo.toml"); - let spl_token_2022_toml = format!("{}", spl_token_2022_toml.display()); - let args = vec!["build-sbf", "--manifest-path", &spl_token_2022_toml]; - let output = Command::new("cargo") - .args(&args) - .output() - .expect("Error running cargo build-sbf"); - if let Ok(output_str) = std::str::from_utf8(&output.stdout) { - let subs = output_str.split('\n'); - for sub in subs { - println!("cargo:warning=(not a warning) {}", sub); + for program_dir in [spl_token_2022_dir, instruction_padding_dir] { + let program_toml = program_dir.join("Cargo.toml"); + let program_toml = format!("{}", program_toml.display()); + let args = vec!["build-sbf", "--manifest-path", &program_toml]; + let output = Command::new("cargo") + .args(&args) + .output() + .expect("Error running cargo build-sbf"); + if let Ok(output_str) = std::str::from_utf8(&output.stdout) { + let subs = output_str.split('\n'); + for sub in subs { + println!("cargo:warning=(not a warning) {}", sub); + } } } diff --git a/token/program-2022-test/tests/cpi_guard.rs b/token/program-2022-test/tests/cpi_guard.rs new file mode 100644 index 00000000000..162734f61ba --- /dev/null +++ b/token/program-2022-test/tests/cpi_guard.rs @@ -0,0 +1,767 @@ +#![cfg(feature = "test-sbf")] + +mod program_test; +use { + program_test::{keypair_clone, TestContext, TokenContext}, + solana_program_test::{ + processor, + tokio::{self, sync::Mutex}, + ProgramTest, + }, + solana_sdk::{ + instruction::InstructionError, pubkey::Pubkey, signature::Signer, signer::keypair::Keypair, + transaction::TransactionError, transport::TransportError, + }, + spl_instruction_padding::instruction::wrap_instruction, + spl_token_2022::{ + error::TokenError, + extension::{ + cpi_guard::{self, CpiGuard}, + ExtensionType, + }, + instruction::{self, AuthorityType}, + processor::Processor as SplToken2022Processor, + }, + spl_token_client::{ + client::ProgramBanksClientProcessTransaction, + token::{Token, TokenError as TokenClientError}, + }, + std::sync::Arc, +}; + +// set up a bank and bank client with spl token 2022 and the instruction padder +// also creates a token with no extensions and inits two token accounts +async fn make_context() -> TestContext { + // TODO this may be removed when we upgrade to a solana version with a fixed `get_stack_height()` stub + if std::env::var("BPF_OUT_DIR").is_err() && std::env::var("SBF_OUT_DIR").is_err() { + panic!("CpiGuard tests MUST be invoked with `cargo test-sbf`, NOT `cargo test --feature test-sbf`. \ + In a non-BPF context, `get_stack_height()` always returns 0, and all tests WILL fail."); + } + + let mut program_test = ProgramTest::new( + "spl_token_2022", + spl_token_2022::id(), + processor!(SplToken2022Processor::process), + ); + + program_test.add_program( + "spl_instruction_padding", + spl_instruction_padding::id(), + processor!(spl_instruction_padding::processor::process), + ); + + let program_context = program_test.start_with_context().await; + let program_context = Arc::new(Mutex::new(program_context)); + + let mut test_context = TestContext { + context: program_context, + token_context: None, + }; + + test_context.init_token_with_mint(vec![]).await.unwrap(); + let token_context = test_context.token_context.as_ref().unwrap(); + + token_context + .token + .create_auxiliary_token_account_with_extension_space( + &token_context.alice, + &token_context.alice.pubkey(), + vec![ExtensionType::CpiGuard], + ) + .await + .unwrap(); + + token_context + .token + .create_auxiliary_token_account(&token_context.bob, &token_context.bob.pubkey()) + .await + .unwrap(); + + test_context +} + +fn client_error(token_error: TokenError) -> TokenClientError { + TokenClientError::Client(Box::new(TransportError::TransactionError( + TransactionError::InstructionError(0, InstructionError::Custom(token_error as u32)), + ))) +} + +#[tokio::test] +async fn test_cpi_guard_enable_disable() { + let context = make_context().await; + let TokenContext { + token, alice, bob, .. + } = context.token_context.unwrap(); + + // enable guard properly + token + .enable_cpi_guard(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + // guard is enabled + let alice_state = token.get_account_info(&alice.pubkey()).await.unwrap(); + let extension = alice_state.get_extension::().unwrap(); + assert!(bool::from(extension.lock_cpi)); + + // attempt to disable through cpi. this fails + let error = token + .process_ixs( + &[wrap_instruction( + spl_instruction_padding::id(), + cpi_guard::instruction::disable_cpi_guard( + &spl_token_2022::id(), + &alice.pubkey(), + &alice.pubkey(), + &[], + ) + .unwrap(), + vec![], + 0, + ) + .unwrap()], + &[&alice], + ) + .await + .unwrap_err(); + assert_eq!(error, client_error(TokenError::CpiGuardSettingsLocked)); + + // guard remains enabled + let alice_state = token.get_account_info(&alice.pubkey()).await.unwrap(); + let extension = alice_state.get_extension::().unwrap(); + assert!(bool::from(extension.lock_cpi)); + + // disable guard properly + token + .disable_cpi_guard(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + // guard is disabled + let alice_state = token.get_account_info(&alice.pubkey()).await.unwrap(); + let extension = alice_state.get_extension::().unwrap(); + assert!(!bool::from(extension.lock_cpi)); + + // attempt to enable through cpi. this fails + let error = token + .process_ixs( + &[wrap_instruction( + spl_instruction_padding::id(), + cpi_guard::instruction::enable_cpi_guard( + &spl_token_2022::id(), + &alice.pubkey(), + &alice.pubkey(), + &[], + ) + .unwrap(), + vec![], + 0, + ) + .unwrap()], + &[&alice], + ) + .await + .unwrap_err(); + assert_eq!(error, client_error(TokenError::CpiGuardSettingsLocked)); + + // guard remains disabled + let alice_state = token.get_account_info(&alice.pubkey()).await.unwrap(); + let extension = alice_state.get_extension::().unwrap(); + assert!(!bool::from(extension.lock_cpi)); + + // enable works with realloc + token + .reallocate( + &bob.pubkey(), + &bob.pubkey(), + &[ExtensionType::CpiGuard], + &[&bob], + ) + .await + .unwrap(); + + token + .enable_cpi_guard(&bob.pubkey(), &bob.pubkey(), &[&bob]) + .await + .unwrap(); + + let bob_state = token.get_account_info(&bob.pubkey()).await.unwrap(); + let extension = bob_state.get_extension::().unwrap(); + assert!(bool::from(extension.lock_cpi)); +} + +#[tokio::test] +async fn test_cpi_guard_transfer() { + let context = make_context().await; + let TokenContext { + token, + token_unchecked, + mint_authority, + alice, + bob, + .. + } = context.token_context.unwrap(); + + let mk_transfer = |authority, do_checked| { + wrap_instruction( + spl_instruction_padding::id(), + if do_checked { + instruction::transfer_checked( + &spl_token_2022::id(), + &alice.pubkey(), + token.get_address(), + &bob.pubkey(), + &authority, + &[], + 1, + 9, + ) + .unwrap() + } else { + #[allow(deprecated)] + instruction::transfer( + &spl_token_2022::id(), + &alice.pubkey(), + &bob.pubkey(), + &authority, + &[], + 1, + ) + .unwrap() + }, + vec![], + 0, + ) + .unwrap() + }; + + let mut amount = 100; + token + .mint_to( + &alice.pubkey(), + &mint_authority.pubkey(), + amount, + &[&mint_authority], + ) + .await + .unwrap(); + + for do_checked in [true, false] { + let token_obj = if do_checked { &token } else { &token_unchecked }; + token_obj + .enable_cpi_guard(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + // transfer works normally with cpi guard enabled + token_obj + .transfer( + &alice.pubkey(), + &bob.pubkey(), + &alice.pubkey(), + 1, + &[&alice], + ) + .await + .unwrap(); + amount -= 1; + + let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap(); + assert_eq!(alice_state.base.amount, amount); + + // user-auth cpi transfer with cpi guard doesnt work + let error = token_obj + .process_ixs(&[mk_transfer(alice.pubkey(), do_checked)], &[&alice]) + .await + .unwrap_err(); + assert_eq!(error, client_error(TokenError::CpiGuardTransferBlocked)); + + let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap(); + assert_eq!(alice_state.base.amount, amount); + + // delegate-auth cpi transfer with cpi guard works + token_obj + .approve( + &alice.pubkey(), + &bob.pubkey(), + &alice.pubkey(), + 1, + &[&alice], + ) + .await + .unwrap(); + + token_obj + .process_ixs(&[mk_transfer(bob.pubkey(), do_checked)], &[&bob]) + .await + .unwrap(); + amount -= 1; + + let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap(); + assert_eq!(alice_state.base.amount, amount); + + // transfer still works through cpi with cpi guard off + token_obj + .disable_cpi_guard(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + token_obj + .process_ixs(&[mk_transfer(alice.pubkey(), do_checked)], &[&alice]) + .await + .unwrap(); + amount -= 1; + + let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap(); + assert_eq!(alice_state.base.amount, amount); + } +} + +#[tokio::test] +async fn test_cpi_guard_burn() { + let context = make_context().await; + let TokenContext { + token, + token_unchecked, + mint_authority, + alice, + bob, + .. + } = context.token_context.unwrap(); + + let mk_burn = |authority, do_checked| { + wrap_instruction( + spl_instruction_padding::id(), + if do_checked { + instruction::burn_checked( + &spl_token_2022::id(), + &alice.pubkey(), + token.get_address(), + &authority, + &[], + 1, + 9, + ) + .unwrap() + } else { + instruction::burn( + &spl_token_2022::id(), + &alice.pubkey(), + token.get_address(), + &authority, + &[], + 1, + ) + .unwrap() + }, + vec![], + 0, + ) + .unwrap() + }; + + let mut amount = 100; + token + .mint_to( + &alice.pubkey(), + &mint_authority.pubkey(), + amount, + &[&mint_authority], + ) + .await + .unwrap(); + + for do_checked in [true, false] { + let token_obj = if do_checked { &token } else { &token_unchecked }; + token_obj + .enable_cpi_guard(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + // burn works normally with cpi guard enabled + token_obj + .burn(&alice.pubkey(), &alice.pubkey(), 1, &[&alice]) + .await + .unwrap(); + amount -= 1; + + let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap(); + assert_eq!(alice_state.base.amount, amount); + + // user-auth cpi burn with cpi guard doesnt work + let error = token_obj + .process_ixs(&[mk_burn(alice.pubkey(), do_checked)], &[&alice]) + .await + .unwrap_err(); + assert_eq!(error, client_error(TokenError::CpiGuardBurnBlocked)); + + let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap(); + assert_eq!(alice_state.base.amount, amount); + + // delegate-auth cpi burn with cpi guard works + token_obj + .approve( + &alice.pubkey(), + &bob.pubkey(), + &alice.pubkey(), + 1, + &[&alice], + ) + .await + .unwrap(); + + token_obj + .process_ixs(&[mk_burn(bob.pubkey(), do_checked)], &[&bob]) + .await + .unwrap(); + amount -= 1; + + let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap(); + assert_eq!(alice_state.base.amount, amount); + + // burn still works through cpi with cpi guard off + token_obj + .disable_cpi_guard(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + token_obj + .process_ixs(&[mk_burn(alice.pubkey(), do_checked)], &[&alice]) + .await + .unwrap(); + amount -= 1; + + let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap(); + assert_eq!(alice_state.base.amount, amount); + } +} + +#[tokio::test] +async fn test_cpi_guard_approve() { + let context = make_context().await; + let TokenContext { + token, + token_unchecked, + alice, + bob, + .. + } = context.token_context.unwrap(); + + let mk_approve = |do_checked| { + wrap_instruction( + spl_instruction_padding::id(), + if do_checked { + instruction::approve_checked( + &spl_token_2022::id(), + &alice.pubkey(), + token.get_address(), + &bob.pubkey(), + &alice.pubkey(), + &[], + 1, + 9, + ) + .unwrap() + } else { + instruction::approve( + &spl_token_2022::id(), + &alice.pubkey(), + &bob.pubkey(), + &alice.pubkey(), + &[], + 1, + ) + .unwrap() + }, + vec![], + 0, + ) + .unwrap() + }; + + for do_checked in [true, false] { + let token_obj = if do_checked { &token } else { &token_unchecked }; + token_obj + .enable_cpi_guard(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + // approve works normally with cpi guard enabled + token_obj + .approve( + &alice.pubkey(), + &bob.pubkey(), + &alice.pubkey(), + 1, + &[&alice], + ) + .await + .unwrap(); + + token_obj + .revoke(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + // approve doesnt work through cpi + let error = token_obj + .process_ixs(&[mk_approve(do_checked)], &[&alice]) + .await + .unwrap_err(); + assert_eq!(error, client_error(TokenError::CpiGuardApproveBlocked)); + + // approve still works through cpi with cpi guard off + token_obj + .disable_cpi_guard(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + token_obj + .process_ixs(&[mk_approve(do_checked)], &[&alice]) + .await + .unwrap(); + + token_obj + .revoke(&alice.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + } +} + +async fn make_close_test_account( + token: &Token, + owner: &S, + authority: Option, +) -> Pubkey { + let account = Keypair::new(); + + token + .create_auxiliary_token_account_with_extension_space( + &account, + &owner.pubkey(), + vec![ExtensionType::CpiGuard], + ) + .await + .unwrap(); + + if authority.is_some() { + token + .set_authority( + &account.pubkey(), + &owner.pubkey(), + authority.as_ref(), + AuthorityType::CloseAccount, + &[owner], + ) + .await + .unwrap(); + } + + token + .enable_cpi_guard(&account.pubkey(), &owner.pubkey(), &[owner]) + .await + .unwrap(); + + account.pubkey() +} + +#[tokio::test] +async fn test_cpi_guard_close_account() { + let context = make_context().await; + let TokenContext { + token, alice, bob, .. + } = context.token_context.unwrap(); + + let mk_close = |account, destination, authority| { + wrap_instruction( + spl_instruction_padding::id(), + instruction::close_account( + &spl_token_2022::id(), + &account, + &destination, + &authority, + &[], + ) + .unwrap(), + vec![], + 0, + ) + .unwrap() + }; + + // test closing through owner and closing through close authority + // the result should be the same eitehr way + for maybe_close_authority in [None, Some(bob.pubkey())] { + let authority = if maybe_close_authority.is_none() { + &alice + } else { + &bob + }; + + // closing normally works + let account = make_close_test_account(&token, &alice, maybe_close_authority).await; + token + .close_account(&account, &bob.pubkey(), &authority.pubkey(), &[authority]) + .await + .unwrap(); + + // cpi close with guard enabled fails if lamports diverted to third party + let account = make_close_test_account(&token, &alice, maybe_close_authority).await; + let error = token + .process_ixs( + &[mk_close(account, bob.pubkey(), authority.pubkey())], + &[authority], + ) + .await + .unwrap_err(); + assert_eq!(error, client_error(TokenError::CpiGuardCloseAccountBlocked)); + + // but close suceeds if lamports are returned to owner + token + .process_ixs( + &[mk_close(account, alice.pubkey(), authority.pubkey())], + &[authority], + ) + .await + .unwrap(); + + // close still works through cpi when guard disabled + let account = make_close_test_account(&token, &alice, maybe_close_authority).await; + token + .disable_cpi_guard(&account, &alice.pubkey(), &[&alice]) + .await + .unwrap(); + + token + .process_ixs( + &[mk_close(account, bob.pubkey(), authority.pubkey())], + &[authority], + ) + .await + .unwrap(); + } +} + +#[derive(Copy, Clone, Debug, PartialEq)] +enum SetAuthTest { + ChangeOwner, + AddCloseAuth, + ChangeCloseAuth, + RemoveCloseAuth, +} + +#[tokio::test] +async fn test_cpi_guard_set_authority() { + let context = make_context().await; + let TokenContext { + token, alice, bob, .. + } = context.token_context.unwrap(); + + // the behavior of cpi guard and close authority is so complicated that its best to test all cases exhaustively + let mut states = vec![]; + for action in [ + SetAuthTest::ChangeOwner, + SetAuthTest::AddCloseAuth, + SetAuthTest::ChangeCloseAuth, + SetAuthTest::RemoveCloseAuth, + ] { + for enable_cpi_guard in [true, false] { + for do_in_cpi in [true, false] { + states.push((action, enable_cpi_guard, do_in_cpi)); + } + } + } + + for state in states { + let (action, enable_cpi_guard, do_in_cpi) = state; + + // make a new account + let account = Keypair::new(); + token + .create_auxiliary_token_account_with_extension_space( + &account, + &alice.pubkey(), + vec![ExtensionType::CpiGuard], + ) + .await + .unwrap(); + + // turn on cpi guard if we are testing that case + // all actions with cpi guard off should succeed unconditionally + // so half of these tests are backwards compat checks + if enable_cpi_guard { + token + .enable_cpi_guard(&account.pubkey(), &alice.pubkey(), &[&alice]) + .await + .unwrap(); + } + + // if we are changing or removing close auth, we need to have one to change/remove + if action == SetAuthTest::ChangeCloseAuth || action == SetAuthTest::RemoveCloseAuth { + token + .set_authority( + &account.pubkey(), + &alice.pubkey(), + Some(&bob.pubkey()), + AuthorityType::CloseAccount, + &[&alice], + ) + .await + .unwrap(); + } + + // this produces the token instruction we want to execute + let (current_authority, new_authority) = match action { + SetAuthTest::ChangeOwner | SetAuthTest::AddCloseAuth => { + (keypair_clone(&alice), Some(bob.pubkey())) + } + SetAuthTest::ChangeCloseAuth => (keypair_clone(&bob), Some(alice.pubkey())), + SetAuthTest::RemoveCloseAuth => (keypair_clone(&bob), None), + }; + let token_instruction = instruction::set_authority( + &spl_token_2022::id(), + &account.pubkey(), + new_authority.as_ref(), + if action == SetAuthTest::ChangeOwner { + AuthorityType::AccountOwner + } else { + AuthorityType::CloseAccount + }, + ¤t_authority.pubkey(), + &[], + ) + .unwrap(); + + // this wraps it or doesnt based on the test case + let instruction = if do_in_cpi { + wrap_instruction(spl_instruction_padding::id(), token_instruction, vec![], 0).unwrap() + } else { + token_instruction + }; + + // and here we go + let result = token + .process_ixs(&[instruction], &[¤t_authority]) + .await; + + // truth table for our cases + match (action, enable_cpi_guard, do_in_cpi) { + // all actions succeed with cpi guard off + (_, false, _) => result.unwrap(), + // ownership cannot be transferred with guard + (SetAuthTest::ChangeOwner, true, false) => assert_eq!( + result.unwrap_err(), + client_error(TokenError::CpiGuardOwnerChangeBlocked) + ), + // all other actions succeed outside cpi with guard + (_, true, false) => result.unwrap(), + // removing a close authority succeeds in cpi with guard + (SetAuthTest::RemoveCloseAuth, true, true) => result.unwrap(), + // changing owner, adding close, or changing close all fail in cpi with guard + (_, true, true) => assert_eq!( + result.unwrap_err(), + client_error(TokenError::CpiGuardSetAuthorityBlocked) + ), + } + } +} diff --git a/token/program-2022/src/error.rs b/token/program-2022/src/error.rs index 41570ef4dc4..041b50eb54e 100644 --- a/token/program-2022/src/error.rs +++ b/token/program-2022/src/error.rs @@ -139,7 +139,6 @@ pub enum TokenError { /// mint and try again #[error("An account can only be closed if its withheld fee balance is zero, harvest fees to the mint and try again")] AccountHasWithheldTransferFees, - /// No memo in previous instruction; required for recipient to receive a transfer #[error("No memo in previous instruction; required for recipient to receive a transfer")] NoMemo, @@ -156,9 +155,38 @@ pub enum TokenError { the associated `maximum_pending_balance_credit_counter`" )] MaximumPendingBalanceCreditCounterExceeded, + + // 40 /// The deposit amount for the confidential extension exceeds the maximum limit #[error("Deposit amount exceeds maximum limit")] MaximumDepositAmountExceeded, + /// CPI Guard cannot be enabled or disabled in CPI + #[error("CPI Guard cannot be enabled or disabled in CPI")] + CpiGuardSettingsLocked, + /// CPI Guard is enabled, and a program attempted to transfer user funds without using a delegate + #[error("CPI Guard is enabled, and a program attempted to transfer user funds via CPI without using a delegate")] + CpiGuardTransferBlocked, + /// CPI Guard is enabled, and a program attempted to burn user funds without using a delegate + #[error( + "CPI Guard is enabled, and a program attempted to burn user funds via CPI without using a delegate" + )] + CpiGuardBurnBlocked, + /// CPI Guard is enabled, and a program attempted to close an account without returning lamports to owner + #[error("CPI Guard is enabled, and a program attempted to close an account via CPI without returning lamports to owner")] + CpiGuardCloseAccountBlocked, + + // 45 + /// CPI Guard is enabled, and a program attempted to approve a delegate + #[error("CPI Guard is enabled, and a program attempted to approve a delegate via CPI")] + CpiGuardApproveBlocked, + /// CPI Guard is enabled, and a program attempted to add or replace an authority + #[error( + "CPI Guard is enabled, and a program attempted to add or replace an authority via CPI" + )] + CpiGuardSetAuthorityBlocked, + /// Account ownership cannot be changed while CPI Guard is enabled + #[error("Account ownership cannot be changed while CPI Guard is enabled")] + CpiGuardOwnerChangeBlocked, } impl From for ProgramError { fn from(e: TokenError) -> Self { @@ -274,6 +302,27 @@ impl PrintProgramError for TokenError { TokenError::MaximumDepositAmountExceeded => { msg!("Deposit amount exceeds maximum limit") } + TokenError::CpiGuardSettingsLocked => { + msg!("CPI Guard status cannot be changed in CPI") + } + TokenError::CpiGuardTransferBlocked => { + msg!("CPI Guard is enabled, and a program attempted to transfer user funds without using a delegate") + } + TokenError::CpiGuardBurnBlocked => { + msg!("CPI Guard is enabled, and a program attempted to burn user funds without using a delegate") + } + TokenError::CpiGuardCloseAccountBlocked => { + msg!("CPI Guard is enabled, and a program attempted to close an account without returning lamports to owner") + } + TokenError::CpiGuardApproveBlocked => { + msg!("CPI Guard is enabled, and a program attempted to approve a delegate") + } + TokenError::CpiGuardSetAuthorityBlocked => { + msg!("CPI Guard is enabled, and a program attempted to add or change an authority") + } + TokenError::CpiGuardOwnerChangeBlocked => { + msg!("Account ownership cannot be changed while CPI Guard is enabled") + } } } } diff --git a/token/program-2022/src/extension/cpi_guard/instruction.rs b/token/program-2022/src/extension/cpi_guard/instruction.rs new file mode 100644 index 00000000000..9570315eb9f --- /dev/null +++ b/token/program-2022/src/extension/cpi_guard/instruction.rs @@ -0,0 +1,100 @@ +use { + crate::{ + check_program_account, + instruction::{encode_instruction, TokenInstruction}, + }, + num_enum::{IntoPrimitive, TryFromPrimitive}, + solana_program::{ + instruction::{AccountMeta, Instruction}, + program_error::ProgramError, + pubkey::Pubkey, + }, +}; + +/// CPI Guard extension instructions +#[derive(Clone, Copy, Debug, PartialEq, IntoPrimitive, TryFromPrimitive)] +#[repr(u8)] +pub enum CpiGuardInstruction { + /// Lock certain token operations from taking place within CPI for this Account, namely: + /// * Transfer and Burn must go through a delegate. + /// * CloseAccount can only return lamports to owner. + /// * SetAuthority can only be used to remove an existing close authority. + /// * Approve is disallowed entirely. + /// + /// In addition, CPI Guard cannot be enabled or disabled via CPI. + /// + /// Accounts expected by this instruction: + /// + /// 0. `[writable]` The account to update. + /// 1. `[signer]` The account's owner. + /// + /// * Multisignature authority + /// 0. `[writable]` The account to update. + /// 1. `[]` The account's multisignature owner. + /// 2. ..2+M `[signer]` M signer accounts. + /// + Enable, + /// Allow all token operations to happen via CPI as normal. + /// + /// Implicitly initializes the extension in the case where it is not present. + /// + /// Accounts expected by this instruction: + /// + /// 0. `[writable]` The account to update. + /// 1. `[signer]` The account's owner. + /// + /// * Multisignature authority + /// 0. `[writable]` The account to update. + /// 1. `[]` The account's multisignature owner. + /// 2. ..2+M `[signer]` M signer accounts. + /// + Disable, +} + +/// Create an `Enable` instruction +pub fn enable_cpi_guard( + token_program_id: &Pubkey, + account: &Pubkey, + owner: &Pubkey, + signers: &[&Pubkey], +) -> Result { + check_program_account(token_program_id)?; + let mut accounts = vec![ + AccountMeta::new(*account, false), + AccountMeta::new_readonly(*owner, signers.is_empty()), + ]; + for signer_pubkey in signers.iter() { + accounts.push(AccountMeta::new_readonly(**signer_pubkey, true)); + } + Ok(encode_instruction( + token_program_id, + accounts, + TokenInstruction::CpiGuardExtension, + CpiGuardInstruction::Enable, + &(), + )) +} + +/// Create a `Disable` instruction +pub fn disable_cpi_guard( + token_program_id: &Pubkey, + account: &Pubkey, + owner: &Pubkey, + signers: &[&Pubkey], +) -> Result { + check_program_account(token_program_id)?; + let mut accounts = vec![ + AccountMeta::new(*account, false), + AccountMeta::new_readonly(*owner, signers.is_empty()), + ]; + for signer_pubkey in signers.iter() { + accounts.push(AccountMeta::new_readonly(**signer_pubkey, true)); + } + Ok(encode_instruction( + token_program_id, + accounts, + TokenInstruction::CpiGuardExtension, + CpiGuardInstruction::Disable, + &(), + )) +} diff --git a/token/program-2022/src/extension/cpi_guard/mod.rs b/token/program-2022/src/extension/cpi_guard/mod.rs new file mode 100644 index 00000000000..46c6981f039 --- /dev/null +++ b/token/program-2022/src/extension/cpi_guard/mod.rs @@ -0,0 +1,39 @@ +use { + crate::{ + extension::{Extension, ExtensionType, StateWithExtensionsMut}, + pod::PodBool, + state::Account, + }, + bytemuck::{Pod, Zeroable}, + solana_program::instruction::{get_stack_height, TRANSACTION_LEVEL_STACK_HEIGHT}, +}; + +/// CPI Guard extension instructions +pub mod instruction; + +/// CPI Guard extension processor +pub mod processor; + +/// CPI Guard extension for Accounts +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)] +pub struct CpiGuard { + /// Lock privileged token operations from happening via CPI + pub lock_cpi: PodBool, +} +impl Extension for CpiGuard { + const TYPE: ExtensionType = ExtensionType::CpiGuard; +} + +/// Determine if CPI Guard is enabled for this account +pub fn cpi_guard_enabled(account_state: &StateWithExtensionsMut) -> bool { + if let Ok(extension) = account_state.get_extension::() { + return extension.lock_cpi.into(); + } + false +} + +/// Determine if we are in CPI +pub fn in_cpi() -> bool { + get_stack_height() > TRANSACTION_LEVEL_STACK_HEIGHT +} diff --git a/token/program-2022/src/extension/cpi_guard/processor.rs b/token/program-2022/src/extension/cpi_guard/processor.rs new file mode 100644 index 00000000000..0e3ae0d9da3 --- /dev/null +++ b/token/program-2022/src/extension/cpi_guard/processor.rs @@ -0,0 +1,98 @@ +use { + crate::{ + check_program_account, + error::TokenError, + extension::{ + cpi_guard::{in_cpi, instruction::CpiGuardInstruction, CpiGuard}, + StateWithExtensionsMut, + }, + instruction::decode_instruction_type, + processor::Processor, + state::Account, + }, + solana_program::{ + account_info::{next_account_info, AccountInfo}, + entrypoint::ProgramResult, + msg, + pubkey::Pubkey, + }, +}; + +fn process_enable_cpi_guard(program_id: &Pubkey, accounts: &[AccountInfo]) -> ProgramResult { + let account_info_iter = &mut accounts.iter(); + let token_account_info = next_account_info(account_info_iter)?; + let owner_info = next_account_info(account_info_iter)?; + let owner_info_data_len = owner_info.data_len(); + + let mut account_data = token_account_info.data.borrow_mut(); + let mut account = StateWithExtensionsMut::::unpack(&mut account_data)?; + + Processor::validate_owner( + program_id, + &account.base.owner, + owner_info, + owner_info_data_len, + account_info_iter.as_slice(), + )?; + + if in_cpi() { + return Err(TokenError::CpiGuardSettingsLocked.into()); + } + + let extension = if let Ok(extension) = account.get_extension_mut::() { + extension + } else { + account.init_extension::(true)? + }; + extension.lock_cpi = true.into(); + Ok(()) +} + +fn process_diasble_cpi_guard(program_id: &Pubkey, accounts: &[AccountInfo]) -> ProgramResult { + let account_info_iter = &mut accounts.iter(); + let token_account_info = next_account_info(account_info_iter)?; + let owner_info = next_account_info(account_info_iter)?; + let owner_info_data_len = owner_info.data_len(); + + let mut account_data = token_account_info.data.borrow_mut(); + let mut account = StateWithExtensionsMut::::unpack(&mut account_data)?; + + Processor::validate_owner( + program_id, + &account.base.owner, + owner_info, + owner_info_data_len, + account_info_iter.as_slice(), + )?; + + if in_cpi() { + return Err(TokenError::CpiGuardSettingsLocked.into()); + } + + let extension = if let Ok(extension) = account.get_extension_mut::() { + extension + } else { + account.init_extension::(true)? + }; + extension.lock_cpi = false.into(); + Ok(()) +} + +pub(crate) fn process_instruction( + program_id: &Pubkey, + accounts: &[AccountInfo], + input: &[u8], +) -> ProgramResult { + check_program_account(program_id)?; + + match decode_instruction_type(input)? { + CpiGuardInstruction::Enable => { + msg!("CpiGuardInstruction::Enable"); + process_enable_cpi_guard(program_id, accounts) + } + CpiGuardInstruction::Disable => { + msg!("CpiGuardInstruction::Disable"); + process_diasble_cpi_guard(program_id, accounts) + } + } +} diff --git a/token/program-2022/src/extension/memo_transfer/instruction.rs b/token/program-2022/src/extension/memo_transfer/instruction.rs index 64520293a85..f299167d27f 100644 --- a/token/program-2022/src/extension/memo_transfer/instruction.rs +++ b/token/program-2022/src/extension/memo_transfer/instruction.rs @@ -31,7 +31,7 @@ pub enum RequiredMemoTransfersInstruction { Enable, /// Stop requiring memos for transfers into this Account. /// - /// Fails if the account does not have the extension present. + /// Implicitly initializes the extension in the case where it is not present. /// /// Accounts expected by this instruction: /// diff --git a/token/program-2022/src/extension/mod.rs b/token/program-2022/src/extension/mod.rs index 7c895fa7f66..7675dd497a4 100644 --- a/token/program-2022/src/extension/mod.rs +++ b/token/program-2022/src/extension/mod.rs @@ -5,6 +5,7 @@ use { error::TokenError, extension::{ confidential_transfer::{ConfidentialTransferAccount, ConfidentialTransferMint}, + cpi_guard::CpiGuard, default_account_state::DefaultAccountState, immutable_owner::ImmutableOwner, interest_bearing_mint::InterestBearingConfig, @@ -33,6 +34,8 @@ use serde::{Deserialize, Serialize}; /// Confidential Transfer extension pub mod confidential_transfer; +/// CPI Guard extension +pub mod cpi_guard; /// Default Account State extension pub mod default_account_state; /// Immutable Owner extension @@ -629,6 +632,8 @@ pub enum ExtensionType { NonTransferable, /// Tokens accrue interest over time, InterestBearingConfig, + /// Locks privileged token operations from happening via CPI + CpiGuard, /// Padding extension used to make an account exactly Multisig::LEN, used for testing #[cfg(test)] AccountPaddingTest = u16::MAX - 1, @@ -669,6 +674,7 @@ impl ExtensionType { ExtensionType::MemoTransfer => pod_get_packed_len::(), ExtensionType::NonTransferable => pod_get_packed_len::(), ExtensionType::InterestBearingConfig => pod_get_packed_len::(), + ExtensionType::CpiGuard => pod_get_packed_len::(), #[cfg(test)] ExtensionType::AccountPaddingTest => pod_get_packed_len::(), #[cfg(test)] @@ -729,7 +735,8 @@ impl ExtensionType { ExtensionType::ImmutableOwner | ExtensionType::TransferFeeAmount | ExtensionType::ConfidentialTransferAccount - | ExtensionType::MemoTransfer => AccountType::Account, + | ExtensionType::MemoTransfer + | ExtensionType::CpiGuard => AccountType::Account, #[cfg(test)] ExtensionType::AccountPaddingTest => AccountType::Account, #[cfg(test)] diff --git a/token/program-2022/src/instruction.rs b/token/program-2022/src/instruction.rs index 116b16097ec..7374f957afa 100644 --- a/token/program-2022/src/instruction.rs +++ b/token/program-2022/src/instruction.rs @@ -605,6 +605,11 @@ pub enum TokenInstruction<'a> { /// See `extension::interest_bearing_mint::instruction::InterestBearingMintInstruction` for /// further details about the extended instructions that share this instruction prefix InterestBearingMintExtension, + /// The common instruction prefix for CPI Guard account extension instructions. + /// + /// See `extension::cpi_guard::instruction::CpiGuardInstruction` for + /// further details about the extended instructions that share this instruction prefix + CpiGuardExtension, } impl<'a> TokenInstruction<'a> { /// Unpacks a byte buffer into a [TokenInstruction](enum.TokenInstruction.html). @@ -735,6 +740,7 @@ impl<'a> TokenInstruction<'a> { 31 => Self::CreateNativeMint, 32 => Self::InitializeNonTransferableMint, 33 => Self::InterestBearingMintExtension, + 34 => Self::CpiGuardExtension, _ => return Err(TokenError::InvalidInstruction.into()), }) } @@ -887,6 +893,9 @@ impl<'a> TokenInstruction<'a> { &Self::InterestBearingMintExtension => { buf.push(33); } + &Self::CpiGuardExtension => { + buf.push(34); + } }; buf } diff --git a/token/program-2022/src/processor.rs b/token/program-2022/src/processor.rs index fa033c5e3df..f7285cbfa7b 100644 --- a/token/program-2022/src/processor.rs +++ b/token/program-2022/src/processor.rs @@ -6,6 +6,7 @@ use { error::TokenError, extension::{ confidential_transfer::{self, ConfidentialTransferAccount}, + cpi_guard::{self, in_cpi, CpiGuard}, default_account_state::{self, DefaultAccountState}, immutable_owner::ImmutableOwner, interest_bearing_mint::{self, InterestBearingConfig}, @@ -348,13 +349,21 @@ impl Processor { } } } - _ => Self::validate_owner( - program_id, - &source_account.base.owner, - authority_info, - authority_info_data_len, - account_info_iter.as_slice(), - )?, + _ => { + Self::validate_owner( + program_id, + &source_account.base.owner, + authority_info, + authority_info_data_len, + account_info_iter.as_slice(), + )?; + + if let Ok(cpi_guard) = source_account.get_extension::() { + if cpi_guard.lock_cpi.into() && in_cpi() { + return Err(TokenError::CpiGuardTransferBlocked.into()); + } + } + } }; // Revisit this later to see if it's worth adding a check to reduce @@ -476,6 +485,12 @@ impl Processor { account_info_iter.as_slice(), )?; + if let Ok(cpi_guard) = source_account.get_extension::() { + if cpi_guard.lock_cpi.into() && in_cpi() { + return Err(TokenError::CpiGuardApproveBlocked.into()); + } + } + source_account.base.delegate = COption::Some(*delegate_info.key); source_account.base.delegated_amount = amount; source_account.pack_base(); @@ -549,6 +564,14 @@ impl Processor { return Err(TokenError::ImmutableOwner.into()); } + if let Ok(cpi_guard) = account.get_extension::() { + if cpi_guard.lock_cpi.into() && in_cpi() { + return Err(TokenError::CpiGuardSetAuthorityBlocked.into()); + } else if cpi_guard.lock_cpi.into() { + return Err(TokenError::CpiGuardOwnerChangeBlocked.into()); + } + } + if let COption::Some(authority) = new_authority { account.base.owner = authority; } else { @@ -571,6 +594,13 @@ impl Processor { authority_info_data_len, account_info_iter.as_slice(), )?; + + if let Ok(cpi_guard) = account.get_extension::() { + if cpi_guard.lock_cpi.into() && in_cpi() && new_authority != COption::None { + return Err(TokenError::CpiGuardSetAuthorityBlocked.into()); + } + } + account.base.close_authority = new_authority; } _ => { @@ -829,13 +859,21 @@ impl Processor { source_account.base.delegate = COption::None; } } - _ => Self::validate_owner( - program_id, - &source_account.base.owner, - authority_info, - authority_info_data_len, - account_info_iter.as_slice(), - )?, + _ => { + Self::validate_owner( + program_id, + &source_account.base.owner, + authority_info, + authority_info_data_len, + account_info_iter.as_slice(), + )?; + + if let Ok(cpi_guard) = source_account.get_extension::() { + if cpi_guard.lock_cpi.into() && in_cpi() { + return Err(TokenError::CpiGuardBurnBlocked.into()); + } + } + } } } @@ -889,6 +927,15 @@ impl Processor { .base .is_owned_by_system_program_or_incinerator() { + if let Ok(cpi_guard) = source_account.get_extension::() { + if cpi_guard.lock_cpi.into() + && in_cpi() + && !cmp_pubkeys(destination_account_info.key, &source_account.base.owner) + { + return Err(TokenError::CpiGuardCloseAccountBlocked.into()); + } + } + Self::validate_owner( program_id, &authority, @@ -1324,6 +1371,9 @@ impl Processor { &input[1..], ) } + TokenInstruction::CpiGuardExtension => { + cpi_guard::processor::process_instruction(program_id, accounts, &input[1..]) + } } }