koibumi-core 0.0.9

//! Functions for validating or performing the proof of work
//! used by the Bitmessage protocol.

use std::{
    convert::TryInto,
    fmt,
    io::{self, Read, Write},
    mem::size_of,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

use serde::{Deserialize, Serialize};
use sha2::Digest;

use crate::{
    config::Config,
    hash::{sha512, Hash512},
    io::{LenBm, ReadFrom, WriteTo},
    time::Time,
    var_type::VarInt,
};

/// Target or trial value of a proof of work.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct Value(u64);

impl fmt::Display for Value {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:016x}", self.0)
    }
}

impl From<u64> for Value {
    fn from(value: u64) -> Self {
        Self(value)
    }
}

impl Value {
    /// Constructs a value from a primitive.
    pub fn new(value: u64) -> Self {
        Self(value)
    }

    /// Returns the value as `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }
}

/// The nonce of a proof of work.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct Nonce(u64);

impl fmt::Display for Nonce {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:016x}", self.0)
    }
}

impl From<u64> for Nonce {
    fn from(value: u64) -> Self {
        Self(value)
    }
}

impl Nonce {
    /// Constructs a nonce from a value.
    pub fn new(value: u64) -> Self {
        Self(value)
    }

    /// Returns the value as `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }
}

impl WriteTo for Nonce {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        self.0.write_to(w)
    }
}

impl ReadFrom for Nonce {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        Ok(Self(u64::read_from(r)?))
    }
}

impl LenBm for Nonce {
    fn len_bm(&self) -> usize {
        self.0.len_bm()
    }
}

/// One of parameters of the proof of work used by the Bitmessage protocol.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Serialize, Deserialize)]
pub struct NonceTrialsPerByte(u64);

impl fmt::Display for NonceTrialsPerByte {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl From<u64> for NonceTrialsPerByte {
    fn from(value: u64) -> Self {
        Self(value)
    }
}

impl NonceTrialsPerByte {
    /// Constructs a nonce trials per byte from a value.
    pub fn new(value: u64) -> Self {
        Self(value)
    }

    /// Returns the value as `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }
}

impl WriteTo for NonceTrialsPerByte {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        let v: VarInt = self.0.into();
        v.write_to(w)
    }
}

impl ReadFrom for NonceTrialsPerByte {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let v = VarInt::read_from(r)?;
        Ok(Self(v.as_u64()))
    }
}

/// One of parameters of the proof of work used by the Bitmessage protocol.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Serialize, Deserialize)]
pub struct PayloadLengthExtraBytes(u64);

impl fmt::Display for PayloadLengthExtraBytes {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl From<u64> for PayloadLengthExtraBytes {
    fn from(value: u64) -> Self {
        Self(value)
    }
}

impl PayloadLengthExtraBytes {
    /// Constructs a payload length extra bytes from a value.
    pub fn new(value: u64) -> Self {
        Self(value)
    }

    /// Returns the value as `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }
}

impl WriteTo for PayloadLengthExtraBytes {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        let v: VarInt = self.0.into();
        v.write_to(w)
    }
}

impl ReadFrom for PayloadLengthExtraBytes {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let v = VarInt::read_from(r)?;
        Ok(Self(v.as_u64()))
    }
}

/// Returns the initial hash of a payload.
pub fn initial_hash(payload: impl AsRef<[u8]>) -> Hash512 {
    sha512(payload)
}

const MIN_TTL: u64 = 300;

fn effective_ttl(expires_time: Time, now: Time) -> u64 {
    if expires_time < now {
        MIN_TTL
    } else {
        let ttl = expires_time.as_secs() - now.as_secs();
        u64::max(ttl, MIN_TTL)
    }
}

fn effective_nonce_trials_per_byte(
    config: &Config,
    nonce_trials_per_byte: NonceTrialsPerByte,
) -> NonceTrialsPerByte {
    NonceTrialsPerByte::max(nonce_trials_per_byte, config.nonce_trials_per_byte())
}

