serai/coins/monero/src/unreduced_scalar.rs

use core::cmp::Ordering;

use std_shims::{
  sync::OnceLock,
  io::{self, *},
};

use curve25519_dalek::scalar::Scalar;

use crate::serialize::*;

static PRECOMPUTED_SCALARS_CELL: OnceLock<[Scalar; 8]> = OnceLock::new();
/// Precomputed scalars used to recover an incorrectly reduced scalar.
#[allow(non_snake_case)]
pub(crate) fn PRECOMPUTED_SCALARS() -> [Scalar; 8] {
  *PRECOMPUTED_SCALARS_CELL.get_or_init(|| {
    let mut precomputed_scalars = [Scalar::ONE; 8];
    for (i, scalar) in precomputed_scalars.iter_mut().enumerate().skip(1) {
      *scalar = Scalar::from(((i * 2) + 1) as u8);
    }
    precomputed_scalars
  })
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct UnreducedScalar(pub [u8; 32]);

impl UnreducedScalar {
  pub fn write<W: Write>(&self, w: &mut W) -> io::Result<()> {
    w.write_all(&self.0)
  }

  pub fn read<R: Read>(r: &mut R) -> io::Result<UnreducedScalar> {
    Ok(UnreducedScalar(read_bytes(r)?))
  }

  pub fn as_bytes(&self) -> &[u8; 32] {
    &self.0
  }

  fn as_bits(&self) -> [u8; 256] {
    let mut bits = [0; 256];
    for (i, bit) in bits.iter_mut().enumerate() {
      *bit = core::hint::black_box(1 & (self.0[i / 8] >> (i % 8)))
    }

    bits
  }

  /// Computes the non-adjacent form of this scalar with width 5.
  ///
  /// This matches Monero's `slide` function and intentionally gives incorrect outputs under
  /// certain conditions in order to match Monero.
  ///
  /// This function does not execute in constant time.
  fn non_adjacent_form(&self) -> [i8; 256] {
    let bits = self.as_bits();
    let mut naf = [0i8; 256];
    for (b, bit) in bits.into_iter().enumerate() {
      naf[b] = bit as i8;
    }

    for i in 0 .. 256 {
      if naf[i] != 0 {
        // if the bit is a one, work our way up through the window
        // combining the bits with this bit.
        for b in 1 .. 6 {
          if (i + b) >= 256 {
            // if we are at the length of the array then break out
            // the loop.
            break;
          }
          // potential_carry - the value of the bit at i+b compared to the bit at i
          let potential_carry = naf[i + b] << b;

          if potential_carry != 0 {
            if (naf[i] + potential_carry) <= 15 {
              // if our current "bit" plus the potential carry is less than 16
              // add it to our current "bit" and set the potential carry bit to 0.
              naf[i] += potential_carry;
              naf[i + b] = 0;
            } else if (naf[i] - potential_carry) >= -15 {
              // else if our current "bit" minus the potential carry is more than -16
              // take it away from our current "bit".
              // we then work our way up through the bits setting ones to zero, when
              // we hit the first zero we change it to one then stop, this is to factor
              // in the minus.
              naf[i] -= potential_carry;
              #[allow(clippy::needless_range_loop)]
              for k in (i + b) .. 256 {
                if naf[k] == 0 {
                  naf[k] = 1;
                  break;
                }
                naf[k] = 0;
              }
            } else {
              break;
            }
          }
        }
      }
    }

    naf
  }

  /// Recover the scalar that an array of bytes was incorrectly interpreted as by Monero's `slide`
  /// function.
  ///
  /// In Borromean range proofs Monero was not checking that the scalars used were
  /// reduced. This lead to the scalar stored being interpreted as a different scalar,
  /// this function recovers that scalar.
  ///
  /// See: https://github.com/monero-project/monero/issues/8438
  pub fn recover_monero_slide_scalar(&self) -> Scalar {
    if self.0[31] & 128 == 0 {
      // Computing the w-NAF of a number can only give an output with 1 more bit than
      // the number, so even if the number isn't reduced, the `slide` function will be
      // correct when the last bit isn't set.
      return Scalar::from_bytes_mod_order(self.0);
    }

    let precomputed_scalars = PRECOMPUTED_SCALARS();

    let mut recovered = Scalar::ZERO;
    for &numb in self.non_adjacent_form().iter().rev() {
      recovered += recovered;
      match numb.cmp(&0) {
        Ordering::Greater => recovered += precomputed_scalars[(numb as usize) / 2],
        Ordering::Less => recovered -= precomputed_scalars[((-numb) as usize) / 2],
        Ordering::Equal => (),
      }
    }
    recovered
  }
}