Add RSAZ ABI tests.

As part of this, move the CPU checks to C.

Change-Id: I17b701e1196c1ca116bbd23e0e669cf603ad464d
Reviewed-on: https://boringssl-review.googlesource.com/c/34626
Reviewed-by: Adam Langley <agl@google.com>
Commit-Queue: David Benjamin <davidben@google.com>
diff --git a/crypto/fipsmodule/bn/asm/rsaz-avx2.pl b/crypto/fipsmodule/bn/asm/rsaz-avx2.pl
index 72aa2a3..51feb69 100755
--- a/crypto/fipsmodule/bn/asm/rsaz-avx2.pl
+++ b/crypto/fipsmodule/bn/asm/rsaz-avx2.pl
@@ -52,7 +52,6 @@
 # versions, but BoringSSL is intended to be used with pre-generated perlasm
 # output, so this isn't useful anyway.
 $avx = 2;
-$addx = 1;
 
 open OUT,"| \"$^X\" \"$xlate\" $flavour \"$output\"";
 *STDOUT = *OUT;
@@ -1474,6 +1473,7 @@
 .type	rsaz_1024_red2norm_avx2,\@abi-omnipotent
 .align	32
 rsaz_1024_red2norm_avx2:
+.cfi_startproc
 	sub	\$-128,$inp	# size optimization
 	xor	%rax,%rax
 ___
@@ -1507,12 +1507,14 @@
 }
 $code.=<<___;
 	ret
+.cfi_endproc
 .size	rsaz_1024_red2norm_avx2,.-rsaz_1024_red2norm_avx2
 
 .globl	rsaz_1024_norm2red_avx2
 .type	rsaz_1024_norm2red_avx2,\@abi-omnipotent
 .align	32
 rsaz_1024_norm2red_avx2:
+.cfi_startproc
 	sub	\$-128,$out	# size optimization
 	mov	($inp),@T[0]
 	mov	\$0x1fffffff,%eax
@@ -1544,6 +1546,7 @@
 	mov	@T[0],`8*($j+2)-128`($out)
 	mov	@T[0],`8*($j+3)-128`($out)
 	ret
+.cfi_endproc
 .size	rsaz_1024_norm2red_avx2,.-rsaz_1024_norm2red_avx2
 ___
 }
@@ -1555,6 +1558,7 @@
 .type	rsaz_1024_scatter5_avx2,\@abi-omnipotent
 .align	32
 rsaz_1024_scatter5_avx2:
+.cfi_startproc
 	vzeroupper
 	vmovdqu	.Lscatter_permd(%rip),%ymm5
 	shl	\$4,$power
@@ -1574,6 +1578,7 @@
 
 	vzeroupper
 	ret
+.cfi_endproc
 .size	rsaz_1024_scatter5_avx2,.-rsaz_1024_scatter5_avx2
 
 .globl	rsaz_1024_gather5_avx2
@@ -1733,27 +1738,6 @@
 }
 
 $code.=<<___;
-.extern	OPENSSL_ia32cap_P
-.globl	rsaz_avx2_eligible
-.type	rsaz_avx2_eligible,\@abi-omnipotent
-.align	32
-rsaz_avx2_eligible:
-	leaq	OPENSSL_ia32cap_P(%rip),%rax
-	mov	8(%rax),%eax
-___
-$code.=<<___	if ($addx);
-	mov	\$`1<<8|1<<19`,%ecx
-	mov	\$0,%edx
-	and	%eax,%ecx
-	cmp	\$`1<<8|1<<19`,%ecx	# check for BMI2+AD*X
-	cmove	%edx,%eax
-___
-$code.=<<___;
-	and	\$`1<<5`,%eax
-	shr	\$5,%eax
-	ret
-.size	rsaz_avx2_eligible,.-rsaz_avx2_eligible
-
 .align	64
 .Land_mask:
 	.quad	0x1fffffff,0x1fffffff,0x1fffffff,0x1fffffff
diff --git a/crypto/fipsmodule/bn/bn_test.cc b/crypto/fipsmodule/bn/bn_test.cc
index be0a86e..b7d98da 100644
--- a/crypto/fipsmodule/bn/bn_test.cc
+++ b/crypto/fipsmodule/bn/bn_test.cc
@@ -86,6 +86,7 @@
 #include <openssl/rand.h>
 
 #include "./internal.h"
