Repository: solana-labs/solana-program-library
Pull request #6550 (2024-04-04 22:31:02)

token-cli: Fix offline signing by making simulated compute unit limit explicit
#### Problem

https://github.com/solana-labs/solana-program-library/pull/6501 broke offline signing -- by default, the token client sets the compute unit limit based on simulation. If it's using offline signing, it'll silently swallow the error and pop the compute unit limit instruction.

So when we offline-sign a transaction, it won't have the compute unit limit instruction, and when we broadcast it, we'll simulate and add the compute unit limit instruction, causing a pre-signer error.

This was the classic overly-clever-but-brittle solution. I discovered it while applying the token-cli changes to the Solana CLI.

#### Solution

Be more explicit:

* don't simulate by default
* allow users to specify `--with-compute-unit-limit SIMULATED` if they want that
* error if someone sets `--with-compute-unit-price` *without* `--with-compute-unit-limit`. This one might be contentious, and I could be talked down to making it a warning. But since we've only just introduced the flag, I think it's better to do the right thing and force people to be explicit
* improve the offline signing test by actually sending the transaction so we catch this in the future

Let me know how you like the `ComputeUnitLimit`, because I'll likely do the same exact thing for the Solana CLI. Except there, `--with-compute-unit-price` has been around for some time, so I would show a warning if someone specifies price without a limit.

CriesofCarrots: Sorry for the slow review here! I'm a little confused. If it's expected that a user would use `SIMULATED` with offline signing, it seems like the `--sign-only` command would need to return the results of simulation so the values could be made specific in the subsequent submission command. Otherwise, there's a risk the signature is invalid if those values change between signing and submission (as you note in the help text). This is a bit gross to me, although we could let the user take that risk (but places more support burden on us).

If it's *not* expected that a user would use `SIMULATED` with offline signing, I think we need to prevent it. And then I don't think the keyword really offers a lot of value, and we should default to simulation in the online case, and require offline signing to be explicit from the get-go. (This is an existing approach/pattern in solana-cli, incidentally, eg split-stake: https://github.com/anza-xyz/agave/blob/180a186c7da1b8d6f5e93ec2cbc5bf2ea17db2cb/cli/src/stake.rs#L551-L560)
joncinque: > If it's not expected that a user would use SIMULATED with offline signing, I think we need to prevent it

That's the way I was thinking about it. We could eventually create a way for someone to simulate a transaction just to get the compute units and pass it in explicitly for offline signing, but that would be a separate PR.

> And then I don't think the keyword really offers a lot of value, and we should default to simulation in the online case, and require offline signing to be explicit from the get-go

Just so I'm sure I understand: require `compute_unit_limit` if both `sign_only` and `compute_unit_price` are set, and then any instruction can also use `compute_unit_limit`. That would fix the issue.

The problem I'm running into is when to use the simulated compute units. Currently, the code *always* simulates, but that's an issue when one of the signers is a pre-signer. So we're left with only simulating for compute units if `compute_unit_price` is set and `compute_unit_limit` isn't, which will work in all cases.

However, this means it isn't possible to specify, "please simulate my transaction and use the compute unit limit from the result", you're required to also set a compute unit price. Someone could get around this by specifying a compute unit price of 0, but that seems confusing. So that's why I added the `SIMULATED` keyword.

Do you think that's a worthwhile use case? I think it is, since the scheduler should have an easier time packing transactions with (low) compute unit limits. But maybe I'm wrong.
CriesofCarrots: > Just so I'm sure I understand: require `compute_unit_limit` if both `sign_only` and `compute_unit_price` are set, and then any instruction can also use `compute_unit_limit`. That would fix the issue.

Yep, this is what I meant. I can't quite tell from your comment, what do you think of that approach?
joncinque: > I can't quite tell from your comment, what do you think of that approach?

Sorry, I'm having a hard time articulating this. Your idea fixes the issue, but there's a knock-on effect: it *requires* someone to specify a compute unit price in order to use the simulated compute unit limit.

I wanted to allow people to specify that they want to use simulated compute units without specifying a compute unit price, which is why I added the `SIMULATED` variant. Do you think that's worth it?
CriesofCarrots: > it requires someone to specify a compute unit price in order to use the simulated compute unit limit.

I must be being stupid, because I don't follow why this is required. Let me re-read the code, more closely, and then reply again.
joncinque: I added some more testing and integrated the changes your requested. Unfortunately, clap v2 doesn't expose a way to say "require this arg if these two other args are set", [required_if](https://docs.rs/clap/2.34.0/clap/struct.Arg.html#method.required_if) only checks if an arg is set to a particular value, and [requires_all](https://docs.rs/clap/2.34.0/clap/struct.Arg.html#method.requires_all) is not conditional
joncinque: >  Also, is the CI failure expected? It doesn't look related, but did happen twice.

It was a legit failure, the override limit that was set for `set-interest-rate` was being clobbered

Patch:
diff --git a/token/cli/src/clap_app.rs b/token/cli/src/clap_app.rs
index f82af2eb7..b9817bff0 100644
--- a/token/cli/src/clap_app.rs
+++ b/token/cli/src/clap_app.rs
@@ -17,7 +17,7 @@ use {
     spl_token_2022::instruction::{AuthorityType, MAX_SIGNERS, MIN_SIGNERS},
     std::{fmt, str::FromStr},
     strum::IntoEnumIterator,
-    strum_macros::{EnumIter, EnumString, IntoStaticStr},
+    strum_macros::{AsRefStr, EnumIter, EnumString, IntoStaticStr},
 };
 
 pub type Error = Box<dyn std::error::Error + Send + Sync>;
@@ -78,7 +78,7 @@ pub const COMPUTE_UNIT_LIMIT_ARG: ArgConstant<'static> = ArgConstant {
 
 pub static VALID_TOKEN_PROGRAM_IDS: [Pubkey; 2] = [spl_token_2022::ID, spl_token::ID];
 
-#[derive(Debug, Clone, Copy, PartialEq, EnumString, IntoStaticStr)]
+#[derive(AsRefStr, Debug, Clone, Copy, PartialEq, EnumString, IntoStaticStr)]
 #[strum(serialize_all = "kebab-case")]
 pub enum CommandName {
     CreateToken,
@@ -350,6 +350,7 @@ where
         Err(e) => Err(e),
     }
 }
