Merge "More Usage Table Tests"

libwvdrmengine/oemcrypto/mock/src/oemcrypto_usage_table_mock.cpp

@@ -427,8 +427,11 @@ OEMCryptoResult UsageTable::LoadUsageEntry(SessionContext* session,
 OEMCryptoResult UsageTable::ShrinkUsageTableHeader(
     uint32_t new_table_size, uint8_t* header_buffer,
     size_t* header_buffer_length) {
-  if (new_table_size > generation_numbers_.size())
+  if (new_table_size > generation_numbers_.size()) {
+    LOGE("OEMCrypto_ShrinkUsageTableHeader: %d > %zd.", new_table_size,
+         generation_numbers_.size());
     return OEMCrypto_ERROR_UNKNOWN_FAILURE;
+  }
   size_t signed_header_size = SignedHeaderSize(new_table_size);
   if (*header_buffer_length < signed_header_size) {
     *header_buffer_length = signed_header_size;
@@ -439,7 +442,7 @@ OEMCryptoResult UsageTable::ShrinkUsageTableHeader(
     LOGE("OEMCrypto_ShrinkUsageTableHeader: buffer null.");
     return OEMCrypto_ERROR_UNKNOWN_FAILURE;
   }
-  for (size_t i = new_table_size; i < sessions_.size(); i++) {
+  for (size_t i = new_table_size; i < sessions_.size(); ++i) {
     if (sessions_[i]) {
       LOGE("ShrinkUsageTableHeader: session open for %d", i);
       return OEMCrypto_ERROR_INVALID_CONTEXT;
@@ -679,6 +682,13 @@ OEMCryptoResult UsageTable::CreateUsageTableHeader(
   }
   *header_buffer_length = signed_header_size;
   if (!LoadGenerationNumber(true)) return OEMCrypto_ERROR_UNKNOWN_FAILURE;
+  // Make sure there are no entries that are currently tied to an open session.
+  for (size_t i=0; i < sessions_.size(); ++i) {
+    if (sessions_[i] != NULL) {
+      LOGE("CreateUsageTableHeader: index %d used by session.", i);
+      return OEMCrypto_ERROR_INVALID_SESSION;
+    }
+  }
   sessions_.clear();
   generation_numbers_.clear();
   header_loaded_ = true;

libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp

@@ -615,14 +615,24 @@ class OEMCryptoSessionTests : public OEMCryptoClientTest {
     }
   }
 
-  void CreateUsageTableHeader() {
+  void CreateUsageTableHeader(bool expect_success = true) {
     size_t header_buffer_length = 0;
-    ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
+    OEMCryptoResult sts =
-              OEMCrypto_CreateUsageTableHeader(NULL, &header_buffer_length));
+        OEMCrypto_CreateUsageTableHeader(NULL, &header_buffer_length);
+    if (expect_success) {
+      ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, sts);
+    } else {
+      ASSERT_NE(OEMCrypto_SUCCESS, sts);
+      if (sts != OEMCrypto_ERROR_SHORT_BUFFER) return;
+    }
     encrypted_usage_header_.resize(header_buffer_length);
-    ASSERT_EQ(OEMCrypto_SUCCESS,
+    sts = OEMCrypto_CreateUsageTableHeader(&encrypted_usage_header_[0],
-              OEMCrypto_CreateUsageTableHeader(&encrypted_usage_header_[0],
+                                           &header_buffer_length);
-                                               &header_buffer_length));
+    if (expect_success) {
+      ASSERT_EQ(OEMCrypto_SUCCESS, sts);
+    } else {
+      ASSERT_NE(OEMCrypto_SUCCESS, sts);
+    }
   }
 
   void EnsureTestKeys() {
@@ -4514,7 +4524,7 @@ TEST_F(UsageTableTest, OnlineMissingEntry) {
   ASSERT_NO_FATAL_FAILURE(s.close());
 }
 
-TEST_F(UsageTableTest, TWoHundredEntries) {
+TEST_F(UsageTableTest, TwoHundredEntries) {
   Session s1;
   ASSERT_NO_FATAL_FAILURE(s1.open());
   ASSERT_NO_FATAL_FAILURE(InstallTestSessionKeys(&s1));
@@ -5065,16 +5075,26 @@ class UsageTableDefragTest : public UsageTableTest {
     ASSERT_NO_FATAL_FAILURE(s->close());
   }
 
-  void ShrinkHeader(uint32_t new_size) {
+  void ShrinkHeader(uint32_t new_size, bool expect_success = true) {
     size_t header_buffer_length = 0;
-    ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
+    OEMCryptoResult sts =
-              OEMCrypto_ShrinkUsageTableHeader(new_size, NULL,
+        OEMCrypto_ShrinkUsageTableHeader(new_size, NULL, &header_buffer_length);
-                                               &header_buffer_length));
+    if (expect_success) {
+      ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, sts);
+    } else {
+      ASSERT_NE(OEMCrypto_SUCCESS, sts);
+      if (sts != OEMCrypto_ERROR_SHORT_BUFFER) return;
+    }
     ASSERT_LT(0u, header_buffer_length);
     encrypted_usage_header_.resize(header_buffer_length);
-    ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_ShrinkUsageTableHeader(
+    sts = OEMCrypto_ShrinkUsageTableHeader(
-                                     new_size, &encrypted_usage_header_[0],
+        new_size, &encrypted_usage_header_[0],
-                                     &header_buffer_length));
+        &header_buffer_length);
+    if (expect_success) {
+      ASSERT_EQ(OEMCrypto_SUCCESS, sts);
+    } else {
+      ASSERT_NE(OEMCrypto_SUCCESS, sts);
+    }
   }
 };
 
@@ -5125,6 +5145,40 @@ TEST_F(UsageTableDefragTest, MoveUsageEntriesToOpenSession) {
                                             OEMCrypto_ERROR_INVALID_SESSION));
 }
 
+TEST_F(UsageTableDefragTest, ShrinkOverOpenSessions) {
+  Session s0;
+  Session s1;
+  LoadFirstLicense(&s0, 0);
+  LoadFirstLicense(&s1, 1);
+  s0.open();
+  ASSERT_NO_FATAL_FAILURE(s0.ReloadUsageEntry());
+  s1.open();
+  ASSERT_NO_FATAL_FAILURE(s1.ReloadUsageEntry());
+  // Since s0 and s1 are open, we can't shrink.
+  ASSERT_NO_FATAL_FAILURE(ShrinkHeader(1, /* expect_success */ false));
+  s1.close();  // Can shrink after closing s1, even if s0 is open.
+  ASSERT_NO_FATAL_FAILURE(ShrinkHeader(1, /* expect_success */ true));
+}
+
+TEST_F(UsageTableDefragTest, EnlargeHeader) {
+  Session s0;
+  Session s1;
+  LoadFirstLicense(&s0, 0);
+  LoadFirstLicense(&s1, 1);
+  // Can only shrink the header -- not make it bigger.
+  ASSERT_NO_FATAL_FAILURE(ShrinkHeader(4, /* expect_success */ false));
+}
+
+TEST_F(UsageTableDefragTest, CreateNewHeaderWhileUsingOldOne) {
+  Session s0;
+  Session s1;
+  LoadFirstLicense(&s0, 0);
+  LoadFirstLicense(&s1, 1);
+  s0.open();
+  ASSERT_NO_FATAL_FAILURE(s0.ReloadUsageEntry());
+  ASSERT_NO_FATAL_FAILURE(CreateUsageTableHeader(/* expect_success */ false));
+}
+
 TEST_F(UsageTableDefragTest, ReloadUsageEntryWrongIndex) {
   Session s0;
   Session s1;