From a66143bec48283744afb993efd8bf1f5885ebd3f Mon Sep 17 00:00:00 2001
From: Ian Benz <ianbenz@google.com>
Date: Thu, 23 Mar 2023 23:01:28 +0000
Subject: [PATCH] Fix OEMCrypto_GenerateRSASignature fuzzer

Ensure OEMCrypto is in the correct state before invoking
OEMCrypto_GenerateRSASignature.

Merged from https://widevine-internal-review.googlesource.com/168857
Merged from https://widevine-internal-review.googlesource.com/172171
Merged from https://widevine-internal-review.googlesource.com/172410
Merged from https://widevine-internal-review.googlesource.com/173130

Change-Id: I774f59af8044939c5a9618de348f2120a834e043
---
 .../test/fuzz_tests/oemcrypto_fuzz_helper.cc  | 18 ++++++++----
 .../test/fuzz_tests/oemcrypto_fuzz_helper.h   | 28 +++++++++++++++----
 .../oemcrypto_generate_rsa_signature_fuzz.cc  | 10 ++++---
 3 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_fuzz_helper.cc b/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_fuzz_helper.cc
index 82a182f7..4761f30b 100644
--- a/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_fuzz_helper.cc
+++ b/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_fuzz_helper.cc
@@ -33,16 +33,24 @@ void InitializeFuzz(SessionUtil& session_util) {
   session_util.EnsureTestROT();
 }
 
-void OEMCryptoLicenseAPIFuzz::Initialize() {
+void SessionFuzz::Initialize() {
   InitializeFuzz(session_util_);
   session_.open();
-  session_util_.InstallTestDrmKey(&session_);
-  session_.GenerateNonce();
+}
+
+void SessionFuzz::Terminate() {
+  session_.close();
+  OEMCrypto_Terminate();
+}
+
+void OEMCryptoLicenseAPIFuzz::Initialize() {
+  session_fuzz_.Initialize();
+  session_fuzz_.InstallTestDrmKey();
+  session_fuzz_.session().GenerateNonce();
 }
 
 void OEMCryptoLicenseAPIFuzz::Terminate() {
-  session_.close();
-  OEMCrypto_Terminate();
+  session_fuzz_.Terminate();
 }
 
 void OEMCryptoLicenseAPIFuzz::LoadLicense(bool generic_crypto_keys) {
diff --git a/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_fuzz_helper.h b/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_fuzz_helper.h
index 036bfa26..c92297b1 100644
--- a/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_fuzz_helper.h
+++ b/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_fuzz_helper.h
@@ -38,9 +38,28 @@ struct FuzzedData {
 // OEMCrypto APIs.
 void InitializeFuzz(SessionUtil& session_util);
 
+class SessionFuzz {
+ public:
+  void Initialize();
+
+  void Terminate();
+
+  void InstallTestDrmKey() {
+    session_util_.InstallTestDrmKey(&session_);
+  }
+
+  Session& session() { return session_; }
+
+  const Session& session() const { return session_; }
+
+ private:
+  SessionUtil session_util_;
+  Session session_;
+};
+
 class OEMCryptoLicenseAPIFuzz {
  public:
-  OEMCryptoLicenseAPIFuzz() : license_messages_(&session_) {}
+  OEMCryptoLicenseAPIFuzz() : license_messages_(&session_fuzz_.session()) {}
 
   void Initialize();
 
@@ -54,15 +73,14 @@ class OEMCryptoLicenseAPIFuzz {
 
   const LicenseRoundTrip& license_messages() const { return license_messages_; }
 
-  Session& session() { return session_; }
+  Session& session() { return session_fuzz_.session(); }
 
-  const Session& session() const { return session_; }
+  const Session& session() const { return session_fuzz_.session(); }
 
  private:
   void LoadLicense(bool generic_crypto_keys);
 
-  SessionUtil session_util_;
-  Session session_;
+  SessionFuzz session_fuzz_;
   LicenseRoundTrip license_messages_;
 };
 
diff --git a/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_generate_rsa_signature_fuzz.cc b/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_generate_rsa_signature_fuzz.cc
index f0e003e9..f30eed87 100644
--- a/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_generate_rsa_signature_fuzz.cc
+++ b/libwvdrmengine/oemcrypto/test/fuzz_tests/oemcrypto_generate_rsa_signature_fuzz.cc
@@ -23,12 +23,14 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
       fuzzed_data.ConsumeRemainingBytes<uint8_t>();
   std::vector<uint8_t> signature(fuzzed_structure.signature_length);
 
-  wvoec::OEMCryptoLicenseAPIFuzz license_api_fuzz;
-  license_api_fuzz.Initialize();
+  wvoec::SessionFuzz session_fuzz;
+  session_fuzz.Initialize();
+  session_fuzz.InstallTestDrmKey();
   OEMCrypto_GenerateRSASignature(
-      license_api_fuzz.session().session_id(), message.data(), message.size(),
+      session_fuzz.session().session_id(), message.data(), message.size(),
       signature.data(), &fuzzed_structure.signature_length,
       fuzzed_structure.padding_scheme);
-  license_api_fuzz.Terminate();
+  session_fuzz.Terminate();
+
   return 0;
 }