diff --git a/lib/crypto/src/arch/mod.rs b/lib/crypto/src/arch/mod.rs index c4952793c34..08a8e0f8cbe 100644 --- a/lib/crypto/src/arch/mod.rs +++ b/lib/crypto/src/arch/mod.rs @@ -54,6 +54,17 @@ pub(crate) fn value_barrier(value: u64) -> u64 { unsafe { core::ptr::read_volatile(&value) } } +/// For a, b < m and n a multiple of four up to 64, by an instruction path: +/// out and the returned top word are a b / 2^(64n) mod m plus at most one +/// m, below 2m and unreduced; `None` where the CPU has no such path. RSA's +/// data is public. +pub(crate) fn rsa_mul(a: &[u64], b: &[u64], m: &[u64], n0: u64, n: usize, out: &mut [u64]) -> Option { + hosted!(x86_64::rsa_mul(a, b, m, n0, n, out), { + let _ = (a, b, m, n0, n, out); + None + }) +} + /// GHASH key powers, (H^(i+1))' at index i in the instruction paths' form. pub(crate) type Powers = [u128; 16]; diff --git a/lib/crypto/src/arch/x86_64.rs b/lib/crypto/src/arch/x86_64.rs index 1778f7a60ca..342c11a66d6 100644 --- a/lib/crypto/src/arch/x86_64.rs +++ b/lib/crypto/src/arch/x86_64.rs @@ -16,6 +16,8 @@ mod aesni; #[cfg(not(target_os = "uefi"))] mod chacha; #[cfg(not(target_os = "uefi"))] +mod montgomery; +#[cfg(not(target_os = "uefi"))] mod poly1305; #[cfg(not(target_os = "uefi"))] mod sha; @@ -40,6 +42,7 @@ impl Features { const VPCLMUL: u32 = 1 << 8; const AVX512VL: u32 = 1 << 9; const AVX512BW: u32 = 1 << 10; + const ADX: u32 = 1 << 11; fn has(self, bits: u32) -> bool { self.0 & bits == bits @@ -64,6 +67,11 @@ impl Features { pub(super) fn bmi2(self) -> bool { self.has(Self::BMI2) } + /// MULX and ADCX/ADOX. + #[cfg(not(target_os = "uefi"))] + pub(super) fn bmi2_adx(self) -> bool { + self.has(Self::BMI2 | Self::ADX) + } /// SHA-NI with the SSSE3 and SSE4.1 shuffles its message schedule takes. #[cfg(not(target_os = "uefi"))] pub(super) fn sha(self) -> bool { @@ -86,7 +94,7 @@ impl Features { /// Intel SDM Vol. 2A, CPUID: leaf 1 ECX bits 1 (PCLMULQDQ), 9 (SSSE3), /// 19 (SSE4.1), 25 (AES), 27 (OSXSAVE) and 28 (AVX); leaf 7 sub-leaf 0 - /// EBX bits 5 (AVX2), 8 (BMI2), 16 (AVX512F), 29 (SHA), 30 (AVX512BW) + /// EBX bits 5 (AVX2), 8 (BMI2), 16 (AVX512F), 19 (ADX), 29 (SHA), 30 (AVX512BW) /// and 31 (AVX512VL), ECX bits 9 (VAES) and 10 (VPCLMULQDQ). Vol. 1 /// ยง13.3: AVX needs XCR0's SSE and AVX state (bits 1 and 2) enabled by /// the OS, AVX-512 its opmask, ZMM_Hi256 and Hi16_ZMM state (bits 5 to 7) @@ -116,6 +124,7 @@ impl Features { | bit(leaf1, 9, Self::SSSE3) | bit(leaf1, 19, Self::SSE41) | bit(leaf7b, 8, Self::BMI2) + | bit(leaf7b, 19, Self::ADX) | bit(leaf7b, 29, Self::SHA); if ymm { bits |= bit(leaf7b, 5, Self::AVX2) | bit(leaf7c, 9, Self::VAES) | bit(leaf7c, 10, Self::VPCLMUL); @@ -151,6 +160,18 @@ fn first_features() -> Features { read } +/// RSA's Montgomery multiplication on BMI2 and ADX, when the CPU has them: +/// out and the returned top word are the unreduced result, below 2m. +#[cfg(not(target_os = "uefi"))] +pub(super) fn rsa_mul(a: &[u64], b: &[u64], m: &[u64], n0: u64, n: usize, out: &mut [u64]) -> Option { + if !features().bmi2_adx() { + return None; + } + assert!(n.is_multiple_of(4) && n <= 64 && a.len() >= n && b.len() >= n && m.len() >= n && out.len() >= n); + // CPUID admitted BMI2 and ADX, and the lengths are checked above. + Some(unsafe { montgomery::mul(a, b, m, n0, n, out) }) +} + /// A vector from 16 bytes by value. With no pointer there is no copy /// precondition, which debug assertions would check per access in hot loops. #[cfg(not(target_os = "uefi"))] @@ -465,7 +486,7 @@ mod tests { for (bit, feature) in [(1, Features::CLMUL), (9, Features::SSSE3), (19, Features::SSE41), (25, Features::AES)] { assert_eq!(Features::decode(!(1 << bit), !0, !0, 0xe7).0 & feature, 0, "leaf 1 bit {bit}"); } - for (bit, feature) in [(5, Features::AVX2), (8, Features::BMI2), (29, Features::SHA)] { + for (bit, feature) in [(5, Features::AVX2), (8, Features::BMI2), (19, Features::ADX), (29, Features::SHA)] { assert_eq!(Features::decode(!0, !(1 << bit), !0, 0xe7).0 & feature, 0, "leaf 7 EBX bit {bit}"); } for (bit, feature) in [(9, Features::VAES), (10, Features::VPCLMUL)] { @@ -481,6 +502,10 @@ mod tests { assert_eq!(f.clmul(), on && std::is_x86_feature_detected!("pclmulqdq")); assert_eq!(f.avx2(), on && std::is_x86_feature_detected!("avx2")); assert_eq!(f.bmi2(), on && std::is_x86_feature_detected!("bmi2")); + assert_eq!( + f.bmi2_adx(), + on && std::is_x86_feature_detected!("bmi2") && std::is_x86_feature_detected!("adx") + ); assert_eq!( f.sha(), on && std::is_x86_feature_detected!("sha") diff --git a/lib/crypto/src/arch/x86_64/montgomery.rs b/lib/crypto/src/arch/x86_64/montgomery.rs new file mode 100644 index 00000000000..278ea0bd0fb --- /dev/null +++ b/lib/crypto/src/arch/x86_64/montgomery.rs @@ -0,0 +1,146 @@ +//! Montgomery multiplication for RSA moduli on BMI2 and ADX: `mulx` with two +//! carry chains, CF through `adcx` and OF through `adox`, so each product +//! row's low and high halves add without waiting on one flag (Intel SDM +//! Vol. 2A, ADCX, ADOX and MULX). The data is public, as all RSA data here. + +use core::arch::asm; + +/// For a, b < m and n a multiple of four up to 64: out and the returned top +/// word are a b / 2^(64n) mod m plus at most one m, below 2m and unreduced. +/// For each word of b, one product row and one reduction row, each four +/// words of a or m per loop step; `t` holds the n + 2 running words. +/// +/// # Safety +/// The CPU has BMI2 and ADX; n is a multiple of four up to 64, and `a`, `b`, +/// `m` and `out` hold at least n words. +#[target_feature(enable = "bmi2,adx")] +pub(super) unsafe fn mul(a: &[u64], b: &[u64], m: &[u64], n0: u64, n: usize, out: &mut [u64]) -> u64 { + let mut t = [0u64; 66]; + for &bi in &b[..n] { + // SAFETY: the rows read n words of `a` and `m` and write n + 2 of `t`; + // they write no other memory and leave the stack alone. + unsafe { row(a.as_ptr(), bi, m.as_ptr(), n0, n, t.as_mut_ptr()) }; + } + out[..n].copy_from_slice(&t[..n]); + t[n] +} + +/// t = (t + a bi + q m) / 2^64 for q = (t + a bi)_0 n0 mod 2^64. +#[target_feature(enable = "bmi2,adx")] +unsafe fn row(a: *const u64, bi: u64, m: *const u64, n0: u64, n: usize, t: *mut u64) { + unsafe { + asm!( + // Product row: t[j] += lo(a_j bi) on CF and hi(a_{j-1} bi) on OF. + "mov {a_ptr}, {a}", + "mov {t_ptr}, {t}", + "mov rcx, {quads}", + "xor {hi:e}, {hi:e}", + "2:", + "mulx {h0}, {lo}, [{a_ptr}]", + "adcx {lo}, [{t_ptr}]", + "adox {lo}, {hi}", + "mov [{t_ptr}], {lo}", + "mulx {hi}, {lo}, [{a_ptr} + 8]", + "adcx {lo}, [{t_ptr} + 8]", + "adox {lo}, {h0}", + "mov [{t_ptr} + 8], {lo}", + "mulx {h0}, {lo}, [{a_ptr} + 16]", + "adcx {lo}, [{t_ptr} + 16]", + "adox {lo}, {hi}", + "mov [{t_ptr} + 16], {lo}", + "mulx {hi}, {lo}, [{a_ptr} + 24]", + "adcx {lo}, [{t_ptr} + 24]", + "adox {lo}, {h0}", + "mov [{t_ptr} + 24], {lo}", + "lea {a_ptr}, [{a_ptr} + 32]", + "lea {t_ptr}, [{t_ptr} + 32]", + "lea rcx, [rcx - 1]", + "jrcxz 3f", + "jmp 2b", + "3:", + // t[n] += hi + CF + OF; t[n + 1] = the carries out of that. + "mov {lo}, 0", + "mov {h0}, [{t_ptr}]", + "adcx {h0}, {lo}", + "adox {h0}, {hi}", + "mov [{t_ptr}], {h0}", + "mov {h0}, 0", + "adcx {h0}, {lo}", + "adox {h0}, {lo}", + "mov [{t_ptr} + 8], {h0}", + // Reduction row: q = t_0 n0; t[j - 1] = t[j] + lo(q m_j) + hi(q m_{j-1}). + "mov rdx, [{t}]", + "imul rdx, {n0}", + "mov {a_ptr}, {m}", + "mov {t_ptr}, {t}", + "mov rcx, {quads}", + "xor {hi:e}, {hi:e}", + // Word 0 of the sum is zero by the choice of q; only its carries + // remain, so it is computed and not stored. + "mulx {hi}, {lo}, [{a_ptr}]", + "adcx {lo}, [{t_ptr}]", + "mulx {h0}, {lo}, [{a_ptr} + 8]", + "adcx {lo}, [{t_ptr} + 8]", + "adox {lo}, {hi}", + "mov [{t_ptr}], {lo}", + "mulx {hi}, {lo}, [{a_ptr} + 16]", + "adcx {lo}, [{t_ptr} + 16]", + "adox {lo}, {h0}", + "mov [{t_ptr} + 8], {lo}", + "mulx {h0}, {lo}, [{a_ptr} + 24]", + "adcx {lo}, [{t_ptr} + 24]", + "adox {lo}, {hi}", + "mov [{t_ptr} + 16], {lo}", + "lea {a_ptr}, [{a_ptr} + 32]", + "lea {t_ptr}, [{t_ptr} + 32]", + "lea rcx, [rcx - 1]", + "jrcxz 5f", + "4:", + "mulx {hi}, {lo}, [{a_ptr}]", + "adcx {lo}, [{t_ptr}]", + "adox {lo}, {h0}", + "mov [{t_ptr} - 8], {lo}", + "mulx {h0}, {lo}, [{a_ptr} + 8]", + "adcx {lo}, [{t_ptr} + 8]", + "adox {lo}, {hi}", + "mov [{t_ptr}], {lo}", + "mulx {hi}, {lo}, [{a_ptr} + 16]", + "adcx {lo}, [{t_ptr} + 16]", + "adox {lo}, {h0}", + "mov [{t_ptr} + 8], {lo}", + "mulx {h0}, {lo}, [{a_ptr} + 24]", + "adcx {lo}, [{t_ptr} + 24]", + "adox {lo}, {hi}", + "mov [{t_ptr} + 16], {lo}", + "lea {a_ptr}, [{a_ptr} + 32]", + "lea {t_ptr}, [{t_ptr} + 32]", + "lea rcx, [rcx - 1]", + "jrcxz 5f", + "jmp 4b", + "5:", + // t[n - 1] = t[n] + hi + CF + OF; t[n] = t[n + 1] + the carries. + "mov {lo}, [{t_ptr}]", + "adcx {lo}, {zero}", + "adox {lo}, {h0}", + "mov [{t_ptr} - 8], {lo}", + "mov {lo}, [{t_ptr} + 8]", + "adcx {lo}, {zero}", + "adox {lo}, {zero}", + "mov [{t_ptr}], {lo}", + a = in(reg) a, + m = in(reg) m, + t = in(reg) t, + n0 = in(reg) n0, + quads = in(reg) n / 4, + zero = in(reg) 0u64, + a_ptr = out(reg) _, + t_ptr = out(reg) _, + lo = out(reg) _, + hi = out(reg) _, + h0 = out(reg) _, + inout("rdx") bi => _, + out("rcx") _, + options(nostack), + ); + } +} diff --git a/lib/crypto/src/rsa.rs b/lib/crypto/src/rsa.rs index 0f9fd7538ed..3ffa2df3950 100644 --- a/lib/crypto/src/rsa.rs +++ b/lib/crypto/src/rsa.rs @@ -78,7 +78,9 @@ impl PublicKey { if bits < 2048 || e < 3 || e & 1 == 0 || modulus[modulus.len() - 1] & 1 == 0 { return Err(Error::InvalidKey); } - let words = modulus.len().div_ceil(8); + // Whole groups of four words, which the instruction path takes; the + // zero words above the modulus only enlarge R. + let words = modulus.len().div_ceil(32) * 4; let m = read(modulus); let mut key = Self { modulus: m, @@ -119,11 +121,23 @@ impl PublicKey { subtract(x, &self.modulus, self.words); } } - /// out = a * b / R mod n, for a, b < n: word-serial Montgomery - /// multiplication with the reduction step interleaved (CIOS; Koc, Acar and - /// Kaliski, "Analyzing and comparing Montgomery multiplication - /// algorithms", IEEE Micro 1996). + /// out = a * b / R mod n, for a, b < n, by the instruction path where + /// the CPU has one. fn mul(&self, a: &[u64; WORDS], b: &[u64; WORDS], out: &mut [u64; WORDS]) { + let n = self.words; + match crate::arch::rsa_mul(a, b, &self.modulus, self.n0, n, out) { + Some(top) => { + if top != 0 || !less(out, &self.modulus, n) { + subtract(out, &self.modulus, n); + } + } + None => self.mul_portable(a, b, out), + } + } + /// Word-serial Montgomery multiplication with the reduction step + /// interleaved (CIOS; Koc, Acar and Kaliski, "Analyzing and comparing + /// Montgomery multiplication algorithms", IEEE Micro 1996). + fn mul_portable(&self, a: &[u64; WORDS], b: &[u64; WORDS], out: &mut [u64; WORDS]) { let n = self.words; let (a, b, m) = (&a[..n], &b[..n], &self.modulus[..n]); let mut t = [0u64; WORDS + 2]; @@ -331,6 +345,53 @@ mod tests { } } + /// The instruction path, where the CPU has one, equals the portable + /// multiplication for the padded word counts 32 to 64 in steps of 8, and + /// at the largest residues. + #[test] + fn instruction_path_matches_the_portable_multiplication() { + // Which path this run compares, for a reader of `--nocapture` output. + let (zero, modulus) = ([0u64; WORDS], [u64::MAX; WORDS]); + let mut out = [0u64; WORDS]; + let kernel = crate::arch::rsa_mul(&zero, &zero, &modulus, 1, 4, &mut out).is_some(); + std::println!("instruction path: {}", if kernel { "the BMI2/ADX kernel" } else { "none; portable against itself" }); + let mut seed = 0x2545f4914f6cdd1du64; + let mut next = move || { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + seed + }; + for len in [256usize, 300, 384, 448, 512] { + let mut modulus: std::vec::Vec = (0..len).map(|_| next() as u8).collect(); + modulus[0] |= 0x80; + modulus[len - 1] |= 1; + let key = PublicKey::from_components(&modulus, &[1, 0, 1]).unwrap(); + let n = BigUint::from_bytes_be(&modulus); + for case in 0..64 { + let mut a = [0u64; WORDS]; + let mut b = [0u64; WORDS]; + let (x, y) = if case == 0 { + (&n - 1u8, &n - 1u8) + } else { + let x: std::vec::Vec = (0..len).map(|_| next() as u8).collect(); + let y: std::vec::Vec = (0..len).map(|_| next() as u8).collect(); + (BigUint::from_bytes_be(&x) % &n, BigUint::from_bytes_be(&y) % &n) + }; + for (i, d) in x.to_u64_digits().iter().enumerate() { + a[i] = *d; + } + for (i, d) in y.to_u64_digits().iter().enumerate() { + b[i] = *d; + } + let (mut ours, mut portable) = ([0; WORDS], [0; WORDS]); + key.mul(&a, &b, &mut ours); + key.mul_portable(&a, &b, &mut portable); + assert_eq!(ours, portable, "{len} bytes, case {case}"); + } + } + } + fn integer(words: &[u64; WORDS]) -> BigUint { BigUint::from_bytes_le(&words.iter().flat_map(|w| w.to_le_bytes()).collect::>()) }