From 4ae4479b4f264bd42da9620366cd1cde9b6da60e Mon Sep 17 00:00:00 2001
From: Justin W Smith <103147162+justsmth@users.noreply.github.com>
Date: Thu, 2 Nov 2023 15:32:32 -0400
Subject: [PATCH] Add BNAssertTest

---
 crypto/CMakeLists.txt                  |  1 +
 crypto/fipsmodule/bn/bn_assert_test.cc | 77 ++++++++++++++++++++++++++
 crypto/fipsmodule/bn/internal.h        |  2 +-
 3 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 crypto/fipsmodule/bn/bn_assert_test.cc

diff --git a/crypto/CMakeLists.txt b/crypto/CMakeLists.txt
index 49227de7d8..3cea99a988 100644
--- a/crypto/CMakeLists.txt
+++ b/crypto/CMakeLists.txt
@@ -719,6 +719,7 @@ if(BUILD_TESTING)
     evp_extra/scrypt_test.cc
     fipsmodule/aes/aes_test.cc
     fipsmodule/bn/bn_test.cc
+    fipsmodule/bn/bn_assert_test.cc
     fipsmodule/cmac/cmac_test.cc
     fipsmodule/ec/ec_test.cc
     fipsmodule/ec/p256-nistz_test.cc
diff --git a/crypto/fipsmodule/bn/bn_assert_test.cc b/crypto/fipsmodule/bn/bn_assert_test.cc
new file mode 100644
index 0000000000..758fd82188
--- /dev/null
+++ b/crypto/fipsmodule/bn/bn_assert_test.cc
@@ -0,0 +1,77 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR ISC
+//
+
+#include <openssl/bn.h>
+#include <openssl/rand.h>
+#include "./internal.h"
+
+#include <gtest/gtest.h>
+
+TEST(BNAssertTest, Assert_fits_in_bytes_large) {
+// TODO: Update Android test harness
+#if !defined(NDEBUG) && !defined(OPENSSL_ANDROID)
+  bssl::UniquePtr<BIGNUM> x(BN_new());
+  uint8_t input[255];
+  OPENSSL_memset(input, 0, sizeof(input));
+  input[0] = 0xaa;
+  input[1] = 0x01;
+  input[254] = 0x01;
+  ASSERT_TRUE(BN_le2bn(input, sizeof(input), x.get()));
+  for (size_t i = 255; i < 260; i++) {
+    bn_assert_fits_in_bytes(x.get(), i);
+  }
+  for (size_t i = 247; i < 255; i++) {
+    EXPECT_DEATH_IF_SUPPORTED(bn_assert_fits_in_bytes(x.get(), i), "");
+  }
+#endif
+}
+
+TEST(BNAssertTest, Assert_fits_in_bytes_small) {
+#if !defined(NDEBUG) && !defined(OPENSSL_ANDROID)
+  bssl::UniquePtr<BIGNUM> x(BN_new());
+  uint8_t input[8];
+  OPENSSL_memset(input, 0, sizeof(input));
+  input[0] = 0xaa;
+  input[1] = 0xbb;
+  input[2] = 0xcc;
+  ASSERT_TRUE(BN_le2bn(input, sizeof(input), x.get()));
+  for (size_t i = 3; i < 10; i++) {
+    bn_assert_fits_in_bytes(x.get(), i);
+  }
+  for (size_t i = 0; i < 3; i++) {
+    EXPECT_DEATH_IF_SUPPORTED(bn_assert_fits_in_bytes(x.get(), i), "");
+  }
+#endif
+}
+
+TEST(BNAssertTest, Assert_fits_in_bytes_zero) {
+#if !defined(NDEBUG) && !defined(OPENSSL_ANDROID)
+  bssl::UniquePtr<BIGNUM> x(BN_new());
+  uint8_t input[8];
+  OPENSSL_memset(input, 0, sizeof(input));
+  ASSERT_TRUE(BN_le2bn(input, sizeof(input), x.get()));
+
+  for (size_t i = 0; i < 10; i++) {
+    bn_assert_fits_in_bytes(x.get(), i);
+  }
+#endif
+}
+
+TEST(BNAssertTest, Assert_fits_in_bytes_boundary) {
+#if !defined(NDEBUG) && !defined(OPENSSL_ANDROID)
+  bssl::UniquePtr<BIGNUM> x(BN_new());
+  uint8_t input[8];
+  OPENSSL_memset(input, 0, sizeof(input));
+  for (size_t i = 0; i < sizeof(input); i++) {
+    input[i] = i * (i + 1) & 0xff;
+  }
+  ASSERT_TRUE(BN_le2bn(input, sizeof(input), x.get()));
+  for (size_t i = 8; i < 18; i++) {
+    bn_assert_fits_in_bytes(x.get(), i);
+  }
+  for (size_t i = 0; i < 8; i++) {
+    EXPECT_DEATH_IF_SUPPORTED(bn_assert_fits_in_bytes(x.get(), i), "");
+  }
+#endif
+}
diff --git a/crypto/fipsmodule/bn/internal.h b/crypto/fipsmodule/bn/internal.h
index 541c585d9b..c545a76e84 100644
--- a/crypto/fipsmodule/bn/internal.h
+++ b/crypto/fipsmodule/bn/internal.h
@@ -271,7 +271,7 @@ int bn_copy_words(BN_ULONG *out, size_t num, const BIGNUM *bn);
 // no-op in release builds, but triggers an assert in debug builds, and
 // declassifies all bytes which are therefore known to be zero in constant-time
 // validation.
-void bn_assert_fits_in_bytes(const BIGNUM *bn, size_t num);
+OPENSSL_EXPORT void bn_assert_fits_in_bytes(const BIGNUM *bn, size_t num);
 
 // bn_mul_add_words multiples |ap| by |w|, adds the result to |rp|, and places
 // the result in |rp|. |ap| and |rp| must both be |num| words long. It returns