+#include "./rsaz_exp.h"
 #include "../../internal.h"
 #include "../../test/abi_test.h"
 #include "../../test/file_test.h"
@@ -2403,3 +2404,35 @@
   }
 }
 #endif   // OPENSSL_BN_ASM_MONT && SUPPORTS_ABI_TEST
+
+#if defined(RSAZ_ENABLED) && defined(SUPPORTS_ABI_TEST)
+TEST_F(BNTest, RSAZABI) {
+  if (!rsaz_avx2_capable()) {
+    return;
+  }
+
+  alignas(64) BN_ULONG table[32 * 18] = {0};
+  alignas(64) BN_ULONG rsaz1[40], rsaz2[40], rsaz3[40], n_rsaz[40];
+  BN_ULONG norm[16], n_norm[16];
+
+  OPENSSL_memset(norm, 0x42, sizeof(norm));
+  OPENSSL_memset(n_norm, 0x99, sizeof(n_norm));
+
+  bssl::UniquePtr<BIGNUM> n(BN_new());
+  ASSERT_TRUE(n);
+  ASSERT_TRUE(bn_set_words(n.get(), n_norm, 16));
+  bssl::UniquePtr<BN_MONT_CTX> mont(
+      BN_MONT_CTX_new_for_modulus(n.get(), nullptr));
+  ASSERT_TRUE(mont);
+  const BN_ULONG k = mont->n0[0];
+
+  CHECK_ABI(rsaz_1024_norm2red_avx2, rsaz1, norm);
+  CHECK_ABI(rsaz_1024_norm2red_avx2, n_rsaz, n_norm);
+  CHECK_ABI(rsaz_1024_sqr_avx2, rsaz2, rsaz1, n_rsaz, k, 1);
+  CHECK_ABI(rsaz_1024_sqr_avx2, rsaz3, rsaz2, n_rsaz, k, 4);
+  CHECK_ABI(rsaz_1024_mul_avx2, rsaz3, rsaz1, rsaz2, n_rsaz, k);
+  CHECK_ABI(rsaz_1024_scatter5_avx2, table, rsaz3, 7);
+  CHECK_ABI(rsaz_1024_gather5_avx2, rsaz1, table, 7);
+  CHECK_ABI(rsaz_1024_red2norm_avx2, norm, rsaz1);
+}
+#endif   // RSAZ_ENABLED && SUPPORTS_ABI_TEST
diff --git a/crypto/fipsmodule/bn/exponentiation.c b/crypto/fipsmodule/bn/exponentiation.c
index 1b9680f..29e7877 100644
--- a/crypto/fipsmodule/bn/exponentiation.c
+++ b/crypto/fipsmodule/bn/exponentiation.c
@@ -117,13 +117,11 @@
 #include <openssl/mem.h>
 
 #include "internal.h"
+#include "rsaz_exp.h"
 
 
 #if !defined(OPENSSL_NO_ASM) && defined(OPENSSL_X86_64)
 #define OPENSSL_BN_ASM_MONT5
