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

// Because all algorithms make use of wildcard matching
#![allow(clippy::match_wildcard_for_single_variants)]

use crate::{
    extended_key::private_key::Error as XPrvError, DerivationPath, ExtendedPrivateKey, PrivateKey,
};
use keyfork_mnemonic_util::{Mnemonic, MnemonicGenerationError};
use serde::{Deserialize, Serialize};

/// An error encountered while deriving a key.
#[derive(Debug, thiserror::Error)]
pub enum DerivationError {
    /// The algorithm requested was not supported. This may occur when a feature adding support for
    /// an algorithm has not been enabled.
    #[error("Algorithm not supported")]
    Algorithm,

    /// A seed was unable to be created from the mnemonic.
    #[error("Unable to create seed from mnemonic: {0}")]
    Mnemonic(#[from] MnemonicGenerationError),

    /// Generating an [`ExtendedPrivateKey`] resulted in an error.
    #[error("{0}")]
    ExtendedPrivateKey(#[from] XPrvError),
}

#[allow(missing_docs)]
pub type Result<T, E = DerivationError> = std::result::Result<T, E>;

/// The algorithm to derive a key for. The choice of algorithm will result in a different resulting
/// derivation.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub enum DerivationAlgorithm {
    #[allow(missing_docs)]
    Ed25519,
    #[allow(missing_docs)]
    Secp256k1,
}

impl DerivationAlgorithm {
    /// Given a mnemonic seed and a derivation path, derive an [`ExtendedPrivateKey`].
    ///
    /// # Errors
    /// The method may error if the derivation fails or if the algorithm is not supported.
    pub fn derive(&self, seed: Vec<u8>, path: &DerivationPath) -> Result<DerivationResponse> {
        match self {
            #[cfg(feature = "ed25519")]
            Self::Ed25519 => {
                let key = ExtendedPrivateKey::<ed25519_dalek::SigningKey>::new(seed)?;
                let derived_key = key.derive_path(path)?;
                Ok(DerivationResponse::with_algo_and_xprv(
                    self.clone(),
                    &derived_key,
                ))
            }
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1 => {
                let key = ExtendedPrivateKey::<k256::SecretKey>::new(seed)?;
                let derived_key = key.derive_path(path)?;
                Ok(DerivationResponse::with_algo_and_xprv(
                    self.clone(),
                    &derived_key,
                ))
            }
            #[allow(unreachable_patterns)]
            _ => Err(DerivationError::Algorithm),
        }
    }
}

impl std::str::FromStr for DerivationAlgorithm {
    type Err = DerivationError;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        Ok(match s {
            "ed25519" => Self::Ed25519,
            "secp256k1" => Self::Secp256k1,
            _ => return Err(DerivationError::Algorithm),
        })
    }
}

/// A derivation request.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct DerivationRequest {
    algorithm: DerivationAlgorithm,
    path: DerivationPath,
}

impl DerivationRequest {
    /// Create a new derivation request.
    pub fn new(algorithm: DerivationAlgorithm, path: &DerivationPath) -> Self {
        Self {
            algorithm,
            path: path.clone(),
        }
    }

    /// Return the path of the derivation request.
    pub fn path(&self) -> &DerivationPath {
        &self.path
    }

    /// Derive an [`ExtendedPrivateKey`] using the seed from the given mnemonic.
    ///
    /// # Errors
    /// The method may error if the derivation fails or if the algorithm is not supported.
    pub fn derive_with_mnemonic(&self, mnemonic: &Mnemonic) -> Result<DerivationResponse> {
        // TODO: passphrase support and/or store passphrase within mnemonic
        self.derive_with_master_seed(mnemonic.seed(None)?)
    }

    /// Derive an [`ExtendedPrivateKey`] using the given seed.
    ///
    /// # Errors
    /// The method may error if the derivation fails or if the algorithm is not supported.
    pub fn derive_with_master_seed(&self, seed: Vec<u8>) -> Result<DerivationResponse> {
        self.algorithm.derive(seed, &self.path)
    }
}

/// A response to a [`DerivationRequest`]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct DerivationResponse {
    /// The algorithm used to derive the data.
    pub algorithm: DerivationAlgorithm,

    /// The derived private key.
    pub data: Vec<u8>,

    /// The chain code, used for further derivation.
    pub chain_code: [u8; 32],

    /// The depth, used for further derivation.
    pub depth: u8,
}

impl DerivationResponse {
    /// Create a [`DerivationResponse`] with the given values.
    pub fn with_algo_and_xprv<T: PrivateKey + Clone>(
        algorithm: DerivationAlgorithm,
        xprv: &ExtendedPrivateKey<T>,
    ) -> Self {
        Self {
            algorithm,
            data: PrivateKey::to_bytes(xprv.private_key()).to_vec(),
            chain_code: xprv.chain_code(),
            depth: xprv.depth(),
        }
    }
}

/// An error when creating a [`DerivationResponse`].
#[derive(Debug, thiserror::Error)]
pub enum TryFromDerivationResponseError {
    /// The algorithm used to derive the data does not match the algorithm of the
    /// [`ExtendedPrivateKey`] being created.
    #[error("incorrect algorithm provided")]
    Algorithm,

    /// An error occurred while creating an [`ExtendedPrivateKey`] from the given response.
    #[error("{0}")]
    ExtendedPrivateKey(#[from] XPrvError),
}

#[cfg(feature = "secp256k1")]
impl TryFrom<&DerivationResponse> for ExtendedPrivateKey<k256::SecretKey> {
    type Error = TryFromDerivationResponseError;

    fn try_from(value: &DerivationResponse) -> std::result::Result<Self, Self::Error> {
        match value.algorithm {
            DerivationAlgorithm::Secp256k1 => {
                Self::new_from_parts(&value.data, value.depth, value.chain_code).map_err(From::from)
            }
            _ => Err(Self::Error::Algorithm),
        }
    }
}

#[cfg(feature = "secp256k1")]
impl TryFrom<DerivationResponse> for ExtendedPrivateKey<k256::SecretKey> {
    type Error = TryFromDerivationResponseError;

    fn try_from(value: DerivationResponse) -> std::result::Result<Self, Self::Error> {
        ExtendedPrivateKey::<k256::SecretKey>::try_from(&value)
    }
}

#[cfg(feature = "ed25519")]
impl TryFrom<&DerivationResponse> for ExtendedPrivateKey<ed25519_dalek::SigningKey> {
    type Error = TryFromDerivationResponseError;

    fn try_from(value: &DerivationResponse) -> std::result::Result<Self, Self::Error> {
        match value.algorithm {
            DerivationAlgorithm::Ed25519 => {
                Self::new_from_parts(&value.data, value.depth, value.chain_code).map_err(From::from)
            }
            _ => Err(Self::Error::Algorithm),
        }
    }
}

#[cfg(feature = "ed25519")]
impl TryFrom<DerivationResponse> for ExtendedPrivateKey<ed25519_dalek::SigningKey> {
    type Error = TryFromDerivationResponseError;

    fn try_from(value: DerivationResponse) -> std::result::Result<Self, Self::Error> {
        ExtendedPrivateKey::<ed25519_dalek::SigningKey>::try_from(&value)
    }
}