koibumi-core 0.0.9

//! Provides methods to encrypt/decrypt Bitmessage messages.

use std::{
    convert::TryInto,
    fmt,
    io::{self, Read, Write},
    mem::size_of,
};

use aes::cipher::block_padding::Pkcs7;
use cbc::cipher::{
    inout::{block_padding::UnpadError, PadError},
    BlockDecryptMut, BlockEncryptMut, KeyIvInit,
};
use hmac::{digest, Mac};
use rand::Rng;

use koibumi_secp256k1 as secp256k1;

use crate::{
    error::TooLongError,
    hash::sha256,
    io::{LenBm, ReadFrom, SizedReadFrom, WriteTo},
    message,
    var_type::VarInt,
};

pub use crate::error::TooShortError;

const KEY_LEN: usize = 32;

/// This error indicates that the operation on a private key failed.
#[derive(Clone, Debug)]
pub enum PrivateKeyError {
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
}

impl fmt::Display for PrivateKeyError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Secp256k1Error(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for PrivateKeyError {}

impl From<secp256k1::Error> for PrivateKeyError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

/// A private key.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct PrivateKey {
    bytes: [u8; KEY_LEN],
}

impl PrivateKey {
    /// Constructs a private key from a byte array.
    pub fn new(bytes: [u8; KEY_LEN]) -> Result<Self, PrivateKeyError> {
        let _sk = secp256k1::SecretKey::from_slice(bytes.as_ref())?;
        Ok(Self { bytes })
    }

    /// Returns the public key derived from this private key.
    pub fn public_key(&self) -> PublicKey {
        let sk = secp256k1::SecretKey::from_slice(self.bytes.as_ref()).unwrap();
        let secp = secp256k1::Secp256k1::new();
        let pk = secp256k1::PublicKey::from_secret_key(&secp, &sk);
        let bytes = pk.serialize_uncompressed();
        let mut x = [0; KEY_LEN];
        let mut y = [0; KEY_LEN];
        x.copy_from_slice(&bytes[1..KEY_LEN + 1]);
        y.copy_from_slice(&bytes[KEY_LEN + 1..]);
        PublicKey::new(x, y).unwrap()
    }
}

impl AsRef<[u8]> for PrivateKey {
    fn as_ref(&self) -> &[u8] {
        &self.bytes
    }
}

/// This error indicates that the operation on a public key failed.
#[derive(Clone, Debug)]
pub enum PublicKeyError {
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
}

impl fmt::Display for PublicKeyError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Secp256k1Error(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for PublicKeyError {}

impl From<secp256k1::Error> for PublicKeyError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

/// A public key.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct PublicKey {
    x: [u8; KEY_LEN],
    y: [u8; KEY_LEN],
}

impl PublicKey {
    /// Constructs a public key from its components.
    pub fn new(x: [u8; KEY_LEN], y: [u8; KEY_LEN]) -> Result<Self, PublicKeyError> {
        let public_key = Self { x, y };
        let _pk = secp256k1::PublicKey::from_slice(&public_key.encode())?;
        Ok(public_key)
    }

    const PREFIX: u8 = 0x04;

    /// Encodes the public key.
    pub fn encode(&self) -> Vec<u8> {
        let mut encoded = Vec::with_capacity(1 + KEY_LEN + KEY_LEN);
        encoded.push(Self::PREFIX);
        encoded.extend_from_slice(&self.x);
        encoded.extend_from_slice(&self.y);
        encoded
    }
}

impl WriteTo for PublicKey {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        self.x.write_to(w)?;
        self.y.write_to(w)?;
        Ok(())
    }
}

impl ReadFrom for PublicKey {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let x = <[u8; KEY_LEN]>::read_from(r)?;
        let y = <[u8; KEY_LEN]>::read_from(r)?;
        Ok(Self { x, y })
    }
}

/// A pair of a public key and a private key.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct KeyPair {
    public_key: PublicKey,
    private_key: PrivateKey,
}

impl KeyPair {
    /// Returns the public key.
    pub fn public_key(&self) -> &PublicKey {
        &self.public_key
    }

    /// Returns the private key.
    pub fn private_key(&self) -> &PrivateKey {
        &self.private_key
    }

