From 2b1583d75f263a5fa0323716a675e2644fd2f86c Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Sat, 23 Mar 2019 17:51:30 +1300
Subject: [PATCH] Deduplicate Sapling key agreement logic

---
 librustzcash/src/rustzcash.rs           | 11 ++++------
 sapling-crypto/src/jubjub/edwards.rs    | 16 +++++++++++++++
 zcash_primitives/src/note_encryption.rs | 27 ++++++++++++++++---------
 3 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/librustzcash/src/rustzcash.rs b/librustzcash/src/rustzcash.rs
index ab9357d..1c84c8d 100644
--- a/librustzcash/src/rustzcash.rs
+++ b/librustzcash/src/rustzcash.rs
@@ -58,7 +58,7 @@ use std::ffi::OsString;
 use std::os::windows::ffi::OsStringExt;
 
 use sapling_crypto::primitives::{ProofGenerationKey, ViewingKey};
-use zcash_primitives::{sapling::spend_sig, JUBJUB};
+use zcash_primitives::{note_encryption::sapling_ka_agree, sapling::spend_sig, JUBJUB};
 use zcash_proofs::{
     load_parameters,
     sapling::{CommitmentTreeWitness, SaplingProvingContext, SaplingVerificationContext},
@@ -536,15 +536,12 @@ pub extern "system" fn librustzcash_sapling_ka_agree(
         Err(_) => return false,
     };
 
-    // Multiply by 8
-    let p = p.mul_by_cofactor(&JUBJUB);
-
-    // Multiply by sk
-    let p = p.mul(sk, &JUBJUB);
+    // Compute key agreement
+    let ka = sapling_ka_agree(&sk, &p);
 
     // Produce result
     let result = unsafe { &mut *result };
-    p.write(&mut result[..]).expect("length is not 32 bytes");
+    result.copy_from_slice(&ka);
 
     true
 }
diff --git a/sapling-crypto/src/jubjub/edwards.rs b/sapling-crypto/src/jubjub/edwards.rs
index 3aa9345..16d21e7 100644
--- a/sapling-crypto/src/jubjub/edwards.rs
+++ b/sapling-crypto/src/jubjub/edwards.rs
@@ -45,6 +45,14 @@ fn convert_subgroup<E: JubjubEngine, S1, S2>(from: &Point<E, S1>) -> Point<E, S2
     }
 }
 
+impl<E: JubjubEngine> From<&Point<E, Unknown>> for Point<E, Unknown>
+{
+    fn from(p: &Point<E, Unknown>) -> Point<E, Unknown>
+    {
+        p.clone()
+    }
+}
+
 impl<E: JubjubEngine> From<Point<E, PrimeOrder>> for Point<E, Unknown>
 {
     fn from(p: Point<E, PrimeOrder>) -> Point<E, Unknown>
@@ -53,6 +61,14 @@ impl<E: JubjubEngine> From<Point<E, PrimeOrder>> for Point<E, Unknown>
     }
 }
 
+impl<E: JubjubEngine> From<&Point<E, PrimeOrder>> for Point<E, Unknown>
+{
+    fn from(p: &Point<E, PrimeOrder>) -> Point<E, Unknown>
+    {
+        convert_subgroup(p)
+    }
+}
+
 impl<E: JubjubEngine, Subgroup> Clone for Point<E, Subgroup>
 {
     fn clone(&self) -> Self {
diff --git a/zcash_primitives/src/note_encryption.rs b/zcash_primitives/src/note_encryption.rs
index be24178..11869f4 100644
--- a/zcash_primitives/src/note_encryption.rs
+++ b/zcash_primitives/src/note_encryption.rs
@@ -128,14 +128,21 @@ fn generate_esk() -> Fs {
     Fs::to_uniform(&buffer[..])
 }
 
-fn sapling_ka_agree(esk: &Fs, pk_d: &edwards::Point<Bls12, PrimeOrder>) -> Vec<u8> {
-    let ka = pk_d
-        .mul(esk.into_repr(), &JUBJUB)
-        .double(&JUBJUB)
-        .double(&JUBJUB)
-        .double(&JUBJUB);
-    let mut result = Vec::with_capacity(32);
-    ka.write(&mut result).expect("length is not 32 bytes");
+pub fn sapling_ka_agree<'a, P>(esk: &Fs, pk_d: &'a P) -> [u8; 32]
+where
+    edwards::Point<Bls12, Unknown>: From<&'a P>,
+{
+    let p: edwards::Point<Bls12, Unknown> = pk_d.into();
+
+    // Multiply by 8
+    let p = p.mul_by_cofactor(&JUBJUB);
+
+    // Multiply by esk
+    let p = p.mul(*esk, &JUBJUB);
+
+    // Produce result
+    let mut result = [0; 32];
+    p.write(&mut result[..]).expect("length is not 32 bytes");
     result
 }
 
@@ -294,7 +301,7 @@ pub fn try_sapling_note_decryption(
     cmu: &Fr,
     enc_ciphertext: &[u8],
 ) -> Option<(Note<Bls12>, PaymentAddress<Bls12>, Memo)> {
-    let shared_secret = sapling_ka_agree(&ivk, &epk);
+    let shared_secret = sapling_ka_agree(ivk, epk);
     let key = kdf_sapling(&shared_secret, &epk);
 
     let mut plaintext = Vec::with_capacity(564);
@@ -328,7 +335,7 @@ pub fn try_sapling_compact_note_decryption(
     cmu: &Fr,
     enc_ciphertext: &[u8],
 ) -> Option<(Note<Bls12>, PaymentAddress<Bls12>)> {
-    let shared_secret = sapling_ka_agree(&ivk, &epk);
+    let shared_secret = sapling_ka_agree(ivk, epk);
     let key = kdf_sapling(&shared_secret, &epk);
 
     let nonce = [0u8; 12];