From 2e451f2e0bdc954b1043424ce4e2d20fcd38d897 Mon Sep 17 00:00:00 2001
From: Alex Dale <sigquit@google.com>
Date: Fri, 1 Apr 2022 20:31:35 -0700
Subject: [PATCH] Added OTA Keybox stress test.

[ Merge of http://go/wvgerrit/150130 ]
[ Cherry-pick of http://ag/17549231 ]

Created a new test for stressing OEMCrypto's ability to generate
OTA Keybox provisioning requests.  This forces the TA to retrieve
keys from KM, generate certificate and sign the request.  This is
intended to find any unexpected system degradation within the
device's TA(s).

Bug: 227542259
Test: oemcrypto_test
Change-Id: Ib34f2f801a7fe74ca67aa0a16f68f9ae326de24e
(cherry picked from commit 8cab20c0c57977483d0b969bcd70f390050048a4)
---
 .../oemcrypto/test/ota_keybox_test.cpp        | 59 +++++++++++++++++--
 1 file changed, 53 insertions(+), 6 deletions(-)

diff --git a/libwvdrmengine/oemcrypto/test/ota_keybox_test.cpp b/libwvdrmengine/oemcrypto/test/ota_keybox_test.cpp
index 1ead628d..6ef7588c 100644
--- a/libwvdrmengine/oemcrypto/test/ota_keybox_test.cpp
+++ b/libwvdrmengine/oemcrypto/test/ota_keybox_test.cpp
@@ -135,6 +135,11 @@ static const uint8_t TestKeyPKCS8[] = {
     0x13, 0xf3, 0xa7, 0xdb, 0xdb, 0x5d, 0x89, 0x2d, 0xd7, 0x02, 0x96, 0xaf,
     0xeb, 0x72, 0x8d, 0xd5, 0x56, 0x3b, 0x3a};
 
+constexpr size_t kInitialOtaKeyboxRequestSize = 8 * 1024;  // 8 kB.
+// TODO(fredgc): Make sure that partners can use a test cert when
+// |use_test_key| parameter of OEMCrypto_GenerateOTARequest() is true.
+constexpr uint32_t kUseTestKey = 1;
+
 // TODO(fredgc): duplicate code. Move to common util package.
 // Return a printable string from data. If all the characters are printable,
 // then just use the string. Otherwise, convert to hex.
@@ -192,9 +197,7 @@ TEST_F(OTAKeyboxProvisioningTest, BasicTest) {
   }
   cout << "             "
        << "OTA Keybox functions supported. Device needs provisioning." << endl;
-  // TODO(fredgc): Make sure that partners can use a test cert when use_test_key
-  // is true.
-  constexpr uint32_t use_test_key = 1;
+
   size_t request_length = 0;
   std::vector<uint8_t> request;
 
@@ -203,12 +206,12 @@ TEST_F(OTAKeyboxProvisioningTest, BasicTest) {
 
   ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
             OEMCrypto_GenerateOTARequest(session_id, nullptr, &request_length,
-                                         use_test_key));
+                                         kUseTestKey));
   ASSERT_NE(request_length, 0u);
   request.resize(request_length);
   ASSERT_EQ(OEMCrypto_SUCCESS,
             OEMCrypto_GenerateOTARequest(session_id, request.data(),
-                                         &request_length, use_test_key));
+                                         &request_length, kUseTestKey));
   ASSERT_GT(request_length, kMinimumRequestLength);
   request.resize(request_length);
   // First 16 bytes should match the label for Option 1 or 2.
@@ -292,9 +295,53 @@ TEST_F(OTAKeyboxProvisioningTest, BasicTest) {
   // Finally, send the response back to the device.
   EXPECT_EQ(OEMCrypto_SUCCESS,
             OEMCrypto_ProcessOTAKeybox(session_id, response.data(),
-                                       response.size(), use_test_key));
+                                       response.size(), kUseTestKey));
   ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_CloseSession(session_id));
   // After installation, the keybox should be valid.
   EXPECT_EQ(OEMCrypto_IsKeyboxValid(), OEMCrypto_SUCCESS);
 }
+
+TEST_F(OTAKeyboxProvisioningTest, StressTest) {
+  static constexpr size_t kStressRequestCount = 1000;
+  const OEMCrypto_ProvisioningMethod method = OEMCrypto_GetProvisioningMethod();
+  if (method != OEMCrypto_Keybox) {
+    std::cout << "Skipping for non keybox-based devices: method = "
+              << ProvisioningMethodName(method) << std::endl;
+    return;
+  }
+  OEMCrypto_SESSION session_id;
+  ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_OpenSession(&session_id))
+      << "Failed to open OTA keybox session";
+  // First request is to ensure operation is supported.
+  std::vector<uint8_t> request(kInitialOtaKeyboxRequestSize);
+  size_t request_size = request.size();
+  OEMCryptoResult result = OEMCrypto_GenerateOTARequest(
+      session_id, request.data(), &request_size, kUseTestKey);
+  if (result == OEMCrypto_ERROR_NOT_IMPLEMENTED) {
+    // Not a failure, just skip the test.
+    OEMCrypto_CloseSession(session_id);
+    std::cout << "Skipping for device which do not support "
+              << "OTA keybox provisioning" << std::endl;
+    return;
+  }
+
+  for (size_t i = 0; i < kStressRequestCount; i++) {
+    request.resize(kInitialOtaKeyboxRequestSize);
+    request_size = request.size();
+    result = OEMCrypto_GenerateOTARequest(session_id, request.data(),
+                                          &request_size, kUseTestKey);
+    if (result == OEMCrypto_ERROR_SHORT_BUFFER) {
+      request.resize(request_size);
+      result = OEMCrypto_GenerateOTARequest(session_id, request.data(),
+                                            &request_size, kUseTestKey);
+    }
+    if (result != OEMCrypto_SUCCESS) {
+      OEMCrypto_CloseSession(session_id);
+      std::cout << "Failed on attempt " << (i + 1) << std::endl;
+      FAIL() << "Failed to generate request: i = " << i;
+      return;
+    }
+  }
+  OEMCrypto_CloseSession(session_id);
+}
 }  // namespace wvoec