fn effective_payload_length_extra_bytes(
    config: &Config,
    payload_length_extra_bytes: PayloadLengthExtraBytes,
) -> PayloadLengthExtraBytes {
    PayloadLengthExtraBytes::max(
        payload_length_extra_bytes,
        config.payload_length_extra_bytes(),
    )
}

fn target_by_ttl(
    payload_length: u64,
    ttl: u64,
    nonce_trials_per_byte: NonceTrialsPerByte,
    payload_length_extra_bytes: PayloadLengthExtraBytes,
) -> Value {
    let length = payload_length + size_of::<u64>() as u64 + payload_length_extra_bytes.as_u64();
    let denominator = nonce_trials_per_byte.as_u64() * (length + (ttl * length / 0x10000));
    const NUMERATOR: u128 = 0x1_0000_0000_0000_0000;
    Value::new((NUMERATOR / denominator as u128) as u64)
}

/// Returns the target of a proof of work.
pub fn target(
    config: &Config,
    payload_len: usize,
    expires_time: Time,
    now: Time,
    nonce_trials_per_byte: NonceTrialsPerByte,
    payload_length_extra_bytes: PayloadLengthExtraBytes,
) -> Value {
    target_by_ttl(
        payload_len as u64,
        effective_ttl(expires_time, now),
        effective_nonce_trials_per_byte(config, nonce_trials_per_byte),
        effective_payload_length_extra_bytes(config, payload_length_extra_bytes),
    )
}

/// An error which can be returne when validating a proof of work.
///
/// This error is used as the error type for the [`validate()`] method.
///
/// [`validate()`]: fn.validate.html
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum ValidateError {
    /// The proof of work was insufficient.
    /// The target value and the actual value are returned
    /// as payloads of this variant.
    Insufficient {
        /// The target value.
        target: Value,
        /// The actual value.
        value: Value,
    },
}

impl fmt::Display for ValidateError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Insufficient { target, value } => write!(
                f,
                "value {:016x} is insufficient for target {:016x}",
                value.as_u64(),
                target.as_u64()
            ),
        }
    }
}

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

/// Validates a proof of work used by the Bitmessage protocol.
pub fn validate(initial_hash: Hash512, target: Value, nonce: Nonce) -> Result<(), ValidateError> {
    let mut hasher = sha2::Sha512::new();
    hasher.update(&nonce.as_u64().to_be_bytes());
    hasher.update(&initial_hash);
    let hash1 = hasher.finalize_reset();
    hasher.update(hash1);

    let trial_value = u64::from_be_bytes(hasher.finalize()[0..8].try_into().unwrap());
    if trial_value <= target.as_u64() {
        Ok(())
    } else {
        Err(ValidateError::Insufficient {
            target,
            value: Value::new(trial_value),
        })
    }
}

/// An error which can be returne when performing a proof of work.
///
/// This error is used as the error type for the [`perform()`] method.
///
/// [`perform()`]: fn.perform.html
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum PerformError {
    /// The proof of work was aborted.
    Canceled,
}

impl fmt::Display for PerformError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Canceled => "canceled".fmt(f),
        }
    }
}

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

/// Performs a proof of work used by the Bitmessage protocol.
pub fn perform(
    initial_hash: Hash512,
    target: Value,
    start: Nonce,
    step: usize,
    cancel: Arc<AtomicBool>,
) -> Result<Nonce, PerformError> {
    let target = target.as_u64();
    let mut nonce: u64 = start.as_u64();
    let mut hasher = sha2::Sha512::new();
    while !cancel.load(Ordering::SeqCst) {
        for _ in 0..0x10000 {
            hasher.update(&nonce.to_be_bytes());
            hasher.update(&initial_hash);
            let hash1 = hasher.finalize_reset();
            hasher.update(hash1);

            let trial_value = u64::from_be_bytes(hasher.finalize_reset()[0..8].try_into().unwrap());
            if trial_value <= target {
                return Ok(nonce.into());
            }
            nonce = nonce.wrapping_add(step as u64);
        }
    }
    Err(PerformError::Canceled)
}