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 { - if other.iter().all(|n| n == &0) { - return Err(PrivateKeyError::NonZero); - } - let other = *other; - // Checked: See above nonzero check - let scalar = Option::::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::::from_bytes(other.into()); + let scalar = Option::>::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::::from(NonZeroScalar::new(derived_scalar)) - .map(Into::into) - .expect(bug!("Should be able to make Key")), - ) + let nonzero_scalar = Option::::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 { - if other.iter().all(|n| n == &0) { - return Err(PublicKeyError::NonZero); - } - // Checked: See above - let scalar = Option::::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::::from_bytes(&other.into()); + let scalar = Option::>::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"))) }