From 34277d4e38e9b710222b91321ae45f90f3bb5816 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@electriccoin.co>
Date: Fri, 17 Dec 2021 03:29:21 +0000
Subject: [PATCH] zcash_note_encryption: Use `*PlaintextBytes` structs in
 `Domain` APIs

`Domain::parse_note_plaintext_without_memo_ivk` is used with both full
note plaintexts and compact notes, so continues to accept a slice. For
all other `Domain` APIs, we constrain the input to `NotePlaintextBytes`
or `OutPlaintextBytes` as appropriate.


Extracted from: https://github.com/zcash/librustzcash/commit/7c1687dcc157df3c28b589d2073104c9c11e5289
---
 src/lib.rs | 32 ++++++++++++++++----------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 062bf7f..788c4e3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -166,19 +166,17 @@ pub trait Domain {
         pk_d: &Self::DiversifiedTransmissionKey,
         esk: &Self::EphemeralSecretKey,
         ephemeral_key: &EphemeralKeyBytes,
-        plaintext: &[u8],
+        plaintext: &NotePlaintextBytes,
     ) -> Option<(Self::Note, Self::Recipient)>;
 
     // &self is passed here in anticipation of future changes
     // to memo handling where the memos may no longer be
     // part of the note plaintext.
-    fn extract_memo(&self, plaintext: &[u8]) -> Self::Memo;
+    fn extract_memo(&self, plaintext: &NotePlaintextBytes) -> Self::Memo;
 
-    fn extract_pk_d(
-        out_plaintext: &[u8; OUT_PLAINTEXT_SIZE],
-    ) -> Option<Self::DiversifiedTransmissionKey>;
+    fn extract_pk_d(out_plaintext: &OutPlaintextBytes) -> Option<Self::DiversifiedTransmissionKey>;
 
-    fn extract_esk(out_plaintext: &[u8; OUT_PLAINTEXT_SIZE]) -> Option<Self::EphemeralSecretKey>;
+    fn extract_esk(out_plaintext: &OutPlaintextBytes) -> Option<Self::EphemeralSecretKey>;
 }
 
 #[cfg(feature = "alloc")]
@@ -420,14 +418,14 @@ fn try_note_decryption_inner<D: Domain, Output: ShieldedOutput<D>>(
     let enc_ciphertext = output.enc_ciphertext();
     assert_eq!(enc_ciphertext.len(), ENC_CIPHERTEXT_SIZE);
 
-    let mut plaintext: [u8; NOTE_PLAINTEXT_SIZE] =
-        enc_ciphertext[..NOTE_PLAINTEXT_SIZE].try_into().unwrap();
+    let mut plaintext =
+        NotePlaintextBytes(enc_ciphertext[..NOTE_PLAINTEXT_SIZE].try_into().unwrap());
 
     ChaCha20Poly1305::new(key.as_ref().into())
         .decrypt_in_place_detached(
             [0u8; 12][..].into(),
             &[],
-            &mut plaintext,
+            &mut plaintext.0,
             enc_ciphertext[NOTE_PLAINTEXT_SIZE..].into(),
         )
         .ok()?;
@@ -437,7 +435,7 @@ fn try_note_decryption_inner<D: Domain, Output: ShieldedOutput<D>>(
         ivk,
         ephemeral_key,
         &output.cmstar_bytes(),
-        &plaintext,
+        &plaintext.0,
     )?;
     let memo = domain.extract_memo(&plaintext);
 
@@ -569,14 +567,14 @@ pub fn try_output_recovery_with_ock<D: Domain, Output: ShieldedOutput<D>>(
     assert_eq!(enc_ciphertext.len(), ENC_CIPHERTEXT_SIZE);
     assert_eq!(out_ciphertext.len(), OUT_CIPHERTEXT_SIZE);
 
-    let mut op = [0; OUT_PLAINTEXT_SIZE];
-    op.copy_from_slice(&out_ciphertext[..OUT_PLAINTEXT_SIZE]);
+    let mut op = OutPlaintextBytes([0; OUT_PLAINTEXT_SIZE]);
+    op.0.copy_from_slice(&out_ciphertext[..OUT_PLAINTEXT_SIZE]);
 
     ChaCha20Poly1305::new(ock.as_ref().into())
         .decrypt_in_place_detached(
             [0u8; 12][..].into(),
             &[],
-            &mut op,
+            &mut op.0,
             out_ciphertext[OUT_PLAINTEXT_SIZE..].into(),
         )
         .ok()?;
@@ -591,14 +589,16 @@ pub fn try_output_recovery_with_ock<D: Domain, Output: ShieldedOutput<D>>(
     // be okay.
     let key = D::kdf(shared_secret, &ephemeral_key);
 
-    let mut plaintext = [0; NOTE_PLAINTEXT_SIZE];
-    plaintext.copy_from_slice(&enc_ciphertext[..NOTE_PLAINTEXT_SIZE]);
+    let mut plaintext = NotePlaintextBytes([0; NOTE_PLAINTEXT_SIZE]);
+    plaintext
+        .0
+        .copy_from_slice(&enc_ciphertext[..NOTE_PLAINTEXT_SIZE]);
 
     ChaCha20Poly1305::new(key.as_ref().into())
         .decrypt_in_place_detached(
             [0u8; 12][..].into(),
             &[],
-            &mut plaintext,
+            &mut plaintext.0,
             enc_ciphertext[NOTE_PLAINTEXT_SIZE..].into(),
         )
         .ok()?;