From 6904c8f933fb74542ebef9b7c889a284c84d4aa4 Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@electriccoin.co>
Date: Thu, 30 Jul 2020 23:07:33 +0800
Subject: [PATCH] Implement plaintext_version_is_valid()

---
 zcash_primitives/src/consensus.rs           |  4 +++
 zcash_primitives/src/note_encryption.rs     | 38 +++++++++++++++++----
 zcash_primitives/src/transaction/builder.rs | 10 ++++--
 3 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/zcash_primitives/src/consensus.rs b/zcash_primitives/src/consensus.rs
index c2d42b756..719231f02 100644
--- a/zcash_primitives/src/consensus.rs
+++ b/zcash_primitives/src/consensus.rs
@@ -13,6 +13,10 @@ pub trait Parameters {
             _ => false,
         }
     }
+
+    fn zip_212_grace_period(&self) -> u32 {
+        32256
+    }
 }
 
 /// Marker struct for the production network.
diff --git a/zcash_primitives/src/note_encryption.rs b/zcash_primitives/src/note_encryption.rs
index 556e22e8b..b819acebc 100644
--- a/zcash_primitives/src/note_encryption.rs
+++ b/zcash_primitives/src/note_encryption.rs
@@ -344,9 +344,9 @@ fn parse_note_plaintext_without_memo<P: consensus::Parameters>(
     plaintext: &[u8],
 ) -> Option<(Note<Bls12>, PaymentAddress<Bls12>)> {
     // Check note plaintext version
-    match plaintext[0] {
-        0x01 => (),
-        _ => return None,
+    match plaintext_version_is_valid(parameters, height, plaintext[0]) {
+        true => (),
+        false => return None,
     }
 
     let mut d = [0u8; 11];
@@ -380,6 +380,32 @@ fn parse_note_plaintext_without_memo<P: consensus::Parameters>(
     Some((note, to))
 }
 
+pub fn plaintext_version_is_valid<P: consensus::Parameters>(
+    parameters: &P,
+    height: u32,
+    leadbyte: u8,
+) -> bool {
+    if parameters.is_nu_active(NetworkUpgrade::Canopy, height) {
+        let grace_period_end_height = parameters
+            .activation_height(NetworkUpgrade::Canopy)
+            .expect("Should have Canopy activation height")
+            + parameters.zip_212_grace_period();
+
+        if height < grace_period_end_height && leadbyte != 0x01 && leadbyte != 0x02 {
+            // non-{0x01,0x02} received after Canopy activation and before grace period has elapsed
+            false
+        } else if height >= grace_period_end_height && leadbyte != 0x02 {
+            // non-0x02 received past (Canopy activation height + grace period)
+            false
+        } else {
+            true
+        }
+    } else {
+        // return false if non-0x01 received when Canopy is not active
+        leadbyte == 0x01
+    }
+}
+
 /// Trial decryption of the full note plaintext by the recipient.
 ///
 /// Attempts to decrypt and validate the given `enc_ciphertext` using the given `ivk`.
@@ -510,9 +536,9 @@ pub fn try_sapling_output_recovery<P: consensus::Parameters>(
     );
 
     // Check note plaintext version
-    match plaintext[0] {
-        0x01 => (),
-        _ => return None,
+    match plaintext_version_is_valid(parameters, height, plaintext[0]) {
+        true => (),
+        false => return None,
     }
 
     let mut d = [0u8; 11];
diff --git a/zcash_primitives/src/transaction/builder.rs b/zcash_primitives/src/transaction/builder.rs
index 36062f39d..81198c914 100644
--- a/zcash_primitives/src/transaction/builder.rs
+++ b/zcash_primitives/src/transaction/builder.rs
@@ -104,13 +104,19 @@ impl SaplingOutput {
             return Err(Error::InvalidAmount);
         }
 
-        let rcm = Fs::random(rng);
+        let rseed = if parameters.is_nu_active(NetworkUpgrade::Canopy, height) {
+            let mut buffer = [0u8; 32];
+            &rng.fill_bytes(&mut buffer);
+            Rseed::AfterZip212(buffer)
+        } else {
+            Rseed::BeforeZip212(Fs::random(rng))
+        };
 
         let note = Note {
             g_d,
             pk_d: to.pk_d().clone(),
             value: value.into(),
-            rseed: Rseed::BeforeZip212(rcm),
+            rseed,
         };
 
         Ok(SaplingOutput {