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

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

#[derive(Debug, thiserror::Error)]
pub enum DerivationError {
    #[error("algorithm not supported")]
    Algorithm,

    #[error("Unable to create seed from mnemonic: {0}")]
    Mnemonic(#[from] MnemonicGenerationError),

    #[error("{0}")]
    ExtendedPrivateKey(#[from] XPrvError),
}

pub type Result<T, E = DerivationError> = std::result::Result<T, E>;

#[derive(Serialize, Deserialize, Clone, Debug)]
pub enum DerivationAlgorithm {
    Ed25519,
    Secp256k1,
}

impl DerivationAlgorithm {
    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),
        })
    }
}

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

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

    pub fn path(&self) -> &DerivationPath {
        &self.path
    }

    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)?)
    }

    pub fn derive_with_master_seed(&self, seed: Vec<u8>) -> Result<DerivationResponse> {
        self.algorithm.derive(seed, &self.path)
    }
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct DerivationResponse {
    pub algorithm: DerivationAlgorithm,
    pub data: Vec<u8>,
    pub chain_code: [u8; 32],
    pub depth: u8,
}

impl DerivationResponse {
    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(),
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum TryFromDerivationResponseError {
    #[error("incorrect algorithm provided")]
    Algorithm,

    #[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)
    }
}