Extract and test the deterministic part of Miller-Rabin.

This way we test not only that we match expectations for primes and
composites but that the core test correctly reports false witnesses. I
made an initial attempt to gather some interesting test input, but
probably one can do better.

Change-Id: I7c29afb534bd6980ef42a893e86d86bd44af8349
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/38164
Commit-Queue: David Benjamin <davidben@google.com>
Reviewed-by: Adam Langley <agl@google.com>
diff --git a/crypto/fipsmodule/bn/bn_test.cc b/crypto/fipsmodule/bn/bn_test.cc
index 0be8396..ea3967b 100644
--- a/crypto/fipsmodule/bn/bn_test.cc
+++ b/crypto/fipsmodule/bn/bn_test.cc
@@ -2279,6 +2279,32 @@
   EXPECT_EQ(1, is_probably_prime_1);
 }
 
+TEST_F(BNTest, MillerRabinIteration) {
+  FileTestGTest(
+      "crypto/fipsmodule/bn/miller_rabin_tests.txt", [&](FileTest *t) {
+        BIGNUMFileTest bn_test(t, /*large_mask=*/0);
+
+        bssl::UniquePtr<BIGNUM> w = bn_test.GetBIGNUM("W");
+        ASSERT_TRUE(w);
+        bssl::UniquePtr<BIGNUM> b = bn_test.GetBIGNUM("B");
+        ASSERT_TRUE(b);
+        bssl::UniquePtr<BN_MONT_CTX> mont(
+            BN_MONT_CTX_new_consttime(w.get(), ctx()));
+        ASSERT_TRUE(mont);
+
+        bssl::BN_CTXScope scope(ctx());
+        BN_MILLER_RABIN miller_rabin;
+        ASSERT_TRUE(bn_miller_rabin_init(&miller_rabin, mont.get(), ctx()));
+        int possibly_prime;
+        ASSERT_TRUE(bn_miller_rabin_iteration(&miller_rabin, &possibly_prime,
+                                              b.get(), mont.get(), ctx()));
+
+        std::string result;
+        ASSERT_TRUE(t->GetAttribute(&result, "Result"));
+        EXPECT_EQ(result, possibly_prime ? "PossiblyPrime" : "Composite");
+      });
+}
+
 TEST_F(BNTest, NumBitsWord) {
   constexpr BN_ULONG kOne = 1;
 
diff --git a/crypto/fipsmodule/bn/internal.h b/crypto/fipsmodule/bn/internal.h
index c1e60fe..d58a2ac 100644
--- a/crypto/fipsmodule/bn/internal.h
+++ b/crypto/fipsmodule/bn/internal.h
@@ -437,6 +437,40 @@
 // of the first several odd primes and zero otherwise.
 int bn_odd_number_is_obviously_composite(const BIGNUM *bn);
 
+// A BN_MILLER_RABIN stores state common to each Miller-Rabin iteration. It is
+// initialized within an existing |BN_CTX| scope and may not be used after
+// that scope is released with |BN_CTX_end|. Field names match those in FIPS
+// 186-4, section C.3.1.
+typedef struct {
+  // w1 is w-1.
+  BIGNUM *w1;
+  // m is (w-1)/2^a.
+  BIGNUM *m;
+  // one_mont is 1 (mod w) in Montgomery form.
+  BIGNUM *one_mont;
+  // w1_mont is w-1 (mod w) in Montgomery form.
+  BIGNUM *w1_mont;
+  // w_bits is BN_num_bits(w).
+  int w_bits;
+  // a is the largest integer such that 2^a divides w-1.
+  int a;
+} BN_MILLER_RABIN;
+
+// bn_miller_rabin_init initializes |miller_rabin| for testing if |mont->N| is
+// prime. It returns one on success and zero on error.
+OPENSSL_EXPORT int bn_miller_rabin_init(BN_MILLER_RABIN *miller_rabin,
+                                        const BN_MONT_CTX *mont, BN_CTX *ctx);
+
+// bn_miller_rabin_iteration performs one Miller-Rabin iteration, checking if
+// |b| is a composite witness for |mont->N|. |miller_rabin| must have been
+// initialized with |bn_miller_rabin_setup|. On success, it returns one and sets
+// |*out_is_possibly_prime| to one if |mont->N| may still be prime or zero if
+// |b| shows it is composite. On allocation or internal failure, it returns
+// zero.
+OPENSSL_EXPORT int bn_miller_rabin_iteration(
+    const BN_MILLER_RABIN *miller_rabin, int *out_is_possibly_prime,
+    const BIGNUM *b, const BN_MONT_CTX *mont, BN_CTX *ctx);
+
 // bn_rshift1_words sets |r| to |a| >> 1, where both arrays are |num| bits wide.
 void bn_rshift1_words(BN_ULONG *r, const BN_ULONG *a, size_t num);
 
diff --git a/crypto/fipsmodule/bn/miller_rabin_tests.txt b/crypto/fipsmodule/bn/miller_rabin_tests.txt
new file mode 100644
index 0000000..437d0f9
--- /dev/null
+++ b/crypto/fipsmodule/bn/miller_rabin_tests.txt
@@ -0,0 +1,322 @@
+# This file contains test vectors for whether B is a Miller-Rabin composite
+# witness for W. W must be odd and B must satisfy 1 <= B <= W-1.
+#
+# The following Python function may be used to check values.
+#
+#  def is_miller_rabin_witness(w, b):
+#      # Variable names taken from FIPS 186-4 C.3.1 but the algorithm skips a
+#      # couple of optimizations in the FIPS formulation.
+#      m = w - 1
+#      a = 0
+#      while m&1 == 0:
+#          a += 1
+#          m //= 2
+#      # b is a composite witness for w iff the following are true:
+#      # - b^m != 1 (mod w)
+#      # - b^(m*2^j) != -1 (mod w), for 0 <= j < a
+#      z = pow(b, m, w)
+#      if z == 1:
+#         # b^m = 1 (mod w)
+#         return False
+#      for j in range(a):
+#         if z == w-1:
+#             # b^(m*2^j) = -1 (mod w)
+#             return False
+#         z = (z * z) % w
+#      # At this point, z is b^(w-1) (mod w). If z is not 1, w has failed the
+#      # Fermat test and is composite. If z is 1, the value of z immediately
+#      # before it became 1 is a non-trivial root of unity and w is composite.
+#      return True
+
+# Exhaustively test a small prime.
+
+Result = PossiblyPrime
+W = 7
+B = 1
+
+Result = PossiblyPrime
+W = 7
+B = 2
+
+Result = PossiblyPrime
+W = 7
+B = 3
+
+Result = PossiblyPrime
+W = 7
+B = 4
+
+Result = PossiblyPrime
+W = 7
+B = 5
+
+Result = PossiblyPrime
+W = 7
+B = 6
+
+
+# Random large inputs which try to cover a few cases. The nontrivial square root
+# case appears to be difficult to hit randomly.
+
+# b^m = w-1
+Result = PossiblyPrime
+W = d6b4ffc7cf70b2a2fc5d6023015875504d40e3dcce7c2e6b762c3de7bb806a5074144e7054198dabf53d23108679ccc541d5a99efeb1d1abaf89e0dbcead2a8b
+B = fabbafdbec6494ddb5ea4bf458536e87082369b0e53a200ed413f3e64b2fddc7c57c565710fbe73fae5b188fce97d8dcca74c2b5d90906c96d3c2c358a735cd
+
+# b^m = w-1
+Result = PossiblyPrime
+W = 52cc61c42b341ad56dc11495e7cb2fe31e506b9e99522efbf44cd7c28468d3833c5e360f3c77b0aa43c0495c4e14665ab0d7cee9294c722f0de47d4401828401
+B = 3bdc9639c0fc2e77ab48d46e0b4ac6529c11c900e8fe4d82d75767c0556feb23d3f42d4924d16876a743feb386b7b84c7fd16a6c252f662faf0024d19972e62f
+
+# b^m = w-1
+Result = PossiblyPrime
+W = cff9897aa7dce0f2afad262b2de57d301305de717f3539c537c4ce062f8cb70df13fbc1eb4a3b9f0958a8810d1ca9042b4f23334b285a15fee3fc66498761d4b
+B = 9ceb43132fddf9ee4104ea1cb3eb2253c1d7f803f05f0305de9e31a17dd75832f47b8bf189a9b7ca0905f2a7470d9c6349080f481ff1708696fa12d972e7d7ba
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 67d1825dad5344170e65247a87aef1634a1b32bdc22f2f04d9d2959767bb5a27610fba55cd607e0f9fdd9fbb0f7f98e40d5e1eb2f52318fb5be4dbfd30d38861
+B = 260fb14724ff80984736859d8755ee98b25bcb56db9fde1db001a1e1273374034c5b75fd60b3710c7a08ce7d390776f010f384d4e32943cf0c477497d53e9e05
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = ad0bc85b58aaa204177aa9431a40929beb1cbea2dd6f66a25cc54600013213b225ba881805661df43f4208965ada7aacc8095d07d3cbef1a7bbfaae8b745f731
+B = 3d9310f20e9c80269fa6830c7e1a6f02fc5c58646001a9ef6b8b3e496602ff22c3dcb2ddb6a221723fc1722ce237fb46f7a7bb2945e415c8839b15a972f076c9
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = b25c917f55f6c7b596921daba919f35039e5d805119c1587e99849dd7104460c86214f162a6f17aea847bc7f3859e59f2991d457059511972ef373d4bc75e309
+B = a1f10b261dee84619b0423201d46af19eef9ec0612cf947c4d5c36c0c4b28207f75967e69452eabad0a5dcd28f27f7a8a7ed9c8b3e5026c6e0ba5634d94c2d44
+
+# b^m = 1
+Result = PossiblyPrime
+W = d3eeb0eff05b6992e9fa61b02755e155f4aae28c6e45ddb874edd86acdd2d83d18a20e0e00d8b8bc94b92d14fc3f41ced6ababe8ac98c7730c075dbe0f699369
+B = 6b7717269c6225203681a1cacec87cacd83003ec6e9e3f04effcc4f86634770c0860e1f2770b8f303719a44949664a1094205a99d95a0856758fed66d690105e
+
+# b^m = 1
+Result = PossiblyPrime
+W = 64561b8d9aa50340c3a01ccb3e6e17f5023513661c012be288f3900a3ca76890e67290b9560fa1d480f9d2aacccca581b5690636665f243fa13aff5d0bff12d3
+B = 1f5ff70d3d60671ebc5fbfca731898a04438053dbc3c841e6335f487e457d92d9efb5d506d5bef6872d58d12b9a41c950bfc38d12ed977c90eacdd6535b811a0
+
+# b^m = 1
+Result = PossiblyPrime
+W = 69c63fbf44df21b0ed0ee929a740c12d1f3f064da0dcd9d509f31fa45fa27d1a759ab5a9f6f1040d7ee90a0b1e68f779273c41ea1c1198fd547ff6bd70c7e787
+B = 5f7996a9bbfd8fd88e472220b70077bfdacdd63d88885134431f024c2acb7126827b174eb093eb5313f07bb5461de9b0feb7d77ca2c39c2a323a150f33ea525f
+
+# End of iteration
+Result = Composite
+W = 28cc3e08c44571c6dcb98a9ab8b4f3e2b16e1f884997d94a3188bcbb7f1b7cdaecdae8329c013ec8f75dc00004da0039943e4262cd080b16a42910102e00dddb
+B = 512061ab1c69931c2fa0bb89d8d09f3c9209230bf927ddd6fb6a72075f967ed3c4dbb5f437bf4d31ca7344782b22011ad56609dc19aed65319bababfc13dd7
+
+# End of iteration
+Result = Composite
+W = 4eeb7b4d371c45fe8586fee3b1efd792176b70f6cc2698dfa1dd028366626febe0199c3c5f77a5c3cad0057a04767383051d41965255d03681b2a37edad34a9b
+B = 4afc2e85f84017b3fd6967a227eb74c8297b40ea02733d9513bff9b3f01081963f25872f4254afc4e9321eea35b2a1e42eadb186fcc84f2f30f4a994350b93b8
+
+# End of iteration
+Result = Composite
+W = 8e35a959555dd2eb66c65cee3c264071d20671f159e1f9896f1d0ceb041905fcf053eacc189de317c3ee6f93901223cbf30d5b7ddbbdab981790e2f6397e6803
+B = 44c0153759309ec4e5b1e59d57c1b126545ef7ea302b6e43561df4d16068b922389d6924f01c945d9080d1f93a0732599bdedae72d6d590839dc0884dd860441
+
+
+# 0x6c1 = 1729 = 7 * 13 * 19 is a Fermat pseudoprime.
+
+# Found non-trivial square root
+Result = Composite
+W = 6c1
+B = b8
+
+# End of iteration
+Result = Composite
+W = 6c1
+B = 111
+
+# End of iteration
+Result = Composite
+W = 6c1
+B = 11d
+
+# Found non-trivial square root
+Result = Composite
+W = 6c1
+B = 19c
+
+# Found non-trivial square root
+Result = Composite
+W = 6c1
+B = 223
+
+# End of iteration
+Result = Composite
+W = 6c1
+B = 3aa
+
+# Found non-trivial square root
+Result = Composite
+W = 6c1
+B = 653
+
+
+# 1729 has a number of false witnesses.
+
+# b^m = 1
+Result = PossiblyPrime
+W = 6c1
+B = 78
+
+# b^m = 1
+Result = PossiblyPrime
+W = 6c1
+B = eb
+
+# b^m = w-1
+Result = PossiblyPrime
+W = 6c1
+B = 178
+
+# b^m = w-1
+Result = PossiblyPrime
+W = 6c1
+B = 178
+
+# b^m = w-1
+Result = PossiblyPrime
+W = 6c1
+B = 1aa
+
+# b^m = 1
+Result = PossiblyPrime
+W = 6c1
+B = 271
+
+# b^m = 1
+Result = PossiblyPrime
+W = 6c1
+B = 2b2
+
+
+# 1 and W-1 are always nonwitnesses.
+Result = PossiblyPrime
+W = 6c1
+B = 1
+
+Result = PossiblyPrime
+W = 6c1
+B = 6c0
+
+
+# https://kconrad.math.uconn.edu/blurbs/ugradnumthy/millerrabin.pdf, examples
+# 3.1 and 3.2 has a complete list of false witnesses for 65 = 0x41 and
+# 85 = 0x55.
+
+# b^m = 1
+Result = PossiblyPrime
+W = 41
+B = 1
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 41
+B = 8
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 41
+B = 12
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 41
+B = 2f
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 41
+B = 39
+
+# b^m = w-1
+Result = PossiblyPrime
+W = 41
+B = 40
+
+# b^m = 1
+Result = PossiblyPrime
+W = 55
+B = 1
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 55
+B = d
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 55
+B = 26
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 55
+B = 2f
+
+# Some b^(m*2^j) = w-1
+Result = PossiblyPrime
+W = 55
+B = 48
+
+# b^m = w-1
+Result = PossiblyPrime
+W = 55
+B = 54
+
+# Other witnesses for 65 and 85 will report composite:
+
+# Found non-trivial square root
+Result = Composite
+W = 41
+B = 2c
+
+# End of iteration
+Result = Composite
+W = 41
+B = 16
+
+# End of iteration
+Result = Composite
+W = 41
+B = 14
+
+# End of iteration
+Result = Composite
+W = 41
+B = 2
+
+# End of iteration
+Result = Composite
+W = 41
+B = 3a
+
+# End of iteration
+Result = Composite
+W = 55
+B = 40
+
+# End of iteration
+Result = Composite
+W = 55
+B = 7
+
+# End of iteration
+Result = Composite
+W = 55
+B = 23
+
+# End of iteration
+Result = Composite
+W = 55
+B = 2e
+
+# End of iteration
+Result = Composite
+W = 55
+B = 2a
diff --git a/crypto/fipsmodule/bn/prime.c b/crypto/fipsmodule/bn/prime.c
index c030c9e..1dd0870 100644
--- a/crypto/fipsmodule/bn/prime.c
+++ b/crypto/fipsmodule/bn/prime.c
@@ -594,6 +594,113 @@
   return bn_trial_division(&prime, bn) && !BN_is_word(bn, prime);
 }
 