+
 struct SignOnlyNeedsFullMintSpec {}
 impl offline::ArgsConfig for SignOnlyNeedsFullMintSpec {
     fn sign_only_arg<'a, 'b>(&self, arg: Arg<'a, 'b>) -> Arg<'a, 'b> {
diff --git a/token/cli/src/command.rs b/token/cli/src/command.rs
index 80615c01e..ee675b4e8 100644
--- a/token/cli/src/command.rs
+++ b/token/cli/src/command.rs
@@ -68,7 +68,7 @@ use {
     },
     spl_token_client::{
         client::{ProgramRpcClientSendTransaction, RpcClientResponse},
-        token::{ExtensionInitializationParams, Token},
+        token::{ComputeUnitLimit, ExtensionInitializationParams, Token},
     },
     spl_token_group_interface::state::TokenGroup,
     spl_token_metadata_interface::state::{Field, TokenMetadata},
@@ -152,11 +152,7 @@ fn config_token_client(
     token: Token<ProgramRpcClientSendTransaction>,
     config: &Config<'_>,
 ) -> Result<Token<ProgramRpcClientSendTransaction>, Error> {
-    let token = if let Some(compute_unit_limit) = config.compute_unit_limit {
-        token.with_compute_unit_limit(compute_unit_limit)
-    } else {
-        token
-    };
+    let token = token.with_compute_unit_limit(config.compute_unit_limit.clone());
 
     let token = if let Some(compute_unit_price) = config.compute_unit_price {
         token.with_compute_unit_price(compute_unit_price)
@@ -432,10 +428,13 @@ async fn command_set_interest_rate(
     rate_bps: i16,
     bulk_signers: Vec<Arc<dyn Signer>>,
 ) -> CommandResult {
+    let mut token = token_client_from_config(config, &token_pubkey, None)?;
     // Because set_interest_rate depends on the time, it can cost more between
     // simulation and execution. To help that, just set a static compute limit
-    let token = base_token_client(config, &token_pubkey, None)?.with_compute_unit_limit(2_500);
-    let token = config_token_client(token, config)?;
+    // if none has been set
+    if !matches!(config.compute_unit_limit, ComputeUnitLimit::Static(_)) {
+        token = token.with_compute_unit_limit(ComputeUnitLimit::Static(2_500));
+    }
 
     if !config.sign_only {
         let mint_account = config.get_account_checked(&token_pubkey).await?;
diff --git a/token/cli/src/config.rs b/token/cli/src/config.rs
index 96072fea0..89137a51b 100644
--- a/token/cli/src/config.rs
+++ b/token/cli/src/config.rs
@@ -20,8 +20,11 @@ use {
         extension::StateWithExtensionsOwned,
         state::{Account, Mint},
     },
-    spl_token_client::client::{
-        ProgramClient, ProgramOfflineClient, ProgramRpcClient, ProgramRpcClientSendTransaction,
+    spl_token_client::{
+        client::{
+            ProgramClient, ProgramOfflineClient, ProgramRpcClient, ProgramRpcClientSendTransaction,
+        },
+        token::ComputeUnitLimit,
     },
     std::{process::exit, rc::Rc, sync::Arc},
 };
@@ -68,7 +71,7 @@ pub struct Config<'a> {
     pub program_id: Pubkey,
     pub restrict_to_program_id: bool,
     pub compute_unit_price: Option<u64>,
-    pub compute_unit_limit: Option<u32>,
+    pub compute_unit_limit: ComputeUnitLimit,
 }
 
 impl<'a> Config<'a> {
@@ -280,9 +283,32 @@ impl<'a> Config<'a> {
                 (default_program_id, false)
             };
 
+        // need to specify a compute limit if compute price and blockhash are specified
+        if matches.is_present(BLOCKHASH_ARG.name)
+            && matches.is_present(COMPUTE_UNIT_PRICE_ARG.name)
+            && !matches.is_present(COMPUTE_UNIT_LIMIT_ARG.name)
+        {
+            clap::Error::with_description(
+                &format!(
+                    "Need to set `{}` if `{}` and `--{}` are set",
+                    COMPUTE_UNIT_LIMIT_ARG.long, COMPUTE_UNIT_PRICE_ARG.long, BLOCKHASH_ARG.long,
+                ),
+                clap::ErrorKind::MissingRequiredArgument,
+            )
+            .exit();
+        }
+
         let nonce_blockhash = value_of(matches, BLOCKHASH_ARG.name);
         let compute_unit_price = value_of(matches, COMPUTE_UNIT_PRICE_ARG.name);
-        let compute_unit_limit = value_of(matches, COMPUTE_UNIT_LIMIT_ARG.name);
+        let compute_unit_limit = value_of(matches, COMPUTE_UNIT_LIMIT_ARG.name)
+            .map(ComputeUnitLimit::Static)
+            .unwrap_or_else(|| {
+                if nonce_blockhash.is_some() {
+                    ComputeUnitLimit::Default
+                } else {
+                    ComputeUnitLimit::Simulated
+                }
+            });
         Self {
             default_signer,
             rpc_client,
diff --git a/token/client/src/token.rs b/token/client/src/token.rs
index 59c90be5d..a9fefc21f 100644
--- a/token/client/src/token.rs
+++ b/token/client/src/token.rs
@@ -332,6 +332,13 @@ impl TokenMemo {
     }
 }
 
+#[derive(Debug, Clone)]
+pub enum ComputeUnitLimit {
+    Default,
+    Simulated,
+    Static(u32),
+}
+
 pub struct Token<T> {
     client: Arc<dyn ProgramClient<T>>,
     pubkey: Pubkey, /* token mint */
@@ -344,7 +351,7 @@ pub struct Token<T> {
     memo: Arc<RwLock<Option<TokenMemo>>>,
     transfer_hook_accounts: Option<Vec<AccountMeta>>,
     compute_unit_price: Option<u64>,
-    compute_unit_limit: Option<u32>,
+    compute_unit_limit: ComputeUnitLimit,
 }
 
 impl<T> fmt::Debug for Token<T> {
@@ -411,7 +418,7 @@ where
             memo: Arc::new(RwLock::new(None)),
             transfer_hook_accounts: None,
             compute_unit_price: None,
-            compute_unit_limit: None,
+            compute_unit_limit: ComputeUnitLimit::Default,
         }
     }
 
@@ -466,8 +473,8 @@ where
         self
     }
 
-    pub fn with_compute_unit_limit(mut self, compute_unit_limit: u32) -> Self {
-        self.compute_unit_limit = Some(compute_unit_limit);
+    pub fn with_compute_unit_limit(mut self, compute_unit_limit: ComputeUnitLimit) -> Self {
+        self.compute_unit_limit = compute_unit_limit;
         self
     }
 
@@ -548,19 +555,16 @@ where
             .simulate_transaction(&transaction)
             .await
             .map_err(TokenError::Client)?;
-        if let Ok(units_consumed) = simulation_result.get_compute_units_consumed() {
-            // Overwrite the compute unit limit instruction with the actual units consumed
-            let compute_unit_limit =
-                u32::try_from(units_consumed).map_err(|x| TokenError::Client(x.into()))?;
-            instructions
-                .last_mut()
-                .expect("Compute budget instruction was added earlier")
-                .data = ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
-        } else {
-            // `get_compute_units_consumed()` fails for offline signing, so we
-            // catch that error and remove the instruction that was added
-            instructions.pop();
-        }
+        let units_consumed = simulation_result
+            .get_compute_units_consumed()
+            .map_err(TokenError::Client)?;
+        // Overwrite the compute unit limit instruction with the actual units consumed
+        let compute_unit_limit =
+            u32::try_from(units_consumed).map_err(|x| TokenError::Client(x.into()))?;
+        instructions
+            .last_mut()
+            .expect("Compute budget instruction was added earlier")
+            .data = ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
         Ok(())
     }
 
@@ -619,13 +623,17 @@ where
         // all instructions have been added to the transaction, so be sure to
         // keep this instruction as the last one before creating and sending the
         // transaction.
-        if let Some(compute_unit_limit) = self.compute_unit_limit {
-            instructions.push(ComputeBudgetInstruction::set_compute_unit_limit(
-                compute_unit_limit,
-            ));
-        } else {
-            self.add_compute_unit_limit_from_simulation(&mut instructions, &blockhash)
-                .await?;
+        match self.compute_unit_limit {
+            ComputeUnitLimit::Default => {}
+            ComputeUnitLimit::Simulated => {
+                self.add_compute_unit_limit_from_simulation(&mut instructions, &blockhash)
+                    .await?;
+            }
+            ComputeUnitLimit::Static(compute_unit_limit) => {
+                instructions.push(ComputeBudgetInstruction::set_compute_unit_limit(
+                    compute_unit_limit,
+                ));
+            }
         }
 
         let message = Message::new_with_blockhash(&instructions, fee_payer, &blockhash);


Test patch:
diff --git a/token/cli/tests/command.rs b/token/cli/tests/command.rs
index 1e9622354..f89d5b3d6 100644
--- a/token/cli/tests/command.rs
+++ b/token/cli/tests/command.rs
@@ -44,11 +44,16 @@ use {
         client::{
             ProgramClient, ProgramOfflineClient, ProgramRpcClient, ProgramRpcClientSendTransaction,
         },
-        token::Token,
+        token::{ComputeUnitLimit, Token},
     },
     spl_token_group_interface::state::{TokenGroup, TokenGroupMember},
     spl_token_metadata_interface::state::TokenMetadata,
-    std::{ffi::OsString, path::PathBuf, str::FromStr, sync::Arc},
+    std::{
+        ffi::{OsStr, OsString},
+        path::PathBuf,
+        str::FromStr,
+        sync::Arc,
+    },
     tempfile::NamedTempFile,
 };
 
@@ -201,7 +206,7 @@ fn test_config_with_default_signer<'a>(
         program_id: *program_id,
         restrict_to_program_id: true,
         compute_unit_price: None,
-        compute_unit_limit: None,
+        compute_unit_limit: ComputeUnitLimit::Simulated,
     }
 }
 
@@ -230,7 +235,7 @@ fn test_config_without_default_signer<'a>(
         program_id: *program_id,
         restrict_to_program_id: true,
         compute_unit_price: None,
-        compute_unit_limit: None,
+        compute_unit_limit: ComputeUnitLimit::Simulated,
     }
 }
 
@@ -441,7 +446,7 @@ where
     process_command(&sub_command, matches, config, wallet_manager, bulk_signers).await
 }
 