-#define RSAZ_ENABLED
-
-#include "rsaz_exp.h"
 
 void bn_mul_mont_gather5(BN_ULONG *rp, const BN_ULONG *ap,
                          const BN_ULONG *table, const BN_ULONG *np,
@@ -974,12 +972,12 @@
   alignas(MOD_EXP_CTIME_MIN_CACHE_LINE_WIDTH) BN_ULONG
     storage[MOD_EXP_CTIME_STORAGE_LEN];
 #endif
-#ifdef RSAZ_ENABLED
-  // If the size of the operands allow it, perform the optimized
-  // RSAZ exponentiation. For further information see
-  // crypto/bn/rsaz_exp.c and accompanying assembly modules.
-  if ((16 == a->width) && (16 == p->width) && (BN_num_bits(m) == 1024) &&
-      rsaz_avx2_eligible()) {
+#if defined(RSAZ_ENABLED)
+  // If the size of the operands allow it, perform the optimized RSAZ
+  // exponentiation. For further information see crypto/fipsmodule/bn/rsaz_exp.c
+  // and accompanying assembly modules.
+  if (a->width == 16 && p->width == 16 && BN_num_bits(m) == 1024 &&
+      rsaz_avx2_preferred()) {
     if (!bn_wexpand(rr, 16)) {
       goto err;
     }
diff --git a/crypto/fipsmodule/bn/rsaz_exp.c b/crypto/fipsmodule/bn/rsaz_exp.c
index c562ea3..7e15aaf 100644
--- a/crypto/fipsmodule/bn/rsaz_exp.c
+++ b/crypto/fipsmodule/bn/rsaz_exp.c
@@ -12,54 +12,16 @@
  * (2) University of Haifa, Israel
  */
 
-#include <openssl/base.h>
-
-#if !defined(OPENSSL_NO_ASM) && defined(OPENSSL_X86_64)
-
 #include "rsaz_exp.h"
 
+#if defined(RSAZ_ENABLED)
+
 #include <openssl/mem.h>
 
 #include "internal.h"
 #include "../../internal.h"
 
 
-// RSAZ represents 1024-bit integers using unsaturated 29-bit limbs stored in
-// 64-bit integers. This requires 36 limbs but padded up to 40.
-//
-// See crypto/bn/asm/rsaz-avx2.pl for further details.
-
-// rsaz_1024_norm2red_avx2 converts |norm| from |BIGNUM| to RSAZ representation
-// and writes the result to |red|.
-void rsaz_1024_norm2red_avx2(BN_ULONG red[40], const BN_ULONG norm[16]);
-
-// rsaz_1024_mul_avx2 computes |a| * |b| mod |n| and writes the result to |ret|.
-// Inputs and outputs are in Montgomery form, using RSAZ's representation. |k|
-// is -|n|^-1 mod 2^64 or |n0| from |BN_MONT_CTX|.
-void rsaz_1024_mul_avx2(BN_ULONG ret[40], const BN_ULONG a[40],
-                        const BN_ULONG b[40], const BN_ULONG n[40], BN_ULONG k);
-
-// rsaz_1024_mul_avx2 computes |a|^(2*|count|) mod |n| and writes the result to
-// |ret|. Inputs and outputs are in Montgomery form, using RSAZ's
-// representation. |k| is -|n|^-1 mod 2^64 or |n0| from |BN_MONT_CTX|.
-void rsaz_1024_sqr_avx2(BN_ULONG ret[40], const BN_ULONG a[40],
-                        const BN_ULONG n[40], BN_ULONG k, int count);
-
-// rsaz_1024_scatter5_avx2 stores |val| at index |i| of |tbl|. |i| must be
-// positive and at most 31. Note the table only uses 18 |BN_ULONG|s per entry
-// instead of 40. It packs two 29-bit limbs into each |BN_ULONG| and only stores
-// 36 limbs rather than the padded 40.
-void rsaz_1024_scatter5_avx2(BN_ULONG tbl[32 * 18], const BN_ULONG val[40],
-                             int i);
-
-// rsaz_1024_gather5_avx2 loads index |i| of |tbl| and writes it to |val|.
-void rsaz_1024_gather5_avx2(BN_ULONG val[40], const BN_ULONG tbl[32 * 18],
-                            int i);
-
-// rsaz_1024_red2norm_avx2 converts |red| from RSAZ to |BIGNUM| representation
-// and writes the result to |norm|.
-void rsaz_1024_red2norm_avx2(BN_ULONG norm[16], const BN_ULONG red[40]);
-
 // one is 1 in RSAZ's representation.
 alignas(64) static const BN_ULONG one[40] = {
     1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
@@ -261,4 +223,4 @@
   OPENSSL_cleanse(storage, MOD_EXP_CTIME_STORAGE_LEN * sizeof(BN_ULONG));
 }
 
-#endif  // OPENSSL_X86_64
+#endif  // RSAZ_ENABLED
diff --git a/crypto/fipsmodule/bn/rsaz_exp.h b/crypto/fipsmodule/bn/rsaz_exp.h
index b6aea14..3b06192 100644
--- a/crypto/fipsmodule/bn/rsaz_exp.h
+++ b/crypto/fipsmodule/bn/rsaz_exp.h
@@ -16,9 +16,18 @@
 #define OPENSSL_HEADER_BN_RSAZ_EXP_H
 
 #include <openssl/bn.h>
+#include <openssl/cpu.h>
 
 #include "internal.h"
 
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if !defined(OPENSSL_NO_ASM) && defined(OPENSSL_X86_64)
+#define RSAZ_ENABLED
+
+
 // RSAZ_1024_mod_exp_avx2 sets |result| to |base_norm| raised to |exponent|
 // modulo |m_norm|. |base_norm| must be fully-reduced and |exponent| must have
 // the high bit set (it is 1024 bits wide). |RR| and |k0| must be |RR| and |n0|,
@@ -31,8 +40,65 @@
                             BN_ULONG k0,
                             BN_ULONG storage_words[MOD_EXP_CTIME_STORAGE_LEN]);
 
-// rsaz_avx2_eligible returns one if |RSAZ_1024_mod_exp_avx2| should be used and
-// zero otherwise.
-int rsaz_avx2_eligible(void);
+OPENSSL_INLINE int rsaz_avx2_capable(void) {
+  const uint32_t *cap = OPENSSL_ia32cap_get();
+  return (cap[2] & (1 << 5)) != 0;  // AVX2
+}
+
+OPENSSL_INLINE int rsaz_avx2_preferred(void) {
+  const uint32_t *cap = OPENSSL_ia32cap_get();
+  static const uint32_t kBMI2AndADX = (1 << 8) | (1 << 19);
+  if ((cap[2] & kBMI2AndADX) == kBMI2AndADX) {
+    // If BMI2 and ADX are available, x86_64-mont5.pl is faster.
+    return 0;
+  }
+  return (cap[2] & (1 << 5)) != 0;  // AVX2
+}
+
+
+// Assembly functions.
+
+// RSAZ represents 1024-bit integers using unsaturated 29-bit limbs stored in
+// 64-bit integers. This requires 36 limbs but padded up to 40.
+//
+// See crypto/bn/asm/rsaz-avx2.pl for further details.
+
+// rsaz_1024_norm2red_avx2 converts |norm| from |BIGNUM| to RSAZ representation
+// and writes the result to |red|.
+void rsaz_1024_norm2red_avx2(BN_ULONG red[40], const BN_ULONG norm[16]);
+
+// rsaz_1024_mul_avx2 computes |a| * |b| mod |n| and writes the result to |ret|.
+// Inputs and outputs are in Montgomery form, using RSAZ's representation. |k|
+// is -|n|^-1 mod 2^64 or |n0| from |BN_MONT_CTX|.
+void rsaz_1024_mul_avx2(BN_ULONG ret[40], const BN_ULONG a[40],
+                        const BN_ULONG b[40], const BN_ULONG n[40], BN_ULONG k);
+
+// rsaz_1024_mul_avx2 computes |a|^(2*|count|) mod |n| and writes the result to
+// |ret|. Inputs and outputs are in Montgomery form, using RSAZ's
+// representation. |k| is -|n|^-1 mod 2^64 or |n0| from |BN_MONT_CTX|.
+void rsaz_1024_sqr_avx2(BN_ULONG ret[40], const BN_ULONG a[40],
+                        const BN_ULONG n[40], BN_ULONG k, int count);
+
+// rsaz_1024_scatter5_avx2 stores |val| at index |i| of |tbl|. |i| must be
+// positive and at most 31. Note the table only uses 18 |BN_ULONG|s per entry
+// instead of 40. It packs two 29-bit limbs into each |BN_ULONG| and only stores
+// 36 limbs rather than the padded 40.
+void rsaz_1024_scatter5_avx2(BN_ULONG tbl[32 * 18], const BN_ULONG val[40],
+                             int i);
+
+// rsaz_1024_gather5_avx2 loads index |i| of |tbl| and writes it to |val|.
+void rsaz_1024_gather5_avx2(BN_ULONG val[40], const BN_ULONG tbl[32 * 18],
+                            int i);
+
+// rsaz_1024_red2norm_avx2 converts |red| from RSAZ to |BIGNUM| representation
+// and writes the result to |norm|.
+void rsaz_1024_red2norm_avx2(BN_ULONG norm[16], const BN_ULONG red[40]);
+
+
+#endif  // !OPENSSL_NO_ASM && OPENSSL_X86_64
+
+#if defined(__cplusplus)
+}  // extern "C"
+#endif
 
 #endif  // OPENSSL_HEADER_BN_RSAZ_EXP_H