authnz-common 0.2.1

//! Signature module.

use crate::{MResult, ServerError};

#[derive(Clone)]
/// Signing key (certificate).
pub struct SignKeypair {
  inner: InnerSignKeypair,
}

#[non_exhaustive]
#[derive(Clone)]
enum InnerSignKeypair {
  #[cfg(feature = "ed25519-utils")]
  Ed25519(Box<ed25519_dalek::SigningKey>),
  #[cfg(feature = "pqc-utils")]
  Dilithium5(Box<authnz_pqc_dilithium::Keypair>),
}

impl InnerSignKeypair {
  fn pack(&self) -> String {
    use base64::{Engine as _, engine::general_purpose::STANDARD};

    #[allow(unreachable_patterns)]
    match self {
      #[cfg(feature = "ed25519-utils")]
      Self::Ed25519(keypair) => STANDARD.encode(keypair.to_keypair_bytes()),
      #[cfg(feature = "pqc-utils")]
      Self::Dilithium5(keypair) => format!(
        "{}::||::{}",
        STANDARD.encode(keypair.public),
        STANDARD.encode(keypair.expose_secret())
      ),
      _ => unreachable!(),
    }
  }

  #[allow(clippy::unwrap_used)]
  fn unpack(cert: impl AsRef<str>) -> MResult<Self> {
    use base64::{Engine as _, engine::general_purpose::STANDARD};

    #[cfg(feature = "pqc-utils")]
    if cert.as_ref().contains("::||::") {
      let parts = cert.as_ref().split("::||::").collect::<Vec<_>>();
      if parts.len() != 2 {
        return Err(
          ServerError::from_private_str("Certificate have no part divider, this is not the certificate!")
            .with_private_str("Invalid certificate!")
            .with_500(),
        );
      }

      let public = parts.first().unwrap();
      let public = STANDARD.decode(public).map_err(|e| {
        ServerError::from_private(e)
          .with_private_str("Can't decode public certificate part!")
          .with_500()
      })?;

      let private = parts.last().unwrap();
      let private = STANDARD.decode(private).map_err(|e| {
        ServerError::from_private(e)
          .with_private_str("Can't decode private certificate part!")
          .with_500()
      })?;

      return Ok(Self::Dilithium5(Box::new(
        authnz_pqc_dilithium::Keypair::restore(&public, &private).map_err(|e| {
          ServerError::from_private(e)
            .with_private_str("Can't restore Dilithium5 keypair!")
            .with_500()
        })?,
      )));
    }

    #[cfg(feature = "ed25519-utils")]
    {
      let keypair = STANDARD.decode(cert.as_ref()).map_err(|e| {
        ServerError::from_private(e)
          .with_private_str("Can't decode Ed25519 certificate!")
          .with_500()
      })?;
      return Ok(Self::Ed25519(Box::new(
        ed25519_dalek::SigningKey::from_keypair_bytes(keypair.as_slice().try_into().map_err(|e| {
          ServerError::from_private(e)
            .with_private_str("Incorrect keypair length!")
            .with_500()
        })?)
        .map_err(|e| {
          ServerError::from_private(e)
            .with_private_str("Incorrect Ed25519 certificate!")
            .with_500()
        })?,
      )));
    }

    #[allow(unreachable_code)]
    Err(ServerError::from_private_str("Enable at least one of `ed25519-utils`, `pqc-utils` to use SignKeypair!").with_500())
  }

  fn public(&self) -> Vec<u8> {
    #[cfg(feature = "pqc-utils")]
    #[allow(irrefutable_let_patterns)]
    if let Self::Dilithium5(keypair) = &self {
      return keypair.public.to_vec();
    }

    #[cfg(feature = "ed25519-utils")]
    #[allow(irrefutable_let_patterns)]
    if let Self::Ed25519(keypair) = &self {
      return keypair.verifying_key().as_bytes().to_vec();
    }

    unreachable!()
  }

  fn private(&self) -> Vec<u8> {
    #[cfg(feature = "pqc-utils")]
    #[allow(irrefutable_let_patterns)]
    if let Self::Dilithium5(keypair) = &self {
      return keypair.expose_secret().to_vec();
    }

    #[cfg(feature = "ed25519-utils")]
    #[allow(irrefutable_let_patterns)]
    if let Self::Ed25519(keypair) = &self {
      return keypair.as_bytes().to_vec();
    }

    unreachable!()
  }
}

