From b7965436baf8e448e0f3d30572e280d6bd0ba5ae Mon Sep 17 00:00:00 2001
From: Kevin Jacobs <kjacobs@mozilla.com>
Date: Wed, 19 Feb 2020 22:19:11 +0000
Subject: [PATCH] Bug 1609751 - Additional tests for mp_comba r=mt

Verify that when clamping, the upper 4 bytes of an `mp_digit` is checked.

Differential Revision: https://phabricator.services.mozilla.com/D58576

--HG--
extra : moz-landing-system : lando
---
 gtests/freebl_gtest/mpi_unittest.cc | 49 +++++++++++++++++++++++++++++
 lib/freebl/mpi/mplogic.h            |  3 ++
 2 files changed, 52 insertions(+)

diff --git a/gtests/freebl_gtest/mpi_unittest.cc b/gtests/freebl_gtest/mpi_unittest.cc
index 56b7454dc9..eabac6ed03 100644
--- a/gtests/freebl_gtest/mpi_unittest.cc
+++ b/gtests/freebl_gtest/mpi_unittest.cc
@@ -13,6 +13,7 @@
 #include <mach/mach.h>
 #endif
 
+#include "mplogic.h"
 #include "mpi.h"
 namespace nss_test {
 
@@ -224,6 +225,54 @@ TEST_F(MPITest, MpiFixlenOctetsTooSmall) {
   }
 }
 
+TEST_F(MPITest, MpiSqrMulClamp) {
+  mp_int a, r, expect;
+  MP_DIGITS(&a) = 0;
+  MP_DIGITS(&r) = 0;
+  MP_DIGITS(&expect) = 0;
+
+  // Comba32 result is 64 mp_digits. *=2 as this is an ascii representation.
+  std::string expect_str((64 * sizeof(mp_digit)) * 2, '0');
+
+  // Set second-highest bit (0x80...^2 == 0x4000...)
+  expect_str.replace(0, 1, "4", 1);
+
+  // Test 32, 16, 8, and 4-1 mp_digit values. 32-4 (powers of two) use the comba
+  // assembly implementation, if enabled and supported. 3-1 use non-comba.
+  int n_digits = 32;
+  while (n_digits > 0) {
+    ASSERT_EQ(MP_OKAY, mp_init(&r));
+    ASSERT_EQ(MP_OKAY, mp_init(&a));
+    ASSERT_EQ(MP_OKAY, mp_init(&expect));
+    ASSERT_EQ(MP_OKAY, mp_read_radix(&expect, expect_str.c_str(), 16));
+
+    ASSERT_EQ(MP_OKAY, mp_set_int(&a, 1));
+    ASSERT_EQ(MP_OKAY, mpl_lsh(&a, &a, (n_digits * sizeof(mp_digit) * 8) - 1));
+
+    ASSERT_EQ(MP_OKAY, mp_sqr(&a, &r));
+    EXPECT_EQ(MP_USED(&expect), MP_USED(&r));
+    EXPECT_EQ(0, mp_cmp(&r, &expect));
+    mp_clear(&r);
+
+    // Take the mul path...
+    ASSERT_EQ(MP_OKAY, mp_init(&r));
+    ASSERT_EQ(MP_OKAY, mp_mul(&a, &a, &r));
+    EXPECT_EQ(MP_USED(&expect), MP_USED(&r));
+    EXPECT_EQ(0, mp_cmp(&r, &expect));
+
+    mp_clear(&a);
+    mp_clear(&r);
+    mp_clear(&expect);
+
+    // Once we're down to 4, check non-powers of two.
+    int sub = n_digits > 4 ? n_digits / 2 : 1;
+    n_digits -= sub;
+
+    // "Shift right" the string (to avoid mutating |expect_str| with MPI).
+    expect_str.resize(expect_str.size() - 2 * 2 * sizeof(mp_digit) * sub);
+  }
+}
+
 // This test is slow. Disable it by default so we can run these tests on CI.
 class DISABLED_MPITest : public ::testing::Test {};
 
diff --git a/lib/freebl/mpi/mplogic.h b/lib/freebl/mpi/mplogic.h
index a4a6b7735d..6a0f67c503 100644
--- a/lib/freebl/mpi/mplogic.h
+++ b/lib/freebl/mpi/mplogic.h
@@ -11,6 +11,7 @@
 #define _H_MPLOGIC_
 
 #include "mpi.h"
+SEC_BEGIN_PROTOS
 
 /*
   The logical operations treat an mp_int as if it were a bit vector,
@@ -49,4 +50,6 @@ mp_err mpl_get_bit(const mp_int *a, mp_size bitNum);
 mp_err mpl_get_bits(const mp_int *a, mp_size lsbNum, mp_size numBits);
 mp_size mpl_significant_bits(const mp_int *a);
 
+SEC_END_PROTOS
+
 #endif /* end _H_MPLOGIC_ */