    /// Generates a random key pair.
    pub fn generate() -> Self {
        let secp = secp256k1::Secp256k1::signing_only();
        let mut rng = rand::rngs::OsRng::default();
        let (sk, pk) = secp.generate_keypair(&mut rng);

        let mut bytes = [0; KEY_LEN];
        bytes.copy_from_slice(&sk[..]);
        let private_key = PrivateKey::new(bytes).unwrap();
        let bytes = pk.serialize_uncompressed();
        let mut x = [0; KEY_LEN];
        let mut y = [0; KEY_LEN];
        x.copy_from_slice(&bytes[1..KEY_LEN + 1]);
        y.copy_from_slice(&bytes[KEY_LEN + 1..]);
        let public_key = PublicKey::new(x, y).unwrap();
        Self {
            public_key,
            private_key,
        }
    }
}

const CURVE_TYPE: u16 = 0x02ca;
type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

/// This error indicates that the encryption failed.
#[derive(Clone, Debug)]
pub enum EncryptError {
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
    /// An error was caught during initialization of the MAC generator.
    /// The actual error caught is returned as a payload of this variant.
    InvalidMacKeyIvLength(digest::InvalidLength),
    /// Padding error.
    PadError(PadError),
}

impl fmt::Display for EncryptError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Secp256k1Error(err) => err.fmt(f),
            Self::InvalidMacKeyIvLength(err) => err.fmt(f),
            Self::PadError(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for EncryptError {}

impl From<secp256k1::Error> for EncryptError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

impl From<digest::InvalidLength> for EncryptError {
    fn from(err: digest::InvalidLength) -> Self {
        Self::InvalidMacKeyIvLength(err)
    }
}

impl From<PadError> for EncryptError {
    fn from(err: PadError) -> Self {
        Self::PadError(err)
    }
}

/// This error indicates
/// that the decryption failed.
#[derive(Clone, Debug)]
pub enum DecryptError {
    /// Indicates that the object was encrypted by using a curve unsupported by this implementation.
    UnsupportedCurve(u16),
    /// Indicates that at least one of the lengths of the components of the public key was invalid.
    /// The actual lengths of the components are returned as payloads of this variant.
    InvalidKeyLength {
        /// The actual length of the component x.
        x_len: usize,
        /// The actual length of the component y.
        y_len: usize,
    },
    /// An error was caught during initialization of the MAC verifier.
    /// The actual error caught is returned as a payload of this variant.
    InvalidMacKeyIvLength(digest::InvalidLength),
    /// Indicates that the secret key is invalid.
    InvalidSecretKey,
    /// An error was caught during verification of the MAC.
    /// The actual error caught is returned as a payload of this variant.
    MacError(digest::MacError),
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
    /// Unpadding error.
    UnpadError(UnpadError),
}

impl fmt::Display for DecryptError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnsupportedCurve(curve) => write!(f, "unsupported curve: {:#04x}", curve),
            Self::InvalidKeyLength { x_len, y_len } => {
                write!(f, "invalid key length x: {}, y: {}", x_len, y_len)
            }
            Self::InvalidMacKeyIvLength(err) => err.fmt(f),
            Self::InvalidSecretKey => "invalid secret key".fmt(f),
            Self::MacError(err) => err.fmt(f),
            Self::Secp256k1Error(err) => err.fmt(f),
            Self::UnpadError(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for DecryptError {}

impl From<digest::InvalidLength> for DecryptError {
    fn from(err: digest::InvalidLength) -> Self {
        Self::InvalidMacKeyIvLength(err)
    }
}

impl From<digest::MacError> for DecryptError {
    fn from(err: digest::MacError) -> Self {
        Self::MacError(err)
    }
}

impl From<secp256k1::Error> for DecryptError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

impl From<UnpadError> for DecryptError {
    fn from(err: UnpadError) -> Self {
        Self::UnpadError(err)
    }
}

const IV_LEN: usize = 16;
const MAC_LEN: usize = 32;

/// A encrypted message.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Encrypted {
    iv: [u8; IV_LEN],
    curve_type: u16,
    //x_length: u16,
    x: Vec<u8>,
    //y_length: u16,
    y: Vec<u8>,
    encrypted: Vec<u8>,
    mac: [u8; MAC_LEN],
}

impl Encrypted {
    /// Encrypts data and returns a encrypted object.
    pub fn encrypt(data: impl AsRef<[u8]>, public_key: &PublicKey) -> Result<Self, EncryptError> {
        let secp = secp256k1::Secp256k1::signing_only();
        let mut rng = rand::rngs::OsRng::default();
        let (sk, pk) = secp.generate_keypair(&mut rng);

        let bytes = pk.serialize_uncompressed();
        let mut x = [0; KEY_LEN];
        let mut y = [0; KEY_LEN];
        x.copy_from_slice(&bytes[1..KEY_LEN + 1]);
        y.copy_from_slice(&bytes[KEY_LEN + 1..]);

        let public_key = secp256k1::PublicKey::from_slice(&public_key.encode())?;

        let shared_secret =
            secp256k1::ecdh::SharedSecret::new_with_hash(&public_key, &sk, |x, _y| {
                secp256k1::ecdh::SharedSecret::from_bytes(x)
            });

        let h = crate::hash::sha512(shared_secret);
        let encryption_key = &h[..32];
        let iv: [u8; IV_LEN] = rng.gen();
        let pt_len = data.as_ref().len();
        let buf_len = (pt_len % 32) * 32;
        let buf_len = if pt_len > buf_len {
            buf_len + 32
        } else {
            buf_len
        };
        let mut buf = vec![0u8; buf_len];
        let encrypted = Aes256CbcEnc::new(encryption_key.into(), iv.as_ref().into())
            .encrypt_padded_b2b_mut::<Pkcs7>(data.as_ref(), &mut buf)?;

        let mac_key = &h[32..];
        let mut mac = hmac::Hmac::<sha2::Sha256>::new_from_slice(mac_key)?;
        mac.update(&iv);
        mac.update(&CURVE_TYPE.to_be_bytes());
        mac.update(&(x.len() as u16).to_be_bytes());
        mac.update(&x);
        mac.update(&(y.len() as u16).to_be_bytes());
        mac.update(&y);
        mac.update(&encrypted);
        let mac_bytes = mac.finalize().into_bytes();

        Ok(Self {
            iv,
            curve_type: CURVE_TYPE,
            x: x.as_ref().try_into().unwrap(),
            y: y.as_ref().try_into().unwrap(),
            encrypted: Vec::from(encrypted),
            mac: mac_bytes[..].try_into().unwrap(),
        })
    }

    /// Decrypts this encrypted object and returns the decrypted bytes.
    pub fn decrypt(&self, private_key: &PrivateKey) -> Result<Vec<u8>, DecryptError> {
        if self.curve_type != CURVE_TYPE {
            return Err(DecryptError::UnsupportedCurve(self.curve_type));
        }
        if self.x.len() != KEY_LEN || self.y.len() != KEY_LEN {
            return Err(DecryptError::InvalidKeyLength {
                x_len: self.x.len(),
                y_len: self.y.len(),
            });
        }
        let mut public_key_bytes = [0; 65];
        public_key_bytes[0] = 0x04;
        public_key_bytes[1..33].copy_from_slice(&self.x);
        public_key_bytes[33..].copy_from_slice(&self.y);
        let public_key = secp256k1::PublicKey::from_slice(&public_key_bytes)?;
        let secret_key = secp256k1::SecretKey::from_slice(private_key.as_ref())?;
        let shared_secret =
            secp256k1::ecdh::SharedSecret::new_with_hash(&public_key, &secret_key, |x, _y| {
                secp256k1::ecdh::SharedSecret::from_bytes(x)
            });

        let h = crate::hash::sha512(shared_secret);
        let mac_key = &h[32..];

        let mut mac = hmac::Hmac::<sha2::Sha256>::new_from_slice(mac_key)?;
        mac.update(&self.iv);
        mac.update(&self.curve_type.to_be_bytes());
        mac.update(&(self.x.len() as u16).to_be_bytes());
        mac.update(&self.x);
        mac.update(&(self.y.len() as u16).to_be_bytes());
        mac.update(&self.y);
        mac.update(&self.encrypted);
        mac.verify(&self.mac.into())?;

        let encryption_key = &h[..32];
        let mut buf = vec![0u8; self.encrypted.len()];
        let decrypted = Aes256CbcDec::new(encryption_key.into(), self.iv.as_ref().into())
            .decrypt_padded_b2b_mut::<Pkcs7>(&self.encrypted, &mut buf)?;
        Ok(Vec::from(decrypted))
    }
}

impl SizedReadFrom for Encrypted {
    fn sized_read_from(r: &mut dyn Read, len: usize) -> io::Result<Self>
    where
        Self: Sized,
    {
        let mut iv = [0; IV_LEN];
        r.read_exact(&mut iv)?;

        let curve_type = u16::read_from(r)?;

        let x_length = u16::read_from(r)?;

        let mut rt = r.take(x_length as u64);
        let mut x = Vec::with_capacity(x_length as usize);
        rt.read_to_end(&mut x)?;

        let y_length = u16::read_from(r)?;

        let mut rt = r.take(y_length as u64);
        let mut y = Vec::with_capacity(y_length as usize);
        rt.read_to_end(&mut y)?;

        let sub_len = IV_LEN
            + size_of::<u16>()
            + size_of::<u16>()
            + x_length as usize
            + size_of::<u16>()
            + y_length as usize
            + MAC_LEN;
        if sub_len > len {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                TooShortError::new(sub_len, len),
            ));
        }
        let encrypted_len = len - sub_len;
        let mut rt = r.take(encrypted_len as u64);
        let mut encrypted = Vec::with_capacity(encrypted_len);
        rt.read_to_end(&mut encrypted)?;

        let mut mac = [0; MAC_LEN];
        r.read_exact(&mut mac)?;

        Ok(Self {
            iv,
            curve_type,
            x,
            y,
            encrypted,
            mac,
        })
    }
}

