From d7d89868491ceb25ec8fa56806964f1545894a58 Mon Sep 17 00:00:00 2001
From: Deirdre Connolly <durumcrustulum@gmail.com>
Date: Fri, 23 Apr 2021 23:17:01 -0400
Subject: [PATCH] Tidy up network handling for Orchard keys/addresses

---
 zebra-chain/src/orchard/address.rs    |  9 +++++----
 zebra-chain/src/orchard/keys.rs       | 23 +++++++++--------------
 zebra-chain/src/orchard/keys/tests.rs |  7 ++-----
 3 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/zebra-chain/src/orchard/address.rs b/zebra-chain/src/orchard/address.rs
index 2da7a5dc2..4f0dafc8e 100644
--- a/zebra-chain/src/orchard/address.rs
+++ b/zebra-chain/src/orchard/address.rs
@@ -64,10 +64,11 @@ mod tests {
     fn derive_keys_and_addresses() {
         zebra_test::init();
 
-        let spending_key = keys::SpendingKey::new(&mut OsRng);
+        let network = Network::Mainnet;
 
-        let full_viewing_key =
-            keys::FullViewingKey::from_spending_key(spending_key, Network::Mainnet);
+        let spending_key = keys::SpendingKey::new(&mut OsRng, network);
+
+        let full_viewing_key = keys::FullViewingKey::from_spending_key(spending_key);
 
         // Default diversifier, where index = 0.
         let diversifier_key = keys::DiversifierKey::from(full_viewing_key);
@@ -78,7 +79,7 @@ mod tests {
         let transmission_key = keys::TransmissionKey::from((incoming_viewing_key, diversifier));
 
         let _orchard_shielded_address = Address {
-            network: Network::Mainnet,
+            network,
             diversifier,
             transmission_key,
         };
diff --git a/zebra-chain/src/orchard/keys.rs b/zebra-chain/src/orchard/keys.rs
index f6f1c8f68..ab669ed00 100644
--- a/zebra-chain/src/orchard/keys.rs
+++ b/zebra-chain/src/orchard/keys.rs
@@ -144,16 +144,6 @@ pub struct SpendingKey {
     bytes: [u8; 32],
 }
 
-impl From<[u8; 32]> for SpendingKey {
-    /// Generate a _SpendingKey_ from existing bytes.
-    fn from(bytes: [u8; 32]) -> Self {
-        Self {
-            network: Network::default(),
-            bytes,
-        }
-    }
-}
-
 impl From<SpendingKey> for [u8; 32] {
     fn from(sk: SpendingKey) -> Self {
         sk.bytes
@@ -203,7 +193,7 @@ impl SpendingKey {
     /// is not zero, else fail.
     ///
     /// [orchardkeycomponents]: https://zips.z.cash/protocol/nu5.pdf#orchardkeycomponents
-    pub fn new<T>(csprng: &mut T) -> Self
+    pub fn new<T>(csprng: &mut T, network: Network) -> Self
     where
         T: RngCore + CryptoRng,
     {
@@ -211,7 +201,7 @@ impl SpendingKey {
             let mut bytes = [0u8; 32];
             csprng.fill_bytes(&mut bytes);
 
-            let sk = Self::from(bytes);
+            let sk = Self::from_bytes(bytes, network);
 
             // "if ask = 0, discard this key and repeat with a new sk"
             if SpendAuthorizingKey::from(sk).0 == pallas::Scalar::zero() {
@@ -221,6 +211,11 @@ impl SpendingKey {
             break sk;
         }
     }
+
+    /// Generate a `SpendingKey` from existing bytes.
+    fn from_bytes(bytes: [u8; 32], network: Network) -> Self {
+        Self { network, bytes }
+    }
 }
 
 /// A Spend authorizing key (_ask_), as described in [protocol specification
@@ -643,11 +638,11 @@ impl FullViewingKey {
     ///
     /// https://zips.z.cash/protocol/nu5.pdf#addressesandkeys
     /// https://zips.z.cash/protocol/nu5.pdf#orchardfullviewingkeyencoding
-    pub fn from_spending_key(sk: SpendingKey, network: Network) -> FullViewingKey {
+    pub fn from_spending_key(sk: SpendingKey) -> FullViewingKey {
         let spend_authorizing_key = SpendAuthorizingKey::from(sk);
 
         Self {
-            network,
+            network: sk.network,
             spend_validating_key: SpendValidatingKey::from(spend_authorizing_key),
             nullifier_deriving_key: NullifierDerivingKey::from(sk),
             ivk_commit_randomness: IvkCommitRandomness::from(sk),
diff --git a/zebra-chain/src/orchard/keys/tests.rs b/zebra-chain/src/orchard/keys/tests.rs
index 8f8bd6cfe..84ec4abca 100644
--- a/zebra-chain/src/orchard/keys/tests.rs
+++ b/zebra-chain/src/orchard/keys/tests.rs
@@ -11,8 +11,7 @@ impl Arbitrary for TransmissionKey {
     fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
         (any::<SpendingKey>())
             .prop_map(|spending_key| {
-                let full_viewing_key =
-                    FullViewingKey::from_spending_key(spending_key, Network::Mainnet);
+                let full_viewing_key = FullViewingKey::from_spending_key(spending_key);
 
                 let diversifier_key = DiversifierKey::from(full_viewing_key);
 
@@ -56,9 +55,7 @@ proptest! {
         prop_assert_eq![full_viewing_key, full_viewing_key_2];
 
         let diversifier_key = DiversifierKey::from(full_viewing_key);
-
-        let mut incoming_viewing_key = IncomingViewingKey::from(full_viewing_key);
-        incoming_viewing_key.network = spending_key.network;
+        let incoming_viewing_key = IncomingViewingKey::from(full_viewing_key);
 
         let ivk_string = incoming_viewing_key.to_string();
         let incoming_viewing_key_2: IncomingViewingKey = ivk_string.parse().unwrap();