Improve logic and testing for UBE

crypto/fipsmodule/rand/new_rand_test.cc

@@ -26,7 +26,7 @@ class newRandTest : public ::testing::Test {
     }
 
     void TearDown() override {
-      allow_mocked_ube_detection_FOR_TESTING(0);
+      disable_mocked_ube_detection_FOR_TESTING();
     }
 
   protected:
@@ -36,7 +36,7 @@ class newRandTest : public ::testing::Test {
 
     void allowMockedUbeIfNecessary(void) {
       if (ube_detection_supported_) {
-        allow_mocked_ube_detection_FOR_TESTING(1);
+        allow_mocked_ube_detection_FOR_TESTING();
       }
     }
 

crypto/ube/internal.h

@@ -31,12 +31,31 @@ extern "C" {
 // entries will immediately return.
 OPENSSL_EXPORT int CRYPTO_get_ube_generation_number(uint64_t *current_generation_number);
 
-// TODO
-// Temporary overrides. Replace with something better. Used atm to test
-// implementation during development.
+// set_fork_generation_number_FOR_TESTING sets the fork generation number to the
+// value |fork_gn|. This value will be the fork generation value used by the UBE
+// logic, overriding the generation number from the real fork detection.
+// |allow_mocked_ube_detection_FOR_TESTING| must have been invoked
+// (once per-process) to allow mocking the fork generation number.
 OPENSSL_EXPORT void set_fork_generation_number_FOR_TESTING(uint64_t fork_gn);
+
+// set_snapsafe_generation_number_FOR_TESTING sets the snapsafe generation
+// number to the value |snapsafe_gn|. This value will be the snapsafe generation
+// value used by the UBE logic, overriding the generation number from the real
+// snapsafe detection.
+// |allow_mocked_ube_detection_FOR_TESTING| must have been invoked (once
+// per-process) to allow mocking the snapsafe generation number.
 OPENSSL_EXPORT void set_snapsafe_generation_number_FOR_TESTING(uint32_t snapsafe_gn);
-OPENSSL_EXPORT void allow_mocked_ube_detection_FOR_TESTING(uint8_t allow);
+
+// allow_mocked_ube_detection_FOR_TESTING allows mocking UBE detection even
+// though real detection is not available. This function must be called in
+// test code when mocking the generation numbers, once per-process. Mocking is
+// process-global i.e. all threads will read the same mocked value.
+OPENSSL_EXPORT void allow_mocked_ube_detection_FOR_TESTING(void);
+
+// disable_mocked_ube_detection_FOR_TESTING disables mocking UBE detection. It
+// also resets any mocked values to default values (0). This function should be
+// invoked when exiting testing.
+OPENSSL_EXPORT void disable_mocked_ube_detection_FOR_TESTING(void);
 
 #if defined(__cplusplus)
 }  // extern C

crypto/ube/ube.c

@@ -39,27 +39,27 @@ static struct CRYPTO_STATIC_MUTEX ube_lock = CRYPTO_STATIC_MUTEX_INIT;
 static uint8_t ube_detection_unavailable = 0;
 static uint8_t allow_mocked_detection = 0;
 
-static uint64_t testing_fork_generation_number = 0;
+static uint64_t override_fork_generation_number = 0;
 void set_fork_generation_number_FOR_TESTING(uint64_t fork_gn) {
-  testing_fork_generation_number = fork_gn;
+  override_fork_generation_number = fork_gn;
 }
-static uint32_t testing_snapsafe_generation_number = 0;
+static uint32_t override_snapsafe_generation_number = 0;
 void set_snapsafe_generation_number_FOR_TESTING(uint32_t snapsafe_gn) {
-  testing_snapsafe_generation_number = snapsafe_gn;
+  override_snapsafe_generation_number = snapsafe_gn;
 }
 
 static int get_snapsafe_generation_number(uint32_t *gn) {
-  if (testing_snapsafe_generation_number != 0) {
-    *gn = testing_snapsafe_generation_number;
+  if (allow_mocked_detection == 1) {
+    *gn = override_snapsafe_generation_number;
     return 1;
   }
 
   return CRYPTO_get_snapsafe_generation(gn);
 }
 
 static int get_fork_generation_number(uint64_t *gn) {
-  if (testing_fork_generation_number != 0) {
-    *gn = testing_fork_generation_number;
+  if (allow_mocked_detection == 1) {
+    *gn = override_fork_generation_number;
     return 1;
   }
 
@@ -187,7 +187,7 @@ int CRYPTO_get_ube_generation_number(uint64_t *current_generation_number) {
   // must be done after attempting to initialize the UBE state. Because
   // initialization might fail and we can short-circuit here.
   if (ube_detection_unavailable == 1 &&
-      allow_mocked_detection == 0) {
+      allow_mocked_detection != 1) {
     return 0;
   }
 
@@ -242,6 +242,12 @@ int CRYPTO_get_ube_generation_number(uint64_t *current_generation_number) {
   return 1;
 }
 
-void allow_mocked_ube_detection_FOR_TESTING(uint8_t allow) {
-  allow_mocked_detection = allow;
+void allow_mocked_ube_detection_FOR_TESTING(void) {
+  allow_mocked_detection = 1;
+}
+
+void disable_mocked_ube_detection_FOR_TESTING(void) {
+  allow_mocked_detection = 0;
+  set_fork_generation_number_FOR_TESTING(0);
+  set_snapsafe_generation_number_FOR_TESTING(0);
 }

crypto/ube/ube_test.cc

@@ -6,7 +6,33 @@
 #include "internal.h"
 #include "../test/test_util.h"
 
-TEST(Ube, BasicTests) {
+class ubeTest : public ::testing::Test {
+  public:
+    void SetUp() override {
+      uint64_t current_generation_number = 0;
+      ube_detection_supported_ = CRYPTO_get_ube_generation_number(
+                                  &current_generation_number);
+    }
+
+    void TearDown() override {
+      disable_mocked_ube_detection_FOR_TESTING();
+    }
+
+  protected:
+    bool UbeIsSupported(void) {
+      return ube_detection_supported_;
+    }
+
+    void allowMockedUbeIfNecessary(void) {
+      if (ube_detection_supported_) {
+        allow_mocked_ube_detection_FOR_TESTING();
+      }
+    }
+
+    bool ube_detection_supported_ = false;
+};
+
+TEST_F(ubeTest, BasicTests) {
   uint64_t generation_number = 0;
   if (CRYPTO_get_ube_generation_number(&generation_number) == 0) {
     // In this case, UBE detection is disabled, so just return
@@ -28,20 +54,26 @@ TEST(Ube, BasicTests) {
   ASSERT_EQ(current_generation_number, generation_number);
 }
 
-TEST(Ube, MockedMethodTests) {
+TEST_F(ubeTest, MockedMethodTests) {
+
+  allowMockedUbeIfNecessary();
+
   uint64_t generation_number = 0;
   uint64_t cached_generation_number = 0;
-  if (CRYPTO_get_ube_generation_number(&generation_number) == 0) {
-    // In this case, UBE detection is disabled, so just return
-    // successfully. The short-circuit feature means we can't mock detection
-    // methods.
-    return;
-  }
 
-  // The fork generation number is initially 1. Use 10 because it's larger...
-  // Configuring specific values must change later on, since we might have tests
-  // running concurrently.
+  // First test incrementing the fork generation number. Pick a starting point
+  // and get initial UBE generation number
   set_fork_generation_number_FOR_TESTING(10);
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+
+  // Should be stable again.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number);
+
+  // Mock a process fork. We used 10 before. Hence, 11 should work.
+  set_fork_generation_number_FOR_TESTING(11);
 
   // Generation number should have incremented once.
   cached_generation_number = generation_number;
@@ -55,8 +87,8 @@ TEST(Ube, MockedMethodTests) {
   ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
   ASSERT_EQ(generation_number, cached_generation_number);
 
-  // Mock another process fork. We used 10 before. Hence, 11 should work.
-  set_fork_generation_number_FOR_TESTING(11);
+  // Mock another process fork with higher increment.
+  set_fork_generation_number_FOR_TESTING(13);
 
   // Generation number should have incremented once.
   cached_generation_number = generation_number;
@@ -70,8 +102,8 @@ TEST(Ube, MockedMethodTests) {
   ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
   ASSERT_EQ(generation_number, cached_generation_number);
 
-  // The snapsafe generation number is initially 1. Again use 10.
-  set_snapsafe_generation_number_FOR_TESTING(10);
+  // Mock another fork event but with a strictly smaller value.
+  set_snapsafe_generation_number_FOR_TESTING(5);
 
   // Generation number should have incremented once.
   cached_generation_number = generation_number;
@@ -85,8 +117,49 @@ TEST(Ube, MockedMethodTests) {
   ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
   ASSERT_EQ(generation_number, cached_generation_number);
 
-  // Mock another snapsafe event. We used 10 before. Hence, 11 should work.
-  set_snapsafe_generation_number_FOR_TESTING(11);
+  // Now test incrementing the snapsafe generation number. Before we used 10 as
+  // the default value, use a smaller number now, say 5.
+  set_snapsafe_generation_number_FOR_TESTING(5);
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+
+  // Should be stable.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number);
+
+  // Mock a snapsafe event. We used 5 before. Hence, 6 should work.
+  set_snapsafe_generation_number_FOR_TESTING(6);
+
+  // Generation number should have incremented once.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number + 1);
+
+  // Should be stable again.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number);
+
+  // Mock another snapsafe event with higher increment.
+  set_snapsafe_generation_number_FOR_TESTING(8);
+
+  // Generation number should have incremented once.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number + 1);
+
+  // Should be stable again.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number);
+
+  // Mock another snapsafe event but with a strictly smaller value.
+  set_snapsafe_generation_number_FOR_TESTING(1);
 
   // Generation number should have incremented once.
   cached_generation_number = generation_number;
@@ -102,23 +175,53 @@ TEST(Ube, MockedMethodTests) {
 
   // Now try to increment both fork and snapsafe generation numbers. We expect
   // to see one increment in the ube generation number and then stability.
-  set_snapsafe_generation_number_FOR_TESTING(20);
   set_fork_generation_number_FOR_TESTING(20);
+  set_snapsafe_generation_number_FOR_TESTING(20);
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+
+  // And that it's now stable.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number);
+
+  // Increment both, this time using higher increments than 1 and different
+  // increment values.
+  set_fork_generation_number_FOR_TESTING(24);
+  set_snapsafe_generation_number_FOR_TESTING(23);
 
-  // Check that ube generation number incremented once.
+  // Generation number should have incremented once.
   cached_generation_number = generation_number;
   generation_number = 0;
   ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
   ASSERT_EQ(generation_number, cached_generation_number + 1);
 
-  // And that it's now stable.
+  // Should be stable again.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number);
+
+  // Try strictly smaller values. These are values previously used as fork and
+  // snapsafe generation number. It should still result in a UBE generation
+  // number increment.
+  set_fork_generation_number_FOR_TESTING(1);
+  set_snapsafe_generation_number_FOR_TESTING(10);
+
+  // Generation number should have incremented once.
+  cached_generation_number = generation_number;
+  generation_number = 0;
+  ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
+  ASSERT_EQ(generation_number, cached_generation_number + 1);
+
+  // Should be stable again.
   cached_generation_number = generation_number;
   generation_number = 0;
   ASSERT_TRUE(CRYPTO_get_ube_generation_number(&generation_number));
   ASSERT_EQ(generation_number, cached_generation_number);
 }
 
-TEST(Ube, ExpectedSupportTests) {
+TEST_F(ubeTest, ExpectedSupportTests) {
   uint64_t generation_number = 0;
   // Operating systems where we expect UBE detection to be enabled.
   if (osIsAmazonLinux()) {