From 53cae4647edced78088088d612c74cdf5dfba43f Mon Sep 17 00:00:00 2001
From: Henry de Valence <hdevalence@hdevalence.ca>
Date: Fri, 20 Dec 2019 19:44:42 -0800
Subject: [PATCH] Make invalid ShieldedData unrepresentable.

ShieldedData objects must have at least one spend or output; using Either
ensures that at least one must be present.  This is similar to the
JoinSplitData case, but slightly more complicated: rather than enforcing that
one list has at least one element (which can be done as `(first, rest)`), here
we need to use Either.   This has the downside that it is possible to construct
multiple equivalent internal representations (choosing whether a spend or
output goes in the `first` slot), but this easily fixed with a custom PartialEq
implementation.
---
 Cargo.lock                                   |  1 +
 zebra-chain/Cargo.toml                       |  3 +-
 zebra-chain/src/transaction/shielded_data.rs | 64 ++++++++++++++++++--
 3 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 494dc11a..f384a89a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1511,6 +1511,7 @@ version = "0.1.0"
 dependencies = [
  "byteorder 1.3.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "chrono 0.4.10 (registry+https://github.com/rust-lang/crates.io-index)",
+ "futures 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "hex 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "proptest 0.9.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "redjubjub 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
diff --git a/zebra-chain/Cargo.toml b/zebra-chain/Cargo.toml
index 606c3600..ed3155fa 100644
--- a/zebra-chain/Cargo.toml
+++ b/zebra-chain/Cargo.toml
@@ -14,6 +14,7 @@ chrono = "0.4"
 hex = "0.4"
 sha2 = "0.8"
 redjubjub = "0.1"
+futures = "0.3"
 
 [dev-dependencies]
-proptest = "0.9"
\ No newline at end of file
+proptest = "0.9"
diff --git a/zebra-chain/src/transaction/shielded_data.rs b/zebra-chain/src/transaction/shielded_data.rs
index 5b4b6724..ba7df5e4 100644
--- a/zebra-chain/src/transaction/shielded_data.rs
+++ b/zebra-chain/src/transaction/shielded_data.rs
@@ -1,3 +1,5 @@
+use futures::future::Either;
+
 // XXX this name seems too long?
 use crate::note_commitment_tree::SaplingNoteTreeRootHash;
 use crate::proofs::Groth16Proof;
@@ -58,12 +60,64 @@ pub struct OutputDescription {
 }
 
 /// Sapling-on-Groth16 spend and output descriptions.
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug)]
 pub struct ShieldedData {
-    /// A sequence of [`SpendDescription`]s for this transaction.
-    pub shielded_spends: Vec<SpendDescription>,
-    /// A sequence of shielded outputs for this transaction.
-    pub shielded_outputs: Vec<OutputDescription>,
+    /// Either a spend or output description.
+    ///
+    /// Storing this separately ensures that it is impossible to construct
+    /// an invalid `ShieldedData` with no spends or outputs.
+    pub first: Either<SpendDescription, OutputDescription>,
+    /// The rest of the [`SpendDescription`]s for this transaction.
+    pub rest_spends: Vec<SpendDescription>,
+    /// The rest of the [`OutputDescription`]s for this transaction.
+    pub rest_outputs: Vec<OutputDescription>,
     /// A signature on the transaction hash.
     pub binding_sig: redjubjub::Signature<Binding>,
 }
+
+impl ShieldedData {
+    /// Iterate over the [`SpendDescription`]s for this transaction.
+    pub fn spends(&self) -> impl Iterator<Item = &SpendDescription> {
+        match self.first {
+            Either::Left(ref spend) => Some(spend),
+            Either::Right(_) => None,
+        }
+        .into_iter()
+        .chain(self.rest_spends.iter())
+    }
+
+    /// Iterate over the [`OutputDescription`]s for this transaction.
+    pub fn outputs(&self) -> impl Iterator<Item = &OutputDescription> {
+        match self.first {
+            Either::Left(_) => None,
+            Either::Right(ref output) => Some(output),
+        }
+        .into_iter()
+        .chain(self.rest_outputs.iter())
+    }
+}
+
+// Technically, it's possible to construct two equivalent representations
+// of a ShieldedData with at least one spend and at least one output, depending
+// on which goes in the `first` slot.  This is annoying but a smallish price to
+// pay for structural validity.
+
+impl std::cmp::PartialEq for ShieldedData {
+    fn eq(&self, other: &Self) -> bool {
+        // First check that the lengths match, so we know it is safe to use zip,
+        // which truncates to the shorter of the two iterators.
+        if self.spends().count() != other.spends().count() {
+            return false;
+        }
+        if self.outputs().count() != other.outputs().count() {
+            return false;
+        }
+
+        // Now check that the binding_sig, spends, outputs match.
+        self.binding_sig == other.binding_sig
+            && self.spends().zip(other.spends()).all(|(a, b)| a == b)
+            && self.outputs().zip(other.outputs()).all(|(a, b)| a == b)
+    }
+}
+
+impl std::cmp::Eq for ShieldedData {}