impl SignKeypair {
  #[cfg(feature = "ed25519-utils")]
  /// Generates Ed25519 keypair for signing usage.
  ///
  /// Store your private key safely!
  pub fn new_ed25519() -> MResult<Self> {
    use rand::TryRngCore;

    let mut csprng = rand::rngs::OsRng;

    let mut secret = ed25519_dalek::SecretKey::default();
    csprng.try_fill_bytes(&mut secret).map_err(|e| {
      ServerError::from_private(e)
        .with_private_str("Can't generate Ed25519 certificate!")
        .with_500()
    })?;

    Ok(Self {
      inner: InnerSignKeypair::Ed25519(Box::new(ed25519_dalek::SigningKey::from_bytes(&secret))),
    })
  }

  #[cfg(feature = "pqc-utils")]
  /// Generates Dilithium (mode 5) keypair for signing usage.
  ///
  /// Store your private key safely!
  pub fn new_dilithium5() -> Self {
    Self {
      inner: InnerSignKeypair::Dilithium5(Box::new(authnz_pqc_dilithium::Keypair::generate())),
    }
  }

  /// Packs keypair to string.
  pub fn pack_keypair(&self) -> String {
    self.inner.pack()
  }

  /// Unpacks keypair from a string.
  pub fn unpack_keypair(keypair: impl AsRef<str>) -> MResult<Self> {
    Ok(Self {
      inner: InnerSignKeypair::unpack(keypair)?,
    })
  }

  /// Signs raw bytes.
  pub fn sign_raw(&self, data: &[u8]) -> Vec<u8> {
    #[allow(unreachable_patterns)]
    match &self.inner {
      #[cfg(feature = "ed25519-utils")]
      InnerSignKeypair::Ed25519(keypair) => ed25519_dalek::Signer::sign(keypair.as_ref(), data).to_vec(),
      #[cfg(feature = "pqc-utils")]
      InnerSignKeypair::Dilithium5(keypair) => keypair.sign(data).to_vec(),
      _ => unreachable!(),
    }
  }

  /// Verifies data by its signature and provided public key.
  pub fn verify_raw(data: &[u8], sign: &[u8], public_key: &[u8]) -> MResult<()> {
    #[cfg(feature = "pqc-utils")]
    if authnz_pqc_dilithium::verify(sign, data, public_key).is_ok() {
      return Ok(());
    }
    #[cfg(feature = "ed25519-utils")]
    if let Ok(pkey) = public_key.try_into()
      && let Ok(vkey) = ed25519_dalek::VerifyingKey::from_bytes(pkey)
      && let Ok(sign) = ed25519_dalek::Signature::from_slice(sign)
      && ed25519_dalek::Verifier::verify(&vkey, data, &sign).is_ok()
    {
      return Ok(());
    }

    Err(ServerError::from_public("Invalid signature!").with_401())
  }

  /// Packs `T` into MessagePack and signs its bytes.
  pub fn sign<T: serde::Serialize>(&self, data: &T) -> MResult<Vec<u8>> {
    rmp_serde::to_vec(data)
      .map_err(|e| {
        ServerError::from_private(e)
          .with_private_str("Can't serialize data to sign!")
          .with_500()
      })
      .map(|v| self.sign_raw(&v))
  }

  /// Verifies any serializable `T` by its MessagePack view's signature and provided public key.
  pub fn verify<T: serde::Serialize>(data: &T, sign: &[u8], public_key: &[u8]) -> MResult<()> {
    let data = rmp_serde::to_vec(data).map_err(|e| {
      ServerError::from_private(e)
        .with_private_str("Can't serialize data to validate signature!")
        .with_500()
    })?;
    Self::verify_raw(&data, sign, public_key)
  }

  /// Verifies header and payload by their signature and provided public key.
  pub fn verify_token(header: &[u8], payload: &[u8], sign: &[u8], public_key: &[u8]) -> MResult<()> {
    let mut data = header.to_vec();
    data.extend_from_slice(payload);
    Self::verify_raw(&data, sign, public_key)
  }

  /// Get public key.
  pub fn public(&self) -> Vec<u8> {
    self.inner.public()
  }

  /// Get private key.
  ///
  /// # Safety
  ///
  /// This method is principally safe, but using `unsafe` keyword to describe the soundness of it's usage.
  pub unsafe fn private(&self) -> Vec<u8> {
    self.inner.private()
  }
}