Perform key type conversion inline with hardened-only derivation

src/arbitrary.rs

@@ -4,6 +4,8 @@
 //!
 //! [arbkd]: https://zips.z.cash/zip-0032#specification-arbitrary-key-derivation
 
+use core::convert::Infallible;
+
 use zcash_spec::PrfExpand;
 
 use crate::{
@@ -16,6 +18,17 @@ struct Arbitrary;
 impl Context for Arbitrary {
     const MKG_DOMAIN: [u8; 16] = *b"ZcashArbitraryKD";
     const CKD_DOMAIN: PrfExpand<([u8; 32], [u8; 4])> = PrfExpand::ARBITRARY_ZIP32_CHILD;
+
+    type Key = [u8; 32];
+    type KeyError = Infallible;
+
+    fn convert_to_key(sk: [u8; 32]) -> Result<Self::Key, Self::KeyError> {
+        Ok(sk)
+    }
+
+    fn key_bytes(sk: &Self::Key) -> &[u8; 32] {
+        sk
+    }
 }
 
 /// An arbitrary extended secret key.
@@ -71,7 +84,7 @@ impl SecretKey {
         ];
 
         Self {
-            inner: HardenedOnlyKey::master(ikm),
+            inner: HardenedOnlyKey::master(ikm).expect("infallible"),
         }
     }
 
@@ -82,7 +95,7 @@ impl SecretKey {
     /// [ckdarb]: https://zips.z.cash/zip-0032#arbitrary-child-key-derivation
     fn derive_child(&self, index: ChildIndex) -> Self {
         Self {
-            inner: self.inner.derive_child(index),
+            inner: self.inner.derive_child(index).expect("infallible"),
         }
     }
 }

src/hardened_only.rs

@@ -12,6 +12,7 @@
 use core::marker::PhantomData;
 
 use blake2b_simd::Params as Blake2bParams;
+use subtle::{Choice, ConstantTimeEq};
 use zcash_spec::PrfExpand;
 
 use crate::{ChainCode, ChildIndex};
@@ -25,6 +26,21 @@ pub trait Context {
     const MKG_DOMAIN: [u8; 16];
     /// The `PrfExpand` domain used during child key derivation.
     const CKD_DOMAIN: PrfExpand<([u8; 32], [u8; 4])>;
+
+    /// The context-specific key type to which the secret key is converted at each layer.
+    ///
+    /// If no specific key type is required, this can be `[u8; 32]`.
+    type Key;
+    /// Errors that can occur when converting the secret key to the context-specific key
+    /// type, if any.
+    ///
+    /// If no specific key type is required, this can be [`std::convert::Infallible`].
+    type KeyError;
+
+    /// Converts the hardened-only secret key to the context-specific key type.
+    fn convert_to_key(sk: [u8; 32]) -> Result<Self::Key, Self::KeyError>;
+    /// Exposes the bytes of the context-specific key type.
+    fn key_bytes(sk: &Self::Key) -> &[u8; 32];
 }
 
 /// An arbitrary extended secret key.
@@ -34,33 +50,34 @@ pub trait Context {
 /// [hkd]: https://zips.z.cash/zip-0032#specification-hardened-only-key-derivation
 #[derive(Clone, Debug)]
 pub struct HardenedOnlyKey<C: Context> {
-    sk: [u8; 32],
+    sk: C::Key,
     chain_code: ChainCode,
     _context: PhantomData<C>,
 }
 
-#[allow(non_snake_case)]
-impl<C: Context> HardenedOnlyKey<C> {
-    /// Constructs a hardened-only extended secret key from its parts.
-    pub fn from_parts(sk: [u8; 32], chain_code: ChainCode) -> Self {
-        Self {
-            sk,
-            chain_code,
-            _context: PhantomData,
-        }
+impl<C> ConstantTimeEq for HardenedOnlyKey<C>
+where
+    C: Context,
+    C::Key: ConstantTimeEq,
+{
+    fn ct_eq(&self, rhs: &Self) -> Choice {
+        self.chain_code.ct_eq(&rhs.chain_code) & self.sk.ct_eq(&rhs.sk)
     }
+}
 
+#[allow(non_snake_case)]
+impl<C: Context> HardenedOnlyKey<C> {
     /// Exposes the parts of this key.
-    pub fn into_parts(self) -> ([u8; 32], ChainCode) {
-        (self.sk, self.chain_code)
+    pub fn parts(&self) -> (&C::Key, &ChainCode) {
+        (&self.sk, &self.chain_code)
     }
 
     /// Generates the master key of a hardened-only extended secret key.
     ///
     /// Defined in [ZIP32: Hardened-only master key generation][mkgh].
     ///
     /// [mkgh]: https://zips.z.cash/zip-0032#hardened-only-master-key-generation
-    pub fn master(ikm: &[&[u8]]) -> Self {
+    pub fn master(ikm: &[&[u8]]) -> Result<Self, C::KeyError> {
         // I := BLAKE2b-512(Context.MKGDomain, IKM)
         let I: [u8; 64] = {
             let mut I = Blake2bParams::new()
@@ -76,43 +93,43 @@ impl<C: Context> HardenedOnlyKey<C> {
         let (I_L, I_R) = I.split_at(32);
 
         // I_L is used as the master secret key sk_m.
-        let sk_m = I_L.try_into().unwrap();
+        let sk_m = C::convert_to_key(I_L.try_into().unwrap())?;
 
         // I_R is used as the master chain code c_m.
         let c_m = ChainCode::new(I_R.try_into().unwrap());
 
-        Self {
+        Ok(Self {
             sk: sk_m,
             chain_code: c_m,
             _context: PhantomData,
-        }
+        })
     }
 
     /// Derives a child key from a parent key at a given index.
     ///
     /// Defined in [ZIP32: Hardened-only child key derivation][ckdh].
     ///
     /// [ckdh]: https://zips.z.cash/zip-0032#hardened-only-child-key-derivation
-    pub fn derive_child(&self, index: ChildIndex) -> Self {
+    pub fn derive_child(&self, index: ChildIndex) -> Result<Self, C::KeyError> {
         // I := PRF^Expand(c_par, [Context.CKDDomain] || sk_par || I2LEOSP(i))
         let I: [u8; 64] = C::CKD_DOMAIN.with(
             self.chain_code.as_bytes(),
-            &self.sk,
+            C::key_bytes(&self.sk),
             &index.index().to_le_bytes(),
         );
 
         let (I_L, I_R) = I.split_at(32);
 
         // I_L is used as the child spending key sk_i.
-        let sk_i = I_L.try_into().unwrap();
+        let sk_i = C::convert_to_key(I_L.try_into().unwrap())?;
 
         // I_R is used as the child chain code c_i.
         let c_i = ChainCode::new(I_R.try_into().unwrap());
 
-        Self {
+        Ok(Self {
             sk: sk_i,
             chain_code: c_i,
             _context: PhantomData,
-        }
+        })
     }
 }