-async fn exec_test_cmd(config: &Config<'_>, args: &[&str]) -> CommandResult {
+async fn exec_test_cmd<T: AsRef<OsStr>>(config: &Config<'_>, args: &[T]) -> CommandResult {
     let default_decimals = format!("{}", spl_token_2022::native_mint::DECIMALS);
     let minimum_signers_help = minimum_signers_help_string();
     let multisig_member_help = multisig_member_help_string();
@@ -2974,10 +2979,17 @@ async fn multisig_transfer(test_validator: &TestValidator, payer: &Keypair) {
     }
 }
 
-async fn offline_multisig_transfer_with_nonce(test_validator: &TestValidator, payer: &Keypair) {
+async fn do_offline_multisig_transfer(
+    test_validator: &TestValidator,
+    payer: &Keypair,
+    compute_unit_price: Option<u64>,
+) {
     let m = 2;
     let n = 3u8;
 
+    let fee_payer_keypair_file = NamedTempFile::new().unwrap();
+    write_keypair_file(payer, &fee_payer_keypair_file).unwrap();
+
     let (multisig_members, multisig_paths): (Vec<_>, Vec<_>) = std::iter::repeat_with(Keypair::new)
         .take(n as usize)
         .map(|s| {
@@ -2988,6 +3000,7 @@ async fn offline_multisig_transfer_with_nonce(test_validator: &TestValidator, pa
         .unzip();
     for program_id in VALID_TOKEN_PROGRAM_IDS.iter() {
         let mut config = test_config_with_default_signer(test_validator, payer, program_id);
+        config.compute_unit_limit = ComputeUnitLimit::Default;
         let token = create_token(&config, payer).await;
         let nonce = create_nonce(&config, payer).await;
 
@@ -3032,58 +3045,121 @@ async fn offline_multisig_transfer_with_nonce(test_validator: &TestValidator, pa
         let program_client: Arc<dyn ProgramClient<ProgramRpcClientSendTransaction>> = Arc::new(
             ProgramOfflineClient::new(blockhash, ProgramRpcClientSendTransaction),
         );
+        let mut args = vec![
+            "spl-token".to_string(),
+            CommandName::Transfer.as_ref().to_string(),
+            token.to_string(),
+            "10".to_string(),
+            destination.to_string(),
+            "--blockhash".to_string(),
+            blockhash.to_string(),
+            "--nonce".to_string(),
+            nonce.to_string(),
+            "--nonce-authority".to_string(),
+            payer.pubkey().to_string(),
+            "--sign-only".to_string(),
+            "--mint-decimals".to_string(),
+            format!("{}", TEST_DECIMALS),
+            "--multisig-signer".to_string(),
+            multisig_paths[1].path().to_str().unwrap().to_string(),
+            "--multisig-signer".to_string(),
+            multisig_members[2].to_string(),
+            "--from".to_string(),
+            source.to_string(),
+            "--owner".to_string(),
+            multisig_pubkey.to_string(),
+            "--fee-payer".to_string(),
+            payer.pubkey().to_string(),
+            "--program-id".to_string(),
+            program_id.to_string(),
+        ];
+        if let Some(compute_unit_price) = compute_unit_price {
+            args.push("--with-compute-unit-price".to_string());
+            args.push(compute_unit_price.to_string());
+            args.push("--with-compute-unit-limit".to_string());
+            args.push(10_000.to_string());
+        }
         config.program_client = program_client;
-        let result = exec_test_cmd(
-            &config,
-            &[
-                "spl-token",
-                CommandName::Transfer.into(),
-                &token.to_string(),
-                "10",
-                &destination.to_string(),
-                "--blockhash",
-                &blockhash.to_string(),
-                "--nonce",
-                &nonce.to_string(),
-                "--nonce-authority",
-                &payer.pubkey().to_string(),
-                "--sign-only",
-                "--mint-decimals",
-                &format!("{}", TEST_DECIMALS),
-                "--multisig-signer",
-                multisig_paths[1].path().to_str().unwrap(),
-                "--multisig-signer",
-                &multisig_members[2].to_string(),
-                "--from",
-                &source.to_string(),
-                "--owner",
-                &multisig_pubkey.to_string(),
-                "--fee-payer",
-                &multisig_members[0].to_string(),
-            ],
-        )
-        .await
-        .unwrap();
+        let result = exec_test_cmd(&config, &args).await.unwrap();
         // the provided signer has a signature, denoted by the pubkey followed
         // by "=" and the signature
-        assert!(result.contains(&format!("{}=", multisig_members[1])));
+        let member_prefix = format!("{}=", multisig_members[1]);
+        let signature_position = result.find(&member_prefix).unwrap();
+        let end_position = result[signature_position..].find('\n').unwrap();
+        let signer = result[signature_position..].get(..end_position).unwrap();
 
         // other three expected signers are absent
         let absent_signers_position = result.find("Absent Signers").unwrap();
         let absent_signers = result.get(absent_signers_position..).unwrap();
-        assert!(absent_signers.contains(&multisig_members[0].to_string()));
         assert!(absent_signers.contains(&multisig_members[2].to_string()));
         assert!(absent_signers.contains(&payer.pubkey().to_string()));
 
         // and nothing else is marked a signer
+        assert!(!absent_signers.contains(&multisig_members[0].to_string()));
         assert!(!absent_signers.contains(&multisig_pubkey.to_string()));
         assert!(!absent_signers.contains(&nonce.to_string()));
         assert!(!absent_signers.contains(&source.to_string()));
         assert!(!absent_signers.contains(&destination.to_string()));
         assert!(!absent_signers.contains(&token.to_string()));
+
+        // now send the transaction
+        let program_client: Arc<dyn ProgramClient<ProgramRpcClientSendTransaction>> = Arc::new(
+            ProgramRpcClient::new(config.rpc_client.clone(), ProgramRpcClientSendTransaction),
+        );
+        config.program_client = program_client;
+        let mut args = vec![
+            "spl-token".to_string(),
+            CommandName::Transfer.as_ref().to_string(),
+            token.to_string(),
+            "10".to_string(),
+            destination.to_string(),
+            "--blockhash".to_string(),
+            blockhash.to_string(),
+            "--nonce".to_string(),
+            nonce.to_string(),
+            "--nonce-authority".to_string(),
+            fee_payer_keypair_file.path().to_str().unwrap().to_string(),
+            "--mint-decimals".to_string(),
+            format!("{}", TEST_DECIMALS),
+            "--multisig-signer".to_string(),
+            multisig_members[1].to_string(),
+            "--multisig-signer".to_string(),
+            multisig_paths[2].path().to_str().unwrap().to_string(),
+            "--from".to_string(),
+            source.to_string(),
+            "--owner".to_string(),
+            multisig_pubkey.to_string(),
+            "--fee-payer".to_string(),
+            fee_payer_keypair_file.path().to_str().unwrap().to_string(),
+            "--program-id".to_string(),
+            program_id.to_string(),
+            "--signer".to_string(),
+            signer.to_string(),
+        ];
+        if let Some(compute_unit_price) = compute_unit_price {
+            args.push("--with-compute-unit-price".to_string());
+            args.push(compute_unit_price.to_string());
+            args.push("--with-compute-unit-limit".to_string());
+            args.push(10_000.to_string());
+        }
+        exec_test_cmd(&config, &args).await.unwrap();
+
+        let account = config.rpc_client.get_account(&source).await.unwrap();
+        let token_account = StateWithExtensionsOwned::<Account>::unpack(account.data).unwrap();
+        let amount = spl_token::ui_amount_to_amount(90.0, TEST_DECIMALS);
+        assert_eq!(token_account.base.amount, amount);
+        let account = config.rpc_client.get_account(&destination).await.unwrap();
+        let token_account = StateWithExtensionsOwned::<Account>::unpack(account.data).unwrap();
+        let amount = spl_token::ui_amount_to_amount(10.0, TEST_DECIMALS);
+        assert_eq!(token_account.base.amount, amount);
     }
 }
 
+async fn offline_multisig_transfer_with_nonce(test_validator: &TestValidator, payer: &Keypair) {
+    do_offline_multisig_transfer(test_validator, payer, None).await;
+    do_offline_multisig_transfer(test_validator, payer, Some(10)).await;
+}
+
 async fn withdraw_excess_lamports_from_multisig(test_validator: &TestValidator, payer: &Keypair) {
     let m = 3;
     let n = 5u8;
@@ -4024,7 +4100,7 @@ async fn compute_budget(test_validator: &TestValidator, payer: &Keypair) {
     for program_id in VALID_TOKEN_PROGRAM_IDS.iter() {
         let mut config = test_config_with_default_signer(test_validator, payer, program_id);
         config.compute_unit_price = Some(42);
-        config.compute_unit_limit = Some(30_000);
+        config.compute_unit_limit = ComputeUnitLimit::Static(30_000);
         run_transfer_test(&config, payer).await;
     }
 }
diff --git a/token/program-2022-test/tests/confidential_transfer.rs b/token/program-2022-test/tests/confidential_transfer.rs
index de1520103..73a1d77f0 100644
--- a/token/program-2022-test/tests/confidential_transfer.rs
+++ b/token/program-2022-test/tests/confidential_transfer.rs
@@ -39,7 +39,7 @@ use {
     },
     spl_token_client::{
         proof_generation::transfer_with_fee_split_proof_data,
-        token::{ExtensionInitializationParams, TokenError as TokenClientError},
+        token::{ComputeUnitLimit, ExtensionInitializationParams, TokenError as TokenClientError},
     },
     std::{convert::TryInto, mem::size_of},
 };
@@ -2526,7 +2526,7 @@ async fn confidential_transfer_transfer_with_split_proof_contexts_in_parallel()
     // With split proofs in parallel, one of the transactions does more work
     // than the other, which isn't caught during the simulation to discover the
     // compute unit limit.
-    let token = token.with_compute_unit_limit(500_000);
+    let token = token.with_compute_unit_limit(ComputeUnitLimit::Static(500_000));
     token
         .confidential_transfer_transfer_with_split_proofs_in_parallel(
             &alice_meta.token_account,
@@ -2948,7 +2948,7 @@ async fn confidential_transfer_transfer_with_fee_and_split_proof_context_in_para
     // With split proofs in parallel, one of the transactions does more work
     // than the other, which isn't caught during the simulation to discover the
     // compute unit limit.
-    let token = token.with_compute_unit_limit(500_000);
+    let token = token.with_compute_unit_limit(ComputeUnitLimit::Static(500_000));
     token
         .confidential_transfer_transfer_with_fee_and_split_proofs_in_parallel(
             &alice_meta.token_account,
diff --git a/token/program-2022-test/tests/program_test.rs b/token/program-2022-test/tests/program_test.rs
index ef106aedf..0d4979383 100644
--- a/token/program-2022-test/tests/program_test.rs
+++ b/token/program-2022-test/tests/program_test.rs
@@ -20,7 +20,7 @@ use {
             ProgramBanksClient, ProgramBanksClientProcessTransaction, ProgramClient,
             SendTransaction, SimulateTransaction,
         },
-        token::{ExtensionInitializationParams, Token, TokenResult},
+        token::{ComputeUnitLimit, ExtensionInitializationParams, Token, TokenResult},
     },
     std::sync::Arc,
 };
@@ -113,7 +113,8 @@ impl TestContext {
             &mint_account.pubkey(),
             Some(decimals),
             Arc::new(keypair_clone(&payer)),
-        );
+        )
+        .with_compute_unit_limit(ComputeUnitLimit::Simulated);
 
         let token_unchecked = Token::new(
             Arc::clone(&client),
@@ -121,7 +122,8 @@ impl TestContext {
             &mint_account.pubkey(),
             None,
             Arc::new(payer),
-        );
+        )
+        .with_compute_unit_limit(ComputeUnitLimit::Simulated);
 
         token
             .create_mint(
@@ -157,7 +159,8 @@ impl TestContext {
             Token::create_native_mint(Arc::clone(&client), &id(), Arc::new(keypair_clone(&payer)))
                 .await?;
         // unchecked native is never needed because decimals is known statically
-        let token_unchecked = Token::new_native(Arc::clone(&client), &id(), Arc::new(payer));
+        let token_unchecked = Token::new_native(Arc::clone(&client), &id(), Arc::new(payer))
+            .with_compute_unit_limit(ComputeUnitLimit::Simulated);
         self.token_context = Some(TokenContext {
             decimals: native_mint::DECIMALS,
             mint_authority: Keypair::new(), /* bogus */