From c36826607e2841c9e2236b70b128394a5f2513fb Mon Sep 17 00:00:00 2001
From: Matt Feddersen <mattfedd@google.com>
Date: Tue, 31 Oct 2023 18:19:19 -0700
Subject: [PATCH] Add CDM integration test for MoveUsageEntry bug

Test: GTEST_FILTER="*PIG*Defrag*" ./jenkins/opk_ta
Bug: 307750348

Change-Id: I5cc1ffb32311d6e46a956e11fc400b6f9510e6d6
---
 .../cdm/core/test/license_holder.cpp          |  7 +++
 libwvdrmengine/cdm/core/test/license_holder.h |  2 +
 .../cdm/core/test/policy_integration_test.cpp | 55 +++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/libwvdrmengine/cdm/core/test/license_holder.cpp b/libwvdrmengine/cdm/core/test/license_holder.cpp
index c2532e2f..30300593 100644
--- a/libwvdrmengine/cdm/core/test/license_holder.cpp
+++ b/libwvdrmengine/cdm/core/test/license_holder.cpp
@@ -75,6 +75,13 @@ void LicenseHolder::ReloadLicense() {
       << "Failed to reload license for " << content_id();
 }
 
+void LicenseHolder::FailReloadLicense() {
+  const CdmResponseType status =
+      cdm_engine_->RestoreKey(session_id_, key_set_id_);
+  ASSERT_NE(KEY_ADDED, status)
+      << "Unexpected success loading license for " << content_id();
+}
+
 void LicenseHolder::GenerateAndPostRenewalRequest(
     const std::string& policy_id) {
   event_listener_.set_renewal_needed(false);
diff --git a/libwvdrmengine/cdm/core/test/license_holder.h b/libwvdrmengine/cdm/core/test/license_holder.h
index be527004..41578567 100644
--- a/libwvdrmengine/cdm/core/test/license_holder.h
+++ b/libwvdrmengine/cdm/core/test/license_holder.h
@@ -66,6 +66,8 @@ class LicenseHolder {
   // ReloadLicense(). Also, the key_set_id must have been set previously. The
   // key_set_id is set by calling LoadLicense(), or by calling set_key_set_id().
   void ReloadLicense();
+  // Attempt to reload a license, but expect a failure.
+  void FailReloadLicense();
   // Generate the renewal request, and send it to the server.
   void GenerateAndPostRenewalRequest(const std::string& policy_id);
   // Fetch the renewal response. This can add a few seconds of latency.
diff --git a/libwvdrmengine/cdm/core/test/policy_integration_test.cpp b/libwvdrmengine/cdm/core/test/policy_integration_test.cpp
index 6b439f6d..7313f4ee 100644
--- a/libwvdrmengine/cdm/core/test/policy_integration_test.cpp
+++ b/libwvdrmengine/cdm/core/test/policy_integration_test.cpp
@@ -73,6 +73,61 @@ TEST_F(CorePIGTest, OfflineWithPST) {
   ASSERT_NO_FATAL_FAILURE(holder.CloseSession());
 }
 
+TEST_F(CorePIGTest, OfflineMultipleLicensesWithDefrag) {
+  const KeyId key_id = "0000000000000000";
+
+  // 1. Open a session, load license, close session
+  LicenseHolder holder1("CDM_OfflineWithPST", &cdm_engine_, config_);
+  holder1.set_can_persist(true);
+  ASSERT_NO_FATAL_FAILURE(holder1.OpenSession());
+  ASSERT_NO_FATAL_FAILURE(holder1.FetchLicense());
+  ASSERT_NO_FATAL_FAILURE(holder1.LoadLicense());
+  ASSERT_NO_FATAL_FAILURE(holder1.CloseSession());
+
+  // 2. Open a session, load license, keep session open
+  LicenseHolder holder2("CDM_OfflineWithPST", &cdm_engine_, config_);
+  holder2.set_can_persist(true);
+  ASSERT_NO_FATAL_FAILURE(holder2.OpenSession());
+  ASSERT_NO_FATAL_FAILURE(holder2.FetchLicense());
+  ASSERT_NO_FATAL_FAILURE(holder2.LoadLicense());
+
+  // 3. Remove first license
+  ASSERT_NO_FATAL_FAILURE(holder1.OpenSession());
+  ASSERT_NO_FATAL_FAILURE(holder1.ReloadLicense());
+  ASSERT_NO_FATAL_FAILURE(holder1.RemoveLicense());
+  ASSERT_NO_FATAL_FAILURE(holder1.CloseSession());
+
+  // 4. Open a session, load license
+  LicenseHolder holder3("CDM_OfflineWithPST", &cdm_engine_, config_);
+  holder3.set_can_persist(true);
+  ASSERT_NO_FATAL_FAILURE(holder3.OpenSession());
+  ASSERT_NO_FATAL_FAILURE(holder3.FetchLicense());
+  ASSERT_NO_FATAL_FAILURE(holder3.LoadLicense());
+  EXPECT_EQ(NO_ERROR, holder3.Decrypt(key_id));
+  ASSERT_NO_FATAL_FAILURE(holder3.CloseSession());
+
+  ASSERT_NO_FATAL_FAILURE(holder2.CloseSession());
+
+  // Ensure first offline license can no longer be used
+  ASSERT_NO_FATAL_FAILURE(holder1.OpenSession());
+  ASSERT_NO_FATAL_FAILURE(holder1.FailReloadLicense());
+  EXPECT_NE(NO_ERROR, holder1.Decrypt(key_id));
+  ASSERT_NO_FATAL_FAILURE(holder1.CloseSession());
+
+  // Ensure second and third offline licenses can be used
+  ASSERT_NO_FATAL_FAILURE(holder2.OpenSession());
+  ASSERT_NO_FATAL_FAILURE(holder2.ReloadLicense());
+  EXPECT_EQ(NO_ERROR, holder2.Decrypt(key_id));
+  ASSERT_NO_FATAL_FAILURE(holder2.RemoveLicense());
+  ASSERT_NO_FATAL_FAILURE(holder2.CloseSession());
+
+  ASSERT_NO_FATAL_FAILURE(holder3.OpenSession());
+  ASSERT_NO_FATAL_FAILURE(holder3.ReloadLicense());
+  EXPECT_EQ(NO_ERROR, holder3.Decrypt(key_id));
+  ASSERT_NO_FATAL_FAILURE(holder3.RemoveLicense());
+  ASSERT_NO_FATAL_FAILURE(holder3.CloseSession());
+}
+
 // This test verifies that the system can download and install license with a
 // key that requires secure buffers. It also verifies that we cannot decrypt to
 // a non-secure buffer using this key, but that we can decrypt to a secure