keyfork/crates/derive/keyfork-derive-util/src/public_key.rs

use crate::private_key::PrivateKeyBytes;

use digest::Digest;
use ripemd::Ripemd160;
use sha2::Sha256;
use thiserror::Error;

use keyfork_bug::bug;

pub(crate) type PublicKeyBytes = [u8; 33];

/// Functions required to use an `ExtendedPublicKey`.
pub trait PublicKey: Sized {
    /// The error returned by [`PublicKey::derive_child()`].
    type Err: std::error::Error;

    /*
     * This may not be doable given ed25519 public keys must be derived from the private key.
    /// Create a Self from bytes.
    fn from_bytes(b: &PublicKeyBytes) -> Self;
    */

    /// Convert a &Self to bytes.
    ///
    /// # Examples
    /// ```rust
    /// # use keyfork_derive_util::{
    /// #   *,
    /// #   private_key::TestPrivateKey as OurPrivateKey,
    /// # };
    /// let key_data: &[u8; 32] = //
    /// #   b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
    /// let private_key = OurPrivateKey::from_bytes(key_data);
    /// let public_key_bytes = private_key.public_key().to_bytes();
    /// ```
    fn to_bytes(&self) -> PublicKeyBytes;

    /// Derive a child [`PublicKey`] with given `PrivateKeyBytes`. The implementation of
    /// derivation is algorithm-specific and a specification should be consulted when implementing
    /// this method.
    ///
    /// # Errors
    ///
    /// An error may be returned if:
    /// * A nonzero `other` is provided.
    /// * An error specific to the given algorithm was encountered.
    fn derive_child(&self, other: PrivateKeyBytes) -> Result<Self, Self::Err>;

    /// Create a BIP-0032/SLIP-0010 fingerprint from the public key.
    ///
    /// # Examples
    /// ```rust
    /// # use keyfork_derive_util::{
    /// #   *,
    /// #   private_key::TestPrivateKey as OurPrivateKey,
    /// # };
    /// let key_data: &[u8; 32] = //
    /// #   b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
    /// let private_key = OurPrivateKey::from_bytes(key_data);
    /// let fingerprint = private_key.public_key().fingerprint();
    /// ```
    fn fingerprint(&self) -> [u8; 4] {
        let hash = Sha256::new().chain_update(self.to_bytes()).finalize();
        let hash = Ripemd160::new().chain_update(hash).finalize();
        // Note: Safety assured by type returned from Ripemd160
        hash[..4]
            .try_into()
            .expect(bug!("Ripemd160 returned too little data"))
    }
}

/// Errors associated with creating and arithmetic on public keys. This specific error is only
/// intended to be used by the implementations in this crate.
#[derive(Clone, Debug, Error)]
pub enum PublicKeyError {
    /// For the given algorithm, the private key must be nonzero.
    #[error("The provided public key must be nonzero, but is not")]
    NonZero,

    /// Public key derivation is unsupported for this algorithm.
    #[error("Public key derivation is unsupported for this algorithm")]
    DerivationUnsupported,
}

#[cfg(feature = "secp256k1")]
use k256::{
    elliptic_curve::{group::prime::PrimeCurveAffine, sec1::ToEncodedPoint},
    AffinePoint, NonZeroScalar,
};

#[cfg(feature = "secp256k1")]
impl PublicKey for k256::PublicKey {
    type Err = PublicKeyError;

    /*
    fn from_bytes(b: &PublicKeyBytes) -> Self {
        Self::from_sec1_bytes(b).expect("Invalid public key bytes")
    }
    */

    fn to_bytes(&self) -> PublicKeyBytes {
        let mut result = [0u8; 33];
        result[..].copy_from_slice(self.to_encoded_point(true).as_bytes());
        result
    }

    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"));

        let point = self.to_projective() + (AffinePoint::generator() * *scalar);
        Ok(Self::from_affine(point.into())
            .expect(bug!("Could not from_affine after scalar arithmetic")))
    }
}

#[cfg(feature = "ed25519")]
use ed25519_dalek::VerifyingKey;

#[cfg(feature = "ed25519")]
impl PublicKey for VerifyingKey {
    type Err = PublicKeyError;

    /*
    fn from_bytes(b: &PublicKeyBytes) -> Self {
        Self::from_bytes(b).expect("Invalid public key bytes")
    }
    */

    fn to_bytes(&self) -> PublicKeyBytes {
        let mut result = [0u8; 33];
        result[1..33].copy_from_slice(&self.to_bytes()[..]);
        result
    }

    fn derive_child(&self, _other: PrivateKeyBytes) -> Result<Self, Self::Err> {
        Err(Self::Err::DerivationUnsupported)
    }
}

#[doc(hidden)]
#[derive(Clone)]
pub struct TestPublicKey {
    pub(crate) key: [u8; 33],
}

impl TestPublicKey {
    #[doc(hidden)]
    #[allow(dead_code)]
    pub fn from_bytes(b: &[u8]) -> Self {
        Self {
            key: b
                .try_into()
                .expect(bug!("invalid size when constructing TestPublicKey")),
        }
    }
}

impl PublicKey for TestPublicKey {
    type Err = PublicKeyError;

    fn to_bytes(&self) -> PublicKeyBytes {
        self.key
    }

    fn derive_child(&self, _other: PrivateKeyBytes) -> Result<Self, Self::Err> {
        // whatever it takes for tests to pass...
        Ok(self.clone())
    }
}