impl WriteTo for Encrypted {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        self.iv.write_to(w)?;
        self.curve_type.write_to(w)?;
        (self.x.len() as u16).write_to(w)?;
        self.x.write_to(w)?;
        (self.y.len() as u16).write_to(w)?;
        self.y.write_to(w)?;
        self.encrypted.write_to(w)?;
        self.mac.write_to(w)?;
        Ok(())
    }
}

/// This error indicates
/// that the signing failed.
#[derive(Clone, Debug)]
pub enum SignError {
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
}

impl fmt::Display for SignError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Secp256k1Error(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for SignError {}

impl From<secp256k1::Error> for SignError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

/// This error indicates
/// that the verification failed.
#[derive(Clone, Debug)]
pub enum VerifyError {
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
}

impl fmt::Display for VerifyError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Secp256k1Error(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for VerifyError {}

impl From<secp256k1::Error> for VerifyError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

/// A signature.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Signature {
    bytes: Vec<u8>,
}

fn sign_sha256(data: impl AsRef<[u8]>, private_key: &PrivateKey) -> Result<Signature, SignError> {
    let hash = sha256(data);
    let message = secp256k1::Message::from_digest_slice(hash.as_ref())?;
    let secret_key = secp256k1::SecretKey::from_slice(private_key.as_ref())?;
    let secp = secp256k1::Secp256k1::signing_only();
    let signature = secp.sign_ecdsa(&message, &secret_key);
    let bytes = Vec::from(signature.serialize_der().as_ref());
    Ok(Signature { bytes })
}

/// Signs on data using a private key.
pub fn sign(data: impl AsRef<[u8]>, private_key: &PrivateKey) -> Result<Signature, SignError> {
    sign_sha256(data, private_key)
}

/// Verifies a signature of data using a public key.
pub fn verify(
    data: impl AsRef<[u8]>,
    signature: &Signature,
    public_key: &PublicKey,
) -> Result<(), VerifyError> {
    let hash = sha256(data);
    let message = secp256k1::Message::from_digest_slice(hash.as_ref())?;
    let mut sig = secp256k1::ecdsa::Signature::from_der(&signature.bytes)?;
    sig.normalize_s();
    let mut public_key_bytes = [0; 65];
    public_key_bytes[0] = 0x04;
    public_key_bytes[1..33].copy_from_slice(&public_key.x);
    public_key_bytes[33..].copy_from_slice(&public_key.y);
    let pk = secp256k1::PublicKey::from_slice(&public_key_bytes)?;
    let secp = secp256k1::Secp256k1::verification_only();
    secp.verify_ecdsa(&message, &sig, &pk)?;
    Ok(())
}

impl WriteTo for Signature {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        let len: VarInt = self.bytes.len().into();
        len.write_to(w)?;
        self.bytes.write_to(w)?;
        Ok(())
    }
}

impl ReadFrom for Signature {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let len = VarInt::read_from(r)?;
        if len.as_u64() > message::Object::MAX_OBJECT_PAYLOAD_LENGTH as u64 {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                TooLongError::new(
                    message::Object::MAX_OBJECT_PAYLOAD_LENGTH,
                    len.as_u64() as usize,
                ),
            ));
        }
        let bytes = Vec::<u8>::sized_read_from(r, len.as_u64() as usize)?;
        Ok(Self { bytes })
    }
}

impl LenBm for Signature {
    fn len_bm(&self) -> usize {
        let len: VarInt = self.bytes.len().into();
        len.len_bm() + self.bytes.len()
    }
}