From 458f4f68c4f1e5f9661c7397ef72a886de58295a Mon Sep 17 00:00:00 2001
From: Alex Dale <sigquit@google.com>
Date: Fri, 19 Aug 2022 18:45:33 -0700
Subject: [PATCH] Allow running tests without test keybox.

[ CDM CL http://go/wvgerrit/156286 ]

Local test scripts can now run the CDM unit tests without using a test
keybox / test RSA key by setting an environment variable:

  DISABLE_TEST_KEYBOX=yes

The default behavior will continue to be that the unit tests will load
a test keybox / test RSA key.

This is to enable testing different provisioning features when the
device uses keyboxes for its ROT.  Provisioning methods like 3.0 and
4.0 were never forcibly required to load a test ROT.

This is change is required for the ability to test ECC-enabled
keyboxes.

Bug: 135283522
Test: run_x86_64_tests and build_and_run_all_unit_tests.sh
Change-Id: I2a85be541deec3f60ab8fa82452a4d4cafaed339
---
 libwvdrmengine/cdm/core/test/test_base.cpp | 44 +++++++++++++++++++---
 libwvdrmengine/cdm/core/test/test_base.h   | 16 ++++++++
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/libwvdrmengine/cdm/core/test/test_base.cpp b/libwvdrmengine/cdm/core/test/test_base.cpp
index 51470ace..5a4540a4 100644
--- a/libwvdrmengine/cdm/core/test/test_base.cpp
+++ b/libwvdrmengine/cdm/core/test/test_base.cpp
@@ -227,15 +227,33 @@ std::string WvCdmTestBase::SignHMAC(const std::string& message,
 
 TestCryptoSession::TestCryptoSession(metrics::CryptoMetrics* crypto_metrics)
     : CryptoSession(crypto_metrics) {
-  // The first CryptoSession should have initialized OEMCrypto.  This is right
-  // after that, so we should tell oemcrypto to use a test keybox.
-  if (session_count() == 1) {
+  MaybeInstallTestKeybox();
+}
+
+TestCryptoSession::TestCryptoSession(metrics::CryptoMetrics* crypto_metrics,
+                                     const TestCryptoSessionConfig* config)
+    : CryptoSession(crypto_metrics), config_(config) {
+  MaybeInstallTestKeybox();
+}
+
+void TestCryptoSession::MaybeInstallTestKeybox() {
+  if (IsTestKeyboxNeeded()) {
     CryptoSession::SetAllowTestKeybox(true);
     ReinitializeForTest();
     WvCdmTestBase::InstallTestRootOfTrust();
   }
 }
 
+bool TestCryptoSession::IsTestKeyboxNeeded() {
+  // The first CryptoSession should have initialized OEMCrypto.  This is right
+  // after that.
+  if (session_count() != 1) return false;
+  // If config is not available, assume keybox is required.
+  if (config_ == nullptr) return true;
+  // Unless disabled, test keybox is required.
+  return !config_->disable_test_keybox;
+}
+
 CdmResponseType TestCryptoSession::GenerateNonce(uint32_t* nonce) {
   CdmResponseType status = CryptoSession::GenerateNonce(nonce);
   for (int i = 0; status != NO_ERROR; i++) {
@@ -248,6 +266,7 @@ CdmResponseType TestCryptoSession::GenerateNonce(uint32_t* nonce) {
 }
 
 class TestCryptoSessionFactory : public CryptoSessionFactory {
+ public:
   CryptoSession* MakeCryptoSession(
       metrics::CryptoMetrics* crypto_metrics) override {
     // We need to add extra locking here because we need to make sure that there
@@ -257,9 +276,18 @@ class TestCryptoSessionFactory : public CryptoSessionFactory {
     // InstallTestRootOfTrust is only called in the constructor of the
     // TestCryptoSession, above.
     std::unique_lock<std::mutex> auto_lock(init_lock_);
-    return new TestCryptoSession(crypto_metrics);
+    return new TestCryptoSession(crypto_metrics, &session_config_);
   }
+
+  void SetDisableTestKeybox(bool disable) {
+    std::unique_lock<std::mutex> auto_lock(init_lock_);
+    session_config_.disable_test_keybox = disable;
+  }
+
+ private:
   std::mutex init_lock_;
+  // Shared with all TestCryptoSession instances created by this factory.
+  TestCryptoSessionConfig session_config_;
 };
 
 void WvCdmTestBase::SetUp() {
@@ -279,7 +307,13 @@ void WvCdmTestBase::SetUp() {
       std::string(test_info->test_case_name()) + "." + test_info->name();
   int overwrite = 1;  // Set value even if already set.
   setenv("MODEL_NAME", model_name.c_str(), overwrite);
-  CryptoSession::SetCryptoSessionFactory(new TestCryptoSessionFactory());
+  TestCryptoSessionFactory* factory = new TestCryptoSessionFactory();
+  CryptoSession::SetCryptoSessionFactory(factory);
+  const char* const disable_test_keybox_flag = getenv("DISABLE_TEST_KEYBOX");
+  if (disable_test_keybox_flag != nullptr &&
+      !strcmp(disable_test_keybox_flag, "yes")) {
+    factory->SetDisableTestKeybox(true);
+  }
   // TODO(fredgc): Add a test version of DeviceFiles.
 }
 
diff --git a/libwvdrmengine/cdm/core/test/test_base.h b/libwvdrmengine/cdm/core/test/test_base.h
index 9119363e..8436724f 100644
--- a/libwvdrmengine/cdm/core/test/test_base.h
+++ b/libwvdrmengine/cdm/core/test/test_base.h
@@ -103,12 +103,28 @@ class WvCdmTestBaseWithEngine : public WvCdmTestBase {
   TestCdmEngine cdm_engine_;
 };
 
+struct TestCryptoSessionConfig {
+  // Disables newly created TestCryptoSession instances from installing
+  // a test keybox.
+  bool disable_test_keybox = false;
+};
+
 class TestCryptoSession : public CryptoSession {
  public:
   explicit TestCryptoSession(metrics::CryptoMetrics* crypto_metrics);
+  TestCryptoSession(metrics::CryptoMetrics* crypto_metrics,
+                    const TestCryptoSessionConfig* config);
   // This intercepts nonce flood errors, which is useful for tests that request
   // many nonces and are not time critical.
   CdmResponseType GenerateNonce(uint32_t* nonce) override;
+
+ private:
+  // Called once when TestCryptoSession is constructed.
+  void MaybeInstallTestKeybox();
+  bool IsTestKeyboxNeeded();
+
+  // An un-owned pointer to the config.
+  const TestCryptoSessionConfig* const config_ = nullptr;
 };
 
 // Given a PSSH data structure, this makes a PSSH string for use in