From 5ff5643887720f7a936ce04acc3cce016cd6c2e7 Mon Sep 17 00:00:00 2001
From: ryan <ryan@distrust.co>
Date: Sun, 14 Apr 2024 23:08:25 -0400
Subject: [PATCH] keyfork-derive-util: allow zeroable input for key derivation

---
 .../keyfork-derive-util/src/private_key.rs    | 31 ++++++++++---------
 .../keyfork-derive-util/src/public_key.rs     | 22 ++++++++-----
 2 files changed, 30 insertions(+), 23 deletions(-)

diff --git a/crates/derive/keyfork-derive-util/src/private_key.rs b/crates/derive/keyfork-derive-util/src/private_key.rs
index 3e8f37e..2a600d3 100644
--- a/crates/derive/keyfork-derive-util/src/private_key.rs
+++ b/crates/derive/keyfork-derive-util/src/private_key.rs
@@ -102,6 +102,10 @@ pub enum PrivateKeyError {
     /// For the given algorithm, the private key must be nonzero.
     #[error("The provided private key must be nonzero, but is not")]
     NonZero,
+
+    /// A scalar could not be constructed for the given algorithm.
+    #[error("A scalar could not be constructed for the given algorithm")]
+    InvalidScalar,
 }
 
 #[cfg(feature = "secp256k1")]
@@ -130,20 +134,19 @@ impl PrivateKey for k256::SecretKey {
     }
 
     fn derive_child(&self, other: &PrivateKeyBytes) -> Result<Self, Self::Err> {
-        if other.iter().all(|n| n == &0) {
-            return Err(PrivateKeyError::NonZero);
-        }
-        let other = *other;
-        // Checked: See above nonzero check
-        let scalar = Option::<NonZeroScalar>::from(NonZeroScalar::from_repr(other.into()))
-            .expect(bug!("Should have been able to get a NonZeroScalar"));
+        use k256::elliptic_curve::ScalarPrimitive;
+        use k256::{Scalar, Secp256k1};
+
+        // Construct a scalar from bytes
+        let scalar = ScalarPrimitive::<Secp256k1>::from_bytes(other.into());
+        let scalar = Option::<ScalarPrimitive<Secp256k1>>::from(scalar);
+        let scalar = scalar.ok_or(PrivateKeyError::InvalidScalar)?;
+        let scalar = Scalar::from(scalar);
 
         let derived_scalar = self.to_nonzero_scalar().as_ref() + scalar.as_ref();
-        Ok(
-            Option::<NonZeroScalar>::from(NonZeroScalar::new(derived_scalar))
-                .map(Into::into)
-                .expect(bug!("Should be able to make Key")),
-        )
+        let nonzero_scalar = Option::<NonZeroScalar>::from(NonZeroScalar::new(derived_scalar))
+            .ok_or(PrivateKeyError::NonZero)?;
+        Ok(Self::from(nonzero_scalar))
     }
 }
 
@@ -202,9 +205,7 @@ impl PrivateKey for TestPrivateKey {
     type Err = PrivateKeyError;
 
     fn from_bytes(b: &PrivateKeyBytes) -> Self {
-        Self {
-            key: *b
-        }
+        Self { key: *b }
     }
 
     fn to_bytes(&self) -> PrivateKeyBytes {
diff --git a/crates/derive/keyfork-derive-util/src/public_key.rs b/crates/derive/keyfork-derive-util/src/public_key.rs
index 0468641..ef89244 100644
--- a/crates/derive/keyfork-derive-util/src/public_key.rs
+++ b/crates/derive/keyfork-derive-util/src/public_key.rs
@@ -77,6 +77,10 @@ pub enum PublicKeyError {
     #[error("The provided public key must be nonzero, but is not")]
     NonZero,
 
+    /// A scalar could not be constructed for the given algorithm.
+    #[error("A scalar could not be constructed for the given algorithm")]
+    InvalidScalar,
+
     /// Public key derivation is unsupported for this algorithm.
     #[error("Public key derivation is unsupported for this algorithm")]
     DerivationUnsupported,
@@ -85,7 +89,7 @@ pub enum PublicKeyError {
 #[cfg(feature = "secp256k1")]
 use k256::{
     elliptic_curve::{group::prime::PrimeCurveAffine, sec1::ToEncodedPoint},
-    AffinePoint, NonZeroScalar,
+    AffinePoint,
 };
 
 #[cfg(feature = "secp256k1")]
@@ -105,14 +109,16 @@ impl PublicKey for k256::PublicKey {
     }
 
     fn derive_child(&self, other: PrivateKeyBytes) -> Result<Self, Self::Err> {
-        if other.iter().all(|n| n == &0) {
-            return Err(PublicKeyError::NonZero);
-        }
-        // Checked: See above
-        let scalar = Option::<NonZeroScalar>::from(NonZeroScalar::from_repr(other.into()))
-            .expect(bug!("Should have been able to get a NonZeroScalar"));
+        use k256::elliptic_curve::ScalarPrimitive;
+        use k256::{Secp256k1, Scalar};
 
-        let point = self.to_projective() + (AffinePoint::generator() * *scalar);
+        // Construct a scalar from bytes
+        let scalar = ScalarPrimitive::<Secp256k1>::from_bytes(&other.into());
+        let scalar = Option::<ScalarPrimitive<Secp256k1>>::from(scalar);
+        let scalar = scalar.ok_or(PublicKeyError::InvalidScalar)?;
+        let scalar = Scalar::from(scalar);
+
+        let point = self.to_projective() + (AffinePoint::generator() * scalar);
         Ok(Self::from_affine(point.into())
             .expect(bug!("Could not from_affine after scalar arithmetic")))
     }