From 96e71198665fcc35929d1be1f1f8347bf12ad166 Mon Sep 17 00:00:00 2001
From: Murisi Tarusenga <murisit@gmail.com>
Date: Fri, 30 Aug 2024 12:47:14 +0200
Subject: [PATCH] Combine the computation of the extended full viewing key
 parts.

---
 app/rust/include/rslib.h     |  1 +
 app/rust/src/zip32_extern.rs | 23 +++++++++++++++++++++++
 app/src/crypto.c             | 26 ++++++--------------------
 app/src/keys_def.h           |  6 ++----
 4 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/app/rust/include/rslib.h b/app/rust/include/rslib.h
index eb22e87..8f5a643 100644
--- a/app/rust/include/rslib.h
+++ b/app/rust/include/rslib.h
@@ -19,3 +19,4 @@ void diversifier_find_valid(uint32_t zip32_account, uint8_t *default_diversifier
 void zip32_dk(uint32_t zip32_account, uint8_t *dk);
 void zip32_chain_code(uint32_t zip32_account, uint8_t *chain_code);
 void zip32_parent_fvk_tag(uint32_t zip32_account, uint8_t *fvk_tag);
+void zip32_xfvk(uint32_t zip32_account, uint8_t *fvk_tag, uint8_t *chain_code, uint8_t *fvk, uint8_t *dk);
diff --git a/app/rust/src/zip32_extern.rs b/app/rust/src/zip32_extern.rs
index ac40ea0..cfa0c3c 100644
--- a/app/rust/src/zip32_extern.rs
+++ b/app/rust/src/zip32_extern.rs
@@ -232,3 +232,26 @@ pub extern "C" fn zip32_parent_fvk_tag(
 
     fvk_tag.copy_from_slice(&tag);
 }
+
+#[no_mangle]
+pub extern "C" fn zip32_xfvk(
+    account: u32,
+    fvk_tag_ptr: *mut FvkTagBytes,
+    cc_ptr: *mut Zip32MasterChainCode,
+    fvk_ptr: *mut FullViewingKey,
+    dk_ptr: *mut DkBytes,
+) {
+    let path = [ZIP32_PURPOSE, ZIP32_COIN_TYPE, account];
+    let fvk_tag = unsafe { &mut *fvk_tag_ptr };
+    let cc = unsafe { &mut *cc_ptr };
+    let fvk_out = unsafe { &mut *fvk_ptr };
+    let dk = unsafe { &mut *dk_ptr };
+
+    let (key_bundle, chain_code, tag) = zip32_sapling_derive(&path);
+
+    fvk_tag.copy_from_slice(&tag);
+    cc.copy_from_slice(&chain_code);
+    let fvk = zip32_sapling_fvk(&key_bundle);
+    fvk_out.to_bytes_mut().copy_from_slice(fvk.to_bytes());
+    dk.copy_from_slice(&key_bundle.dk());
+}
diff --git a/app/src/crypto.c b/app/src/crypto.c
index acd9f6b..924735f 100644
--- a/app/src/crypto.c
+++ b/app/src/crypto.c
@@ -508,27 +508,15 @@ static zxerr_t computeKeys(keys_t * saplingKeys) {
         return zxerr_no_data;
     }
 
-    // Compute ask, nsk, ovk
-    zip32_child_ask_nsk(hdPath[2], saplingKeys->ask, saplingKeys->nsk);
-    zip32_ovk(hdPath[2], saplingKeys->ovk);
+    // Compute chain code, fvk, parent fvk tag, dk
+    zip32_xfvk(hdPath[2], saplingKeys->parent_fvk_tag, saplingKeys->chain_code, saplingKeys->fvk, saplingKeys->dk);
 
-    // Compute ak, nk, ivk
-    CHECK_PARSER_OK(generate_key(saplingKeys->ask, SpendingKeyGenerator, saplingKeys->ak));
-    CHECK_PARSER_OK(generate_key(saplingKeys->nsk, ProofGenerationKeyGenerator, saplingKeys->nk));
-    CHECK_PARSER_OK(computeIVK(saplingKeys->ak, saplingKeys->nk, saplingKeys->ivk));
+    // Compute ask, nsk
+    zip32_child_ask_nsk(hdPath[2], saplingKeys->ask, saplingKeys->nsk);
 
     // Compute diversifier
     diversifier_find_valid(hdPath[2], saplingKeys->diversifier);
 
-    // Compute dk
-    zip32_dk(hdPath[2], saplingKeys->dk);
-
-    // Compute chain code
-    zip32_chain_code(hdPath[2], saplingKeys->chain_code);
-
-    // Compute parent full viewing key tag
-    zip32_parent_fvk_tag(hdPath[2], saplingKeys->parent_fvk_tag);
-
     // Compute address
     get_pkd(hdPath[2], saplingKeys->diversifier, saplingKeys->address);
 
@@ -556,9 +544,7 @@ __Z_INLINE zxerr_t copyKeys(keys_t *saplingKeys, key_kind_e requestedKeys, uint8
             memcpy(output + 1, saplingKeys->parent_fvk_tag, TAG_LENGTH);
             memcpy(output + 5, &hdPath[hdPathLen - 1], TAG_LENGTH);
             memcpy(output + 9, saplingKeys->chain_code, KEY_LENGTH);
-            memcpy(output + 41, saplingKeys->ak, KEY_LENGTH);
-            memcpy(output + 73, saplingKeys->nk, KEY_LENGTH);
-            memcpy(output + 105, saplingKeys->ovk, KEY_LENGTH);
+            memcpy(output + 41, saplingKeys->fvk, KEY_LENGTH*3);
             memcpy(output + 137, saplingKeys->dk, KEY_LENGTH);
             break;
 
@@ -566,7 +552,7 @@ __Z_INLINE zxerr_t copyKeys(keys_t *saplingKeys, key_kind_e requestedKeys, uint8
             if (outputLen < 2 * KEY_LENGTH) {
                 return zxerr_buffer_too_small;
             }
-            memcpy(output, saplingKeys->ak, KEY_LENGTH);
+            memcpy(output, saplingKeys->fvk, KEY_LENGTH);
             memcpy(output + KEY_LENGTH, saplingKeys->nsk, KEY_LENGTH);
             break;
 
diff --git a/app/src/keys_def.h b/app/src/keys_def.h
index 72d0bbf..04cd036 100644
--- a/app/src/keys_def.h
+++ b/app/src/keys_def.h
@@ -51,16 +51,14 @@ typedef uint8_t ivk_t[KEY_LENGTH];
 typedef uint8_t ovk_t[KEY_LENGTH];
 typedef uint8_t d_t[DIVERSIFIER_LENGTH];
 typedef uint8_t fvk_tag_t[TAG_LENGTH];
+typedef uint8_t fvk_t[KEY_LENGTH*3];
 
 typedef uint8_t public_address_t[KEY_LENGTH];
 
 typedef struct {
     ask_t ask;
-    ak_t ak;
     nsk_t nsk;
-    nk_t nk;
-    ivk_t ivk;
-    ovk_t ovk;
+    fvk_t fvk;
     d_t diversifier;
     dk_t dk;
     chain_code_t chain_code;