From b9aa6956b6d25fd7ab605013ec160ae9fcb9f250 Mon Sep 17 00:00:00 2001
From: Jon C <me@jonc.dev>
Date: Mon, 25 Nov 2024 12:38:07 +0100
Subject: [PATCH] token-client: Check simulation results for error

#### Problem

If the simulation for compute unit limit fails in the token-cli, it may
continue to try to broadcast a transaction since the error falls
through. Although that transaction should fail too, at least we don't
run the risk of broadcasting something known to be bad.

#### Summary of changes

Add an `err()` function on the simulation results, and abort if the
simulation for compute units fails.
---
 token/client/src/client.rs | 13 +++++++++++++
 token/client/src/token.rs  |  3 +++
 2 files changed, 16 insertions(+)

diff --git a/token/client/src/client.rs b/token/client/src/client.rs
index 12f3faf17a4..0ead22863e6 100644
--- a/token/client/src/client.rs
+++ b/token/client/src/client.rs
@@ -26,6 +26,7 @@ pub trait SimulateTransaction {
 /// Trait for the output of a simulation
 pub trait SimulationResult {
     fn get_compute_units_consumed(&self) -> ProgramClientResult<u64>;
+    fn err(self) -> Option<ProgramClientError>;
 }
 
 /// Extends basic `SendTransaction` trait with function `send` where client is
@@ -78,6 +79,9 @@ impl SimulationResult for BanksTransactionResultWithSimulation {
             .map(|x| x.units_consumed)
             .ok_or("No simulation results found".into())
     }
+    fn err(self) -> Option<ProgramClientError> {
+        self.result.and_then(|r| r.err().map(|e| e.into()))
+    }
 }
 
 impl SimulateTransaction for ProgramBanksClientProcessTransaction {
@@ -165,6 +169,15 @@ impl SimulationResult for RpcClientResponse {
                 .ok_or("No simulation results found".into()),
         }
     }
+    fn err(self) -> Option<ProgramClientError> {
+        match self {
+            // `Transaction` is the result of an offline simulation. The error
+            // should be properly handled by a caller that supports offline
+            // signing
+            Self::Signature(_) | Self::Transaction(_) => Some("Not a simulation result".into()),
+            Self::Simulation(simulation_result) => simulation_result.err.map(|e| e.into()),
+        }
+    }
 }
 
 impl SimulateTransaction for ProgramRpcClientSendTransaction {
diff --git a/token/client/src/token.rs b/token/client/src/token.rs
index 20cae87c1fe..7e41ea447c9 100644
--- a/token/client/src/token.rs
+++ b/token/client/src/token.rs
@@ -588,6 +588,9 @@ where
         let units_consumed = simulation_result
             .get_compute_units_consumed()
             .map_err(TokenError::Client)?;
+        if let Some(err) = simulation_result.err() {
+            return Err(TokenError::Client(err));
+        }
         // 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()))?;