use core::ops::Range;
use libcrux_hacl_rs::bignum::bignum64;
use super::{MAX_PRIVATE_MODULUS_BYTES, MAX_PRIVATE_MODULUS_WORDS, Mont, RsaPrivateKey};
use crate::Error;
use crate::low::{self, PosInt, ct};
use crate::mid::rng::RandomSource;
use crate::mid::rsa_key_size::RsaKeySize;
pub(super) fn generate_key(
size: RsaKeySize,
candidate_random: &mut dyn RandomSource,
checks_random: &mut dyn RandomSource,
) -> Result<RsaPrivateKey, Error> {
let len = size.prime_bits() / 64;
let sieve = Sieve::new();
let e = Mont::new(&[PUBLIC_EXPONENT.into()]);
loop {
let mut p = random_prime(size, &sieve, candidate_random, checks_random)?;
let mut q = random_prime(size, &sieve, candidate_random, checks_random)?;
let (p, q) = (&mut p.0[..len], &mut q.0[..len]);
if public(bignum64::eq_mask(len as u32, p, q)) != 0 {
return Err(Error::RngFailed);
}
let swap = bignum64::lt_mask(len as u32, p, q);
for (p, q) in p.iter_mut().zip(q.iter_mut()) {
let t = (*p ^ *q) & swap;
*p ^= t;
*q ^= t;
}
let mut k = Derivation::default();
if !k.derive(&e, p, q) {
continue;
}
return RsaPrivateKey::new(
public(PosInt::from_words(p)?),
public(PosInt::from_words(q)?),
public(PosInt::from_words(&k.d[..len * 2])?),
public(PosInt::from_words(&k.dp[..len])?),
public(PosInt::from_words(&k.dq[..len])?),
public(PosInt::from_words(&k.iqmp[..len])?),
public(PosInt::from_words(&k.n[..len * 2])?),
PUBLIC_EXPONENT.into(),
);
}
}
const PUBLIC_EXPONENT: u32 = 0x10001;
struct Derivation {
p_1: [u64; MAX_PRIVATE_MODULUS_WORDS],
q_1: [u64; MAX_PRIVATE_MODULUS_WORDS],
phi: [u64; MAX_PRIVATE_MODULUS_WORDS * 2],
p_2: [u64; MAX_PRIVATE_MODULUS_WORDS],
d: [u64; MAX_PRIVATE_MODULUS_WORDS * 2],
dp: [u64; MAX_PRIVATE_MODULUS_WORDS],
dq: [u64; MAX_PRIVATE_MODULUS_WORDS],
iqmp: [u64; MAX_PRIVATE_MODULUS_WORDS],
n: [u64; MAX_PRIVATE_MODULUS_WORDS * 2],
}
impl Derivation {
fn derive(&mut self, e: &Mont, p: &[u64], q: &[u64]) -> bool {
let len = p.len();
let w = len as u32;
let (p_1, q_1) = (&mut self.p_1[..len], &mut self.q_1[..len]);
p_1.copy_from_slice(p);
p_1[0] ^= 1;
q_1.copy_from_slice(q);
q_1[0] ^= 1;
let phi = &mut self.phi[..len * 2];
bignum64::mul(w, p_1, q_1, phi);
let ok = invert_e(e, phi, &mut self.d[..len * 2])
& invert_e(e, p_1, &mut self.dp[..len])
& invert_e(e, q_1, &mut self.dq[..len]);
if !ok {
return false;
}
let p_2 = &mut self.p_2[..len];
bignum64::sub(w, p, &ONE_TWO[1][..len], p_2);
let ctx = Mont::new_secret(p, len);
bignum64::mod_exp_consttime_precomp(&ctx, q, w * 64, p_2, &mut self.iqmp[..len]);
bignum64::mul(w, p, q, &mut self.n[..len * 2]);
true
}
}
impl Default for Derivation {
fn default() -> Self {
Self {
p_1: [0; MAX_PRIVATE_MODULUS_WORDS],
q_1: [0; MAX_PRIVATE_MODULUS_WORDS],
phi: [0; MAX_PRIVATE_MODULUS_WORDS * 2],
p_2: [0; MAX_PRIVATE_MODULUS_WORDS],
d: [0; MAX_PRIVATE_MODULUS_WORDS * 2],
dp: [0; MAX_PRIVATE_MODULUS_WORDS],
dq: [0; MAX_PRIVATE_MODULUS_WORDS],
iqmp: [0; MAX_PRIVATE_MODULUS_WORDS],
n: [0; MAX_PRIVATE_MODULUS_WORDS * 2],
}
}
}
impl Drop for Derivation {
fn drop(&mut self) {
low::zeroise(&mut self.p_1);
low::zeroise(&mut self.q_1);
low::zeroise(&mut self.phi);
low::zeroise(&mut self.p_2);
low::zeroise(&mut self.d);
low::zeroise(&mut self.dp);
low::zeroise(&mut self.dq);
low::zeroise(&mut self.iqmp);
low::zeroise(&mut self.n);
}
}
fn invert_e(e: &Mont, m: &[u64], out: &mut [u64]) -> bool {
let len = m.len();
let mut x = [0u64];
for word in m.iter().rev() {
bignum64::mod_precomp(e, &[*word, x[0]], &mut x);
}
if public(low::bignum_eq(&x, &[0])) {
return false;
}
let e_2 = u64::from(PUBLIC_EXPONENT) - 2;
let mut inv = [0u64];
bignum64::mod_exp_consttime_precomp(e, &x, 17, &[e_2], &mut inv);
let mut scratch = InvertScratch::default();
scratch.t[0] = u64::from(PUBLIC_EXPONENT) - inv[0];
let wide = &mut scratch.wide[..len * 2];
bignum64::mul(len as u32, m, &scratch.t[..len], wide);
let mut carry = 1;
for word in wide[..len].iter_mut() {
let (sum, c) = word.overflowing_add(carry);
*word = sum;
carry = u64::from(c);
}
divide_exact_e(&wide[..len], out);
true
}
struct InvertScratch {
t: [u64; MAX_PRIVATE_MODULUS_WORDS * 2],
wide: [u64; MAX_PRIVATE_MODULUS_WORDS * 4],
}
impl Default for InvertScratch {
fn default() -> Self {
Self {
t: [0; MAX_PRIVATE_MODULUS_WORDS * 2],
wide: [0; MAX_PRIVATE_MODULUS_WORDS * 4],
}
}
}
impl Drop for InvertScratch {
fn drop(&mut self) {
low::zeroise(&mut self.t);
low::zeroise(&mut self.wide);
}
}
fn divide_exact_e(x: &[u64], out: &mut [u64]) {
const E: u64 = PUBLIC_EXPONENT as u64;
const E_INV: u64 = {
let mut inv = E;
let mut i = 0;
while i < 5 {
inv = inv.wrapping_mul(2u64.wrapping_sub(E.wrapping_mul(inv)));
i += 1;
}
assert!(inv.wrapping_mul(E) == 1);
inv
};
let mut borrow = 0u64;
for (q, x) in out.iter_mut().zip(x) {
let (s, b) = x.overflowing_sub(borrow);
*q = s.wrapping_mul(E_INV);
borrow = ((u128::from(*q) * u128::from(E)) >> 64) as u64 + u64::from(b);
}
}
struct Prime([u64; MAX_PRIVATE_MODULUS_WORDS]);
impl Drop for Prime {
fn drop(&mut self) {
low::zeroise(&mut self.0);
}
}
fn random_prime(
size: RsaKeySize,
sieve: &Sieve,
candidate_random: &mut dyn RandomSource,
checks_random: &mut dyn RandomSource,
) -> Result<Prime, Error> {
let len = size.prime_bits() / 64;
let mut buffer = [0u8; MAX_PRIVATE_MODULUS_BYTES];
let mut candidate = Prime([0; MAX_PRIVATE_MODULUS_WORDS]);
loop {
candidate_random.fill(&mut buffer[..len * 8])?;
buffer[0] |= 0b1100_0000;
buffer[len * 8 - 1] |= 0b0000_0001;
for (word, bytes) in candidate
.0
.iter_mut()
.zip(buffer[..len * 8].rchunks_exact(8))
{
*word = u64::from_be_bytes(bytes.try_into().unwrap());
}
ct::secret_slice(&candidate.0[..len]);
if is_prime(&candidate.0[..len], size, sieve, checks_random)? {
return Ok(candidate);
}
}
}
fn is_prime(
candidate: &[u64],
size: RsaKeySize,
sieve: &Sieve,
rng: &mut dyn RandomSource,
) -> Result<bool, Error> {
if !sieve.is_coprime(candidate) {
return Ok(false);
}
miller_rabin(candidate, size, rng)
}
fn miller_rabin(w: &[u64], size: RsaKeySize, rng: &mut dyn RandomSource) -> Result<bool, Error> {
let len = w.len();
let ctx = Mont::new_secret(w, len);
let mut s = MillerRabinScratch::default();
let w1 = &mut s.w1[..len];
w1.copy_from_slice(w);
w1[0] ^= 1;
let mut a = 1;
while public((w1[a / 64] >> (a % 64)) & 1) == 0 {
a += 1;
}
let (shift_words, shift_bits) = (a / 64, a % 64);
for (i, m) in s.m[..len].iter_mut().enumerate() {
let lo = w1.get(i + shift_words).copied().unwrap_or(0);
let hi = w1.get(i + shift_words + 1).copied().unwrap_or(0);
*m = match shift_bits {
0 => lo,
_ => (lo >> shift_bits) | (hi << (64 - shift_bits)),
};
}
let w_len = len * 64 - public(w[len - 1].leading_zeros()) as usize;
let mut iter = 1..=size.miller_rabin_rounds();
loop {
let mut buffer = [0u8; MAX_PRIVATE_MODULUS_BYTES];
let w_bytes = w_len.div_ceil(8);
rng.fill(&mut buffer[..w_bytes])?;
if w_len & 7 != 0 {
buffer[0] &= (1u8 << (w_len & 7)).saturating_sub(1);
}
let b = PosInt::<MAX_PRIVATE_MODULUS_WORDS>::from_bytes(&buffer[..w_bytes])?;
let b = b.words_at_width(len)?;
let suitable = !bignum64::lt_mask(len as u32, b, &ONE_TWO[1][..len])
& bignum64::lt_mask(len as u32, b, &s.w1[..len]);
if public(suitable) == 0 {
continue;
}
if !s.check_base(&ctx, b, a) {
return Ok(false);
}
if iter.next().is_none() {
return Ok(true);
}
}
}
struct MillerRabinScratch {
w1: [u64; MAX_PRIVATE_MODULUS_WORDS],
m: [u64; MAX_PRIVATE_MODULUS_WORDS],
z: [u64; MAX_PRIVATE_MODULUS_WORDS],
wide: [u64; MAX_PRIVATE_MODULUS_WORDS * 2],
}
impl MillerRabinScratch {
fn check_base(&mut self, ctx: &Mont, b: &[u64], a: usize) -> bool {
let len = b.len();
let (w1, z) = (&self.w1[..len], &mut self.z[..len]);
let one = &ONE_TWO[0][..len];
let eq = |z: &[u64], x: &[u64]| public(bignum64::eq_mask(len as u32, z, x)) != 0;
bignum64::mod_exp_consttime_precomp(ctx, b, len as u32 * 64, &self.m[..len], z);
if eq(z, one) || eq(z, w1) {
return true;
}
for _j in 1..a {
let wide = &mut self.wide[..len * 2];
bignum64::sqr(len as u32, z, wide);
bignum64::mod_precomp(ctx, wide, z);
if eq(z, w1) {
return true;
}
if eq(z, one) {
return false;
}
}
false
}
}
impl Default for MillerRabinScratch {
fn default() -> Self {
Self {
w1: [0; MAX_PRIVATE_MODULUS_WORDS],
m: [0; MAX_PRIVATE_MODULUS_WORDS],
z: [0; MAX_PRIVATE_MODULUS_WORDS],
wide: [0; MAX_PRIVATE_MODULUS_WORDS * 2],
}
}
}
impl Drop for MillerRabinScratch {
fn drop(&mut self) {
low::zeroise(&mut self.w1);
low::zeroise(&mut self.m);
low::zeroise(&mut self.z);
low::zeroise(&mut self.wide);
}
}
struct Sieve {
products: Vec<(Mont, Range<usize>)>,
}
impl Sieve {
fn new() -> Self {
let mut products = Vec::new();
let mut start = 0;
while start < SMALL_PRIMES.len() {
let mut product = 1u64;
let mut end = start;
while let Some(next) = SMALL_PRIMES
.get(end)
.and_then(|prime| product.checked_mul(*prime))
{
product = next;
end += 1;
}
products.push((Mont::new(&[product]), start..end));
start = end;
}
Self { products }
}
fn is_coprime(&self, candidate: &[u64]) -> bool {
for (ctx, primes) in &self.products {
let mut r = [0u64];
for word in candidate.iter().rev() {
bignum64::mod_precomp(ctx, &[*word, r[0]], &mut r);
}
let mut divisible = 0;
for (inv, limit) in SMALL_PRIME_DIVISORS[primes.clone()].iter() {
let (_, above) = limit.overflowing_sub(r[0].wrapping_mul(*inv));
divisible |= core::hint::black_box(u64::from(above).wrapping_sub(1));
}
if public(divisible) != 0 {
return false;
}
}
true
}
}
const SMALL_PRIMES: [u64; 130] = {
let mut primes = [0; 130];
let (mut found, mut n) = (0, 3);
while n < 743 {
let mut d = 3;
while d * d <= n && n % d != 0 {
d += 2;
}
if d * d > n {
primes[found] = n;
found += 1;
}
n += 2;
}
assert!(found == primes.len());
primes
};
const SMALL_PRIME_DIVISORS: [(u64, u64); 130] = {
let mut divisors = [(0, 0); 130];
let mut i = 0;
while i < divisors.len() {
let s = SMALL_PRIMES[i];
let mut inv = s;
let mut j = 0;
while j < 5 {
inv = inv.wrapping_mul(2u64.wrapping_sub(s.wrapping_mul(inv)));
j += 1;
}
assert!(inv.wrapping_mul(s) == 1);
divisors[i] = (inv, u64::MAX / s);
i += 1;
}
divisors
};
const ONE_TWO: [[u64; MAX_PRIVATE_MODULUS_WORDS]; 2] = {
let mut r = [[0; MAX_PRIVATE_MODULUS_WORDS]; 2];
r[0][0] = 1;
r[1][0] = 2;
r
};
fn public<T>(v: T) -> T {
ct::into_public(v)
}
#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
use super::*;
use crate::high::hash::Sha256;
use crate::high::hmac_drbg::HmacDrbg;
use crate::mid::rng::{ChainRandomSource, SliceRandomSource, SystemRandom};
fn read_hex(data: &str) -> Vec<u8> {
data.split_whitespace()
.flat_map(|line| {
(0..line.len())
.step_by(2)
.map(|i| u8::from_str_radix(&line[i..i + 2], 16).unwrap())
})
.collect()
}
fn check_primes(data: &str, expected: bool) {
let sieve = Sieve::new();
for line in data.lines().filter(|line| !line.starts_with("#")) {
let candidate =
PosInt::<MAX_PRIVATE_MODULUS_WORDS>::from_bytes(&read_hex(line)).unwrap();
assert_eq!(
is_prime(
candidate.as_words(),
RsaKeySize::Rsa2048,
&sieve,
&mut HmacDrbg::<Sha256>::new(b"", b"", b"")
),
Ok(expected),
"{line}"
);
}
}
#[test]
fn valid_primes() {
check_primes(include_str!("../../../testdata/valid-primes.txt"), true);
}
#[test]
fn invalid_primes() {
check_primes(include_str!("../../../testdata/invalid-primes.txt"), false);
}
#[test]
fn sieve_matches_remainder() {
let sieve = Sieve::new();
assert_eq!(sieve.products.last().unwrap().1.end, SMALL_PRIMES.len());
for s in SMALL_PRIMES.map(u128::from) {
let top = u128::MAX / s * s;
for x in [s, s * 743, top, top - 1, top - s + 1, s << 64] {
let coprime = SMALL_PRIMES.iter().all(|s| x % u128::from(*s) != 0);
let words = [x as u64, (x >> 64) as u64];
assert_eq!(sieve.is_coprime(&words), coprime, "{x:#x}");
}
}
}
#[test]
fn inverts_e() {
let e = Mont::new(&[PUBLIC_EXPONENT.into()]);
for m in [
[3u64, 0],
[6, 0],
[u64::MAX - 3, 1],
[0x1234_5678_9abc_def0, 0xf_ffff_ffff],
] {
let mut d = [0u64; 2];
assert!(invert_e(&e, &m, &mut d));
let (m, d) = (
(u128::from(m[1]) << 64) | u128::from(m[0]),
(u128::from(d[1]) << 64) | u128::from(d[0]),
);
assert!(d < m);
assert_eq!(d * u128::from(PUBLIC_EXPONENT) % m, 1);
}
assert!(!invert_e(&e, &[0x10001 * 7, 0], &mut [0; 2]));
}
#[test]
fn hosed_rng() {
assert_eq!(
generate_key(
RsaKeySize::Rsa2048,
&mut SliceRandomSource(&[]),
&mut SystemRandom
)
.err(),
Some(Error::RngFailed)
);
assert_eq!(
generate_key(
RsaKeySize::Rsa2048,
&mut SystemRandom,
&mut SliceRandomSource(&[]),
)
.err(),
Some(Error::RngFailed)
);
}
#[test]
fn p_q_equal() {
let mut bytes = read_hex(include_str!(
"../../../testdata/rsa.phi-not-coprime-e.2048.txt"
));
assert_eq!(bytes.len(), 256);
bytes.drain(128..);
bytes.extend_from_within(..);
assert_eq!(
generate_key(
RsaKeySize::Rsa2048,
&mut SliceRandomSource(&bytes),
&mut SystemRandom
)
.err(),
Some(Error::RngFailed),
);
}
#[test]
fn phi_not_coprime_e() {
let bytes = read_hex(include_str!(
"../../../testdata/rsa.phi-not-coprime-e.2048.txt"
));
let mut fixed_source = SliceRandomSource(&bytes);
let mut retry_source = SystemRandom;
let mut candidate_source = ChainRandomSource::First(&mut fixed_source, &mut retry_source);
generate_key(
RsaKeySize::Rsa2048,
&mut candidate_source,
&mut SystemRandom,
)
.unwrap();
}
}