From 21eca164d89e862e1122c5444239432e181ae04c Mon Sep 17 00:00:00 2001
From: Deirdre Connolly <deirdre@zfnd.org>
Date: Sat, 18 Apr 2020 21:38:21 -0400
Subject: [PATCH] Refine Output description ephemeral_key to
 jubjub::AffinePoint

And impl Arbitrary for Output to support better generation of those points in proptests.
---
 zebra-chain/src/transaction/serialize.rs     |  4 +-
 zebra-chain/src/transaction/shielded_data.rs | 43 +++++++++++++++++---
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/zebra-chain/src/transaction/serialize.rs b/zebra-chain/src/transaction/serialize.rs
index 25fcdaabd52..bcaf42c1069 100644
--- a/zebra-chain/src/transaction/serialize.rs
+++ b/zebra-chain/src/transaction/serialize.rs
@@ -332,7 +332,7 @@ impl ZcashSerialize for Output {
     fn zcash_serialize<W: io::Write>(&self, mut writer: W) -> Result<(), io::Error> {
         writer.write_all(&self.cv[..])?;
         writer.write_all(&self.cmu[..])?;
-        writer.write_all(&self.ephemeral_key[..])?;
+        writer.write_all(&self.ephemeral_key.to_bytes())?;
         self.enc_ciphertext.zcash_serialize(&mut writer)?;
         self.out_ciphertext.zcash_serialize(&mut writer)?;
         self.zkproof.zcash_serialize(&mut writer)?;
@@ -345,7 +345,7 @@ impl ZcashDeserialize for Output {
         Ok(Output {
             cv: reader.read_32_bytes()?,
             cmu: reader.read_32_bytes()?,
-            ephemeral_key: reader.read_32_bytes()?,
+            ephemeral_key: jubjub::AffinePoint::from_bytes(reader.read_32_bytes()?).unwrap(),
             enc_ciphertext: shielded_data::EncryptedCiphertext::zcash_deserialize(&mut reader)?,
             out_ciphertext: shielded_data::OutCiphertext::zcash_deserialize(&mut reader)?,
             zkproof: Groth16Proof::zcash_deserialize(&mut reader)?,
diff --git a/zebra-chain/src/transaction/shielded_data.rs b/zebra-chain/src/transaction/shielded_data.rs
index fcaed938bc4..256c1395696 100644
--- a/zebra-chain/src/transaction/shielded_data.rs
+++ b/zebra-chain/src/transaction/shielded_data.rs
@@ -75,8 +75,7 @@ impl Arbitrary for Spend {
 /// A _Output Description_, as described in [protocol specification §7.4][ps].
 ///
 /// [ps]: https://zips.z.cash/protocol/protocol.pdf#outputencoding
-#[derive(Clone, Debug, PartialEq, Eq)]
-#[cfg_attr(test, derive(Arbitrary))]
+#[derive(Clone, Debug, PartialEq)]
 pub struct Output {
     /// A value commitment to the value of the input note.
     ///
@@ -87,9 +86,7 @@ pub struct Output {
     /// XXX refine to a specific type.
     pub cmu: [u8; 32],
     /// An encoding of an ephemeral Jubjub public key.
-    ///
-    /// XXX refine to a Jubjub key agreement type, not RedJubjub.
-    pub ephemeral_key: [u8; 32],
+    pub ephemeral_key: jubjub::AffinePoint,
     /// A ciphertext component for the encrypted output note.
     pub enc_ciphertext: EncryptedCiphertext,
     /// A ciphertext component for the encrypted output note.
@@ -98,6 +95,42 @@ pub struct Output {
     pub zkproof: Groth16Proof,
 }
 
+impl Eq for Output {}
+
+#[cfg(test)]
+impl Arbitrary for Output {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        (
+            array::uniform32(any::<u8>()),
+            array::uniform32(any::<u8>()),
+            array::uniform32(any::<u8>()).prop_filter("Valid jubjub::AffinePoint", |b| {
+                jubjub::AffinePoint::from_bytes(*b).is_some().unwrap_u8() == 1
+            }),
+            any::<EncryptedCiphertext>(),
+            any::<OutCiphertext>(),
+            any::<Groth16Proof>(),
+        )
+            .prop_map(
+                |(cv, cmu, ephemeral_key_bytes, enc_ciphertext, out_ciphertext, zkproof)| {
+                    return Self {
+                        cv,
+                        cmu,
+                        ephemeral_key: jubjub::AffinePoint::from_bytes(ephemeral_key_bytes)
+                            .unwrap(),
+                        enc_ciphertext,
+                        out_ciphertext,
+                        zkproof,
+                    };
+                },
+            )
+            .boxed()
+    }
+
+    type Strategy = BoxedStrategy<Self>;
+}
+
 /// Sapling-on-Groth16 spend and output descriptions.
 #[derive(Clone, Debug)]
 pub struct ShieldedData {