From dbb5eeb704474b4f41a078ba8a5aca049af65e66 Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Tue, 16 Jul 2024 15:43:21 -0600
Subject: [PATCH] Fix a potential crash related to varying behavior between
 change strategies.

It was possible for `GreedyInputSelector` to crash if the change
strategy being used for input selection were to place a ZIP 320
ephemeral output anywhere but as the last element in the returned change
values. This has been replaced by an error; the error will also be
returned if the change strategy returns more than one ephemeral output
in the change values.
---
 .../src/data_api/wallet/input_selection.rs    | 43 +++++++++++++------
 zcash_client_backend/src/proposal.rs          |  9 ++++
 2 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/zcash_client_backend/src/data_api/wallet/input_selection.rs b/zcash_client_backend/src/data_api/wallet/input_selection.rs
index a3518795c..7ffc87382 100644
--- a/zcash_client_backend/src/data_api/wallet/input_selection.rs
+++ b/zcash_client_backend/src/data_api/wallet/input_selection.rs
@@ -33,7 +33,7 @@ use crate::{
 #[cfg(feature = "transparent-inputs")]
 use {
     crate::{
-        fees::{ChangeValue, EphemeralBalance},
+        fees::EphemeralBalance,
         proposal::{Step, StepOutput, StepOutputIndex},
         zip321::Payment,
     },
@@ -328,6 +328,11 @@ pub struct GreedyInputSelector<DbT, ChangeT> {
 impl<DbT, ChangeT: ChangeStrategy> GreedyInputSelector<DbT, ChangeT> {
     /// Constructs a new greedy input selector that uses the provided change strategy to determine
     /// change values and fee amounts.
+    ///
+    /// The [`ChangeStrategy`] provided must produce exactly one ephemeral change value when
+    /// computing a transaction balance if an [`EphemeralBalance::Output`] value is provided for
+    /// its ephemeral balance, or the resulting [`GreedyInputSelector`] will return an error when
+    /// attempting to construct a transaction proposal that requires such an output.
     pub fn new(change_strategy: ChangeT, dust_output_policy: DustOutputPolicy) -> Self {
         GreedyInputSelector {
             change_strategy,
@@ -605,19 +610,29 @@ where
                         //   a single additional ephemeral output to the transparent pool.
                         // * `tr1` spends from that ephemeral output to each TEX output.
 
-                        // The ephemeral output should always be at the last change index.
-                        assert_eq!(
-                            *balance.proposed_change().last().expect("nonempty"),
-                            ChangeValue::ephemeral_transparent(
-                                ephemeral_balance
-                                    .and_then(|b| b.ephemeral_output_amount())
-                                    .expect("ephemeral output is present")
-                            )
-                        );
-                        let ephemeral_stepoutput = StepOutput::new(
-                            0,
-                            StepOutputIndex::Change(balance.proposed_change().len() - 1),
-                        );
+                        // Find exactly one ephemeral change output.
+                        let ephemeral_outputs = balance
+                            .proposed_change()
+                            .iter()
+                            .enumerate()
+                            .filter(|(_, c)| c.is_ephemeral())
+                            .collect::<Vec<_>>();
+
+                        let ephemeral_value = ephemeral_balance
+                            .and_then(|b| b.ephemeral_output_amount())
+                            .expect("ephemeral output balance exists (constructed above)");
+
+                        let ephemeral_output_index = match &ephemeral_outputs[..] {
+                            [(i, change_value)] if change_value.value() == ephemeral_value => {
+                                Ok(*i)
+                            }
+                            _ => Err(InputSelectorError::Proposal(
+                                ProposalError::EphemeralOutputsInvalid,
+                            )),
+                        }?;
+
+                        let ephemeral_stepoutput =
+                            StepOutput::new(0, StepOutputIndex::Change(ephemeral_output_index));
 
                         let tr0 = TransactionRequest::from_indexed(
                             transaction_request
diff --git a/zcash_client_backend/src/proposal.rs b/zcash_client_backend/src/proposal.rs
index e67e3e676..e88dbcf8d 100644
--- a/zcash_client_backend/src/proposal.rs
+++ b/zcash_client_backend/src/proposal.rs
@@ -55,6 +55,10 @@ pub enum ProposalError {
     /// The proposal included a payment to a TEX address and a spend from a shielded input in the same step.
     #[cfg(feature = "transparent-inputs")]
     PaysTexFromShielded,
+    /// The change strategy provided to input selection failed to correctly generate an ephemeral
+    /// change output when needed for sending to a TEX address.
+    #[cfg(feature = "transparent-inputs")]
+    EphemeralOutputsInvalid,
 }
 
 impl Display for ProposalError {
@@ -114,6 +118,11 @@ impl Display for ProposalError {
                 f,
                 "The proposal included a payment to a TEX address and a spend from a shielded input in the same step.",
             ),
+            #[cfg(feature = "transparent-inputs")]
+            ProposalError::EphemeralOutputsInvalid => write!(
+                f,
+                "The change strategy provided to input selection failed to correctly generate an ephemeral change output when needed for sending to a TEX address."
+            ),
         }
     }
 }