keyfork/keyfork-derive-util/src/extended_key/private_key.rs

use crate::{DerivationIndex, DerivationPath, ExtendedPublicKey, PrivateKey, PublicKey};

use hmac::{Hmac, Mac};
use serde::{Deserialize, Serialize};
use sha2::Sha512;
use thiserror::Error;

const KEY_SIZE: usize = 256;

/// Errors associated with creating or deriving Extended Private Keys.
#[derive(Error, Clone, Debug)]
pub enum Error {
    /// The seed has an unsuitable length; supported lengths are 16 bytes, 32 bytes, or 64 bytes.
    #[error("Seed had an unsuitable length: {0}")]
    BadSeedLength(usize),

    /// The maximum depth for key derivation has been reached. The supported maximum depth is 255.
    #[error("Reached maximum depth for key derivation")]
    Depth,

    /// This should never happen. HMAC keys should be able to take any size input.
    #[error("Invalid length for HMAC key while generating master key (report me!)")]
    HmacInvalidLength(#[from] hmac::digest::InvalidLength),

    /// An unknown error occurred while deriving a child key.
    #[error("Unknown error while deriving child key")]
    Derivation,

    /// The algorithm used mandates hardened derivation only.
    #[error("The algorithm used mandates hardened derivation only")]
    HardenedDerivationRequired,

    /// The given slice was of an inappropriate size to create a Private Key.
    #[error("The given slice was of an inappropriate size to create a Private Key")]
    InvalidSliceError(#[from] std::array::TryFromSliceError),
}

type Result<T, E = Error> = std::result::Result<T, E>;
type ChainCode = [u8; 32];
type HmacSha512 = Hmac<Sha512>;

/// Extended private keys derived using BIP-0032.
///
/// Generic over types implementing [`PrivateKey`].
#[derive(Clone, Serialize, Deserialize)]
pub struct ExtendedPrivateKey<K: PrivateKey + Clone> {
    /// The internal private key data.
    private_key: K,
    depth: u8,
    chain_code: ChainCode,
}

impl<K: PrivateKey + Clone> std::fmt::Debug for ExtendedPrivateKey<K> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ExtendedPrivateKey")
            .field("private_key", &"obscured")
            .field("depth", &self.depth)
            .field("chain_code", &self.chain_code)
            .finish()
    }
}

impl<K> ExtendedPrivateKey<K>
where
    K: PrivateKey + Clone,
{
    /// Generate a new [`ExtendedPrivateKey`] from a seed, ideally from a 12-word or 24-word
    /// mnemonic, but may take 16-byte seeds.
    ///
    /// # Panics
    ///
    /// The method performs unchecked `try_into()` operations on a constant-sized slice.
    ///
    /// # Errors
    ///
    /// An error may be returned if:
    /// * The given seed had an incorrect length.
    /// * A `HmacSha512` can't be constructed - this should be impossible.
    pub fn new(seed: impl AsRef<[u8]>) -> Result<Self> {
        Self::new_internal(seed.as_ref())
    }

    fn new_internal(seed: &[u8]) -> Result<Self> {
        let len = seed.len();
        if ![16, 32, 64].contains(&len) {
            return Err(Error::BadSeedLength(len));
        }

        let hash = HmacSha512::new_from_slice(&K::key().bytes().collect::<Vec<_>>())?
            .chain_update(seed)
            .finalize()
            .into_bytes();
        let (private_key, chain_code) = hash.split_at(KEY_SIZE / 8);

        Self::new_from_parts(
            private_key,
            0,
            // Checked: chain_code is always the same length, hash is static size
            chain_code.try_into().expect("Invalid chain code length"),
        )
    }

    pub fn new_from_parts(seed: &[u8], depth: u8, chain_code: [u8; 32]) -> Result<Self> {
        Ok(Self {
            private_key: K::from_bytes(seed.try_into()?),
            depth,
            chain_code,
        })
    }

    /// Returns a reference to the [`PrivateKey`].
    pub fn private_key(&self) -> &K {
        &self.private_key
    }

    pub fn extended_public_key(&self) -> ExtendedPublicKey<K::PublicKey> {
        ExtendedPublicKey::new(self.public_key(), self.chain_code)
    }

    /// Return a public key for the current [`PrivateKey`].
    pub fn public_key(&self) -> K::PublicKey {
        self.private_key.public_key()
    }

    /// Returns the current depth.
    pub fn depth(&self) -> u8 {
        self.depth
    }

    /// Returns a copy of the current chain code.
    pub fn chain_code(&self) -> [u8; 32] {
        self.chain_code
    }

    /// Derive a child using the given [`DerivationPath`].
    ///
    /// # Errors
    ///
    /// An error may be returned under the same circumstances as
    /// [`ExtendedPrivateKey::derive_child`].
    pub fn derive_path(&self, path: &DerivationPath) -> Result<Self> {
        if path.path.is_empty() {
            Ok(self.clone())
        } else {
            path.iter()
                .try_fold(self.clone(), |key, index| key.derive_child(index))
        }
    }

    /// Derive a child with a given [`DerivationIndex`].
    ///
    /// # Panics
    ///
    /// The method performs unchecked `try_into()` operations on constant-sized slice.
    ///
    /// # Errors
    ///
    /// An error may be returned if:
    ///
    /// * The depth exceeds the maximum depth [`u8::MAX`].
    /// * A `HmacSha512` can't be constructed - this should be impossible.
    /// * Deriving a child key fails. Check the documentation for your [`PrivateKey`].
    pub fn derive_child(&self, index: &DerivationIndex) -> Result<Self> {
        let depth = self.depth.checked_add(1).ok_or(Error::Depth)?;

        let mut hmac =
            HmacSha512::new_from_slice(&self.chain_code).map_err(Error::HmacInvalidLength)?;
        if index.is_hardened() {
            hmac.update(&[0]);
            hmac.update(&self.private_key.to_bytes());
        } else if !K::requires_hardened_derivation() {
            hmac.update(&self.private_key.public_key().to_bytes());
        } else {
            return Err(Error::HardenedDerivationRequired);
        }
        hmac.update(&index.to_bytes());
        let result = hmac.finalize().into_bytes();
        let (private_key, chain_code) = result.split_at(KEY_SIZE / 8);

        let private_key = self
            .private_key
            .derive_child(
                &private_key
                    .try_into()
                    .expect("Invalid length for private key"),
            )
            .map_err(|_| Error::Derivation)?;

        Ok(Self {
            private_key,
            depth,
            chain_code: chain_code
                .try_into()
                .expect("Invalid length for chain code"),
        })
    }
}