+int bn_miller_rabin_init(BN_MILLER_RABIN *miller_rabin, const BN_MONT_CTX *mont,
+                         BN_CTX *ctx) {
+  // This function corresponds to steps 1 through 3 of FIPS 186-4, C.3.1.
+  const BIGNUM *w = &mont->N;
+  // Note we do not call |BN_CTX_start| in this function. We intentionally
+  // allocate values in the containing scope so they outlive this function.
+  miller_rabin->w1 = BN_CTX_get(ctx);
+  miller_rabin->m = BN_CTX_get(ctx);
+  miller_rabin->one_mont = BN_CTX_get(ctx);
+  miller_rabin->w1_mont = BN_CTX_get(ctx);
+  if (miller_rabin->w1 == NULL ||
+      miller_rabin->m == NULL ||
+      miller_rabin->one_mont == NULL ||
+      miller_rabin->w1_mont == NULL) {
+    return 0;
+  }
+
+  // See FIPS 186-4, C.3.1, steps 1 through 3.
+  if (!bn_usub_consttime(miller_rabin->w1, w, BN_value_one())) {
+    return 0;
+  }
+  miller_rabin->a = BN_count_low_zero_bits(miller_rabin->w1);
+  if (!bn_rshift_secret_shift(miller_rabin->m, miller_rabin->w1,
+                              miller_rabin->a, ctx)) {
+    return 0;
+  }
+  miller_rabin->w_bits = BN_num_bits(w);
+
+  // Precompute some values in Montgomery form.
+  if (!bn_one_to_montgomery(miller_rabin->one_mont, mont, ctx) ||
+      // w - 1 is -1 mod w, so we can compute it in the Montgomery domain, -R,
+      // with a subtraction. (|one_mont| cannot be zero.)
+      !bn_usub_consttime(miller_rabin->w1_mont, w, miller_rabin->one_mont)) {
+    return 0;
+  }
+
+  return 1;
+}
+
+int bn_miller_rabin_iteration(const BN_MILLER_RABIN *miller_rabin,
+                              int *out_is_possibly_prime, const BIGNUM *b,
+                              const BN_MONT_CTX *mont, BN_CTX *ctx) {
+  // This function corresponds to steps 4.3 through 4.5 of FIPS 186-4, C.3.1.
+  int ret = 0;
+  BN_CTX_start(ctx);
+
+  // Step 4.3. We use Montgomery-encoding for better performance and to avoid
+  // timing leaks.
+  const BIGNUM *w = &mont->N;
+  BIGNUM *z = BN_CTX_get(ctx);
+  if (z == NULL ||
+      !BN_mod_exp_mont_consttime(z, b, miller_rabin->m, w, ctx, mont) ||
+      !BN_to_montgomery(z, z, mont, ctx)) {
+    goto err;
+  }
+
+  // loop_done is all ones if the loop has completed and all zeros otherwise.
+  crypto_word_t loop_done = 0;
+  // is_possibly_prime is all ones if we have determined |b| is not a composite
+  // witness for |w|. This is equivalent to going to step 4.7 in the original
+  // algorithm.
+  crypto_word_t is_possibly_prime = 0;
+
+  // Step 4.4. If z = 1 or z = w-1, b is not a composite witness and w is still
+  // possibly prime.
+  loop_done = BN_equal_consttime(z, miller_rabin->one_mont) |
+              BN_equal_consttime(z, miller_rabin->w1_mont);
+  loop_done = 0 - loop_done;   // Make it all zeros or all ones.
+  is_possibly_prime = loop_done;  // Go to step 4.7 if |loop_done|.
+
+  // Step 4.5.
+  //
+  // To avoid leaking |a|, we run the loop to |w_bits| and mask off all
+  // iterations once |j| = |a|.
+  for (int j = 1; j < miller_rabin->w_bits; j++) {
+    loop_done |= constant_time_eq_int(j, miller_rabin->a);
+
+    // Step 4.5.1.
+    if (!BN_mod_mul_montgomery(z, z, z, mont, ctx)) {
+      goto err;
+    }
+
+    // Step 4.5.2. If z = w-1 and the loop is not done, this is not a composite
+    // witness.
+    crypto_word_t z_is_w1_mont =
+        BN_equal_consttime(z, miller_rabin->w1_mont) & ~loop_done;
+    z_is_w1_mont = 0 - z_is_w1_mont;  // Make it all zeros or all ones.
+    loop_done |= z_is_w1_mont;
+    is_possibly_prime |= z_is_w1_mont;  // Go to step 4.7 if |z_is_w1_mont|.
+
+    // Step 4.5.3. If z = 1 and the loop is not done, the previous value of z
+    // was not -1. There are no non-trivial square roots of 1 modulo a prime, so
+    // w is composite and we may exit in variable time.
+    if (BN_equal_consttime(z, miller_rabin->one_mont) & ~loop_done) {
+      assert(!is_possibly_prime);
+      break;
+    }
+  }
+
+  *out_is_possibly_prime = is_possibly_prime & 1;
+  ret = 1;
+
+err:
+  BN_CTX_end(ctx);
+  return ret;
+}
+
 int BN_primality_test(int *out_is_probably_prime, const BIGNUM *w,
                       int iterations, BN_CTX *ctx, int do_trial_division,
                       BN_GENCB *cb) {
@@ -646,36 +753,13 @@
 
   // See C.3.1 from FIPS 186-4.
   int ret = 0;
-  BN_MONT_CTX *mont = NULL;
   BN_CTX_start(ctx);
-  BIGNUM *w1 = BN_CTX_get(ctx);
-  if (w1 == NULL ||
-      !bn_usub_consttime(w1, w, BN_value_one())) {
-    goto err;
-  }
-
-  // Write w1 as m * 2^a (Steps 1 and 2).
-  int w_len = BN_num_bits(w);
-  int a = BN_count_low_zero_bits(w1);
-  BIGNUM *m = BN_CTX_get(ctx);
-  if (m == NULL ||
-      !bn_rshift_secret_shift(m, w1, a, ctx)) {
-    goto err;
-  }
-
-  // Montgomery setup for computations mod w. Additionally, compute 1 and w - 1
-  // in the Montgomery domain for later comparisons.
   BIGNUM *b = BN_CTX_get(ctx);
-  BIGNUM *z = BN_CTX_get(ctx);
-  BIGNUM *one_mont = BN_CTX_get(ctx);
-  BIGNUM *w1_mont = BN_CTX_get(ctx);
-  mont = BN_MONT_CTX_new_consttime(w, ctx);
-  if (b == NULL || z == NULL || one_mont == NULL || w1_mont == NULL ||
-      mont == NULL ||
-      !bn_one_to_montgomery(one_mont, mont, ctx) ||
-      // w - 1 is -1 mod w, so we can compute it in the Montgomery domain, -R,
-      // with a subtraction. (|one_mont| cannot be zero.)
-      !bn_usub_consttime(w1_mont, w, one_mont)) {
+  BN_MONT_CTX *mont = BN_MONT_CTX_new_consttime(w, ctx);
+  BN_MILLER_RABIN miller_rabin;
+  if (b == NULL || mont == NULL ||
+      // Steps 1-3.
+      !bn_miller_rabin_init(&miller_rabin, mont, ctx)) {
     goto err;
   }
 
@@ -713,64 +797,22 @@
   for (int i = 1; (i <= BN_PRIME_CHECKS_BLINDED) |
                   constant_time_lt_w(uniform_iterations, iterations);
        i++) {
+    // Step 4.1-4.2
     int is_uniform;
-    if (// Step 4.1-4.2
-        !bn_rand_secret_range(b, &is_uniform, 2, w1) ||
-        // Step 4.3
-        !BN_mod_exp_mont_consttime(z, b, m, w, ctx, mont)) {
-      goto err;
+    if (!bn_rand_secret_range(b, &is_uniform, 2, miller_rabin.w1)) {
+        goto err;
     }
     uniform_iterations += is_uniform;
 
-    // loop_done is all ones if the loop has completed and all zeros otherwise.
-    crypto_word_t loop_done = 0;
-    // next_iteration is all ones if we should continue to the next iteration
-    // (|b| is not a composite witness for |w|). This is equivalent to going to
-    // step 4.7 in the original algorithm.
-    crypto_word_t next_iteration = 0;
-
-    // Step 4.4. If z = 1 or z = w-1, mask off the loop and continue to the next
-    // iteration (go to step 4.7).
-    loop_done = BN_equal_consttime(z, BN_value_one()) |
-                BN_equal_consttime(z, w1);
-    loop_done = 0 - loop_done;   // Make it all zeros or all ones.
-    next_iteration = loop_done;  // Go to step 4.7 if |loop_done|.
-
-    // Step 4.5. We use Montgomery-encoding for better performance and to avoid
-    // timing leaks.
-    if (!BN_to_montgomery(z, z, mont, ctx)) {
+    // Steps 4.3-4.5
+    int is_possibly_prime = 0;
+    if (!bn_miller_rabin_iteration(&miller_rabin, &is_possibly_prime, b, mont,
+                                   ctx)) {
       goto err;
     }
 
-    // To avoid leaking |a|, we run the loop to |w_len| and mask off all
-    // iterations once |j| = |a|.
-    for (int j = 1; j < w_len; j++) {
-      loop_done |= constant_time_eq_int(j, a);
-
-      // Step 4.5.1.
-      if (!BN_mod_mul_montgomery(z, z, z, mont, ctx)) {
-        goto err;
-      }
-
-      // Step 4.5.2. If z = w-1 and the loop is not done, run through the next
-      // iteration.
-      crypto_word_t z_is_w1_mont = BN_equal_consttime(z, w1_mont) & ~loop_done;
-      z_is_w1_mont = 0 - z_is_w1_mont;  // Make it all zeros or all ones.
-      loop_done |= z_is_w1_mont;
-      next_iteration |= z_is_w1_mont;  // Go to step 4.7 if |z_is_w1_mont|.
-
-      // Step 4.5.3. If z = 1 and the loop is not done, w is composite and we
-      // may exit in variable time.
-      if (BN_equal_consttime(z, one_mont) & ~loop_done) {
-        assert(!next_iteration);
-        break;
-      }
-    }
-
-    if (!next_iteration) {
+    if (!is_possibly_prime) {
       // Step 4.6. We did not see z = w-1 before z = 1, so w must be composite.
-      // (For any prime, the value of z immediately preceding 1 must be -1.
-      // There are no non-trivial square roots of 1 modulo a prime.)
       *out_is_probably_prime = 0;
       ret = 1;
       goto err;
diff --git a/sources.cmake b/sources.cmake
index 8dc65e6..db14456 100644
--- a/sources.cmake
+++ b/sources.cmake
@@ -47,6 +47,7 @@
   crypto/evp/scrypt_tests.txt
   crypto/fipsmodule/aes/aes_tests.txt
   crypto/fipsmodule/bn/bn_tests.txt
+  crypto/fipsmodule/bn/miller_rabin_tests.txt
   crypto/fipsmodule/ec/ec_scalar_base_mult_tests.txt
   crypto/fipsmodule/ec/p256-x86_64_tests.txt
   crypto/fipsmodule/ecdsa/ecdsa_sign_tests.txt