Handle key types in D3D11CdmProxy to handle the keyblobs correctly

- CdmProxy::SetKey takes a key type. This should be used to manage
  possibly different key blobs for different key types.
- Calling CdmProxyContext::GetD3D11DecryptContext requires passing in
  a key type as well, to specify which key type it's requesting for.
- Moved CdmProxyContext from media/base to media/cdm because 1) it
  doesn't need to be in media/base and 2) it depends on cdm_proxy.h
  now, to access CdmProxy::KeyType.

Cq-Include-Trybots: luci.chromium.try:android_optional_gpu_tests_rel;luci.chromium.try:linux_optional_gpu_tests_rel;luci.chromium.try:mac_optional_gpu_tests_rel;luci.chromium.try:win_optional_gpu_tests_rel
Change-Id: I4064e9426bc8efac3fe667806a9f3713881b31ac
Reviewed-on: https://chromium-review.googlesource.com/1234332
Commit-Queue: Rintaro Kuroiwa <rkuroiwa@chromium.org>
Reviewed-by: Xiaohan Wang <xhwang@chromium.org>
Cr-Commit-Position: refs/heads/master@{#593222}

media/base/BUILD.gn

@@ -98,8 +98,6 @@ jumbo_source_set("base") {
     "cdm_promise.h",
     "cdm_promise_adapter.cc",
     "cdm_promise_adapter.h",
-    "cdm_proxy_context.cc",
-    "cdm_proxy_context.h",
     "cdm_session_tracker.cc",
     "cdm_session_tracker.h",
     "channel_mixer.cc",

media/cdm/BUILD.gn

@@ -41,6 +41,8 @@ source_set("cdm") {
     # when EnabledIf attribute is supported in mojom.
     "cdm_proxy.cc",
     "cdm_proxy.h",
+    "cdm_proxy_context.cc",
+    "cdm_proxy_context.h",
     "cenc_utils.cc",
     "cenc_utils.h",
     "default_cdm_factory.cc",

media/cdm/cdm_proxy_context.cc

@@ -2,7 +2,7 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-#include "media/base/cdm_proxy_context.h"
+#include "media/cdm/cdm_proxy_context.h"
 
 #include "build/build_config.h"
 
@@ -13,7 +13,8 @@ CdmProxyContext::~CdmProxyContext() {}
 
 #if defined(OS_WIN)
 base::Optional<CdmProxyContext::D3D11DecryptContext>
-CdmProxyContext::GetD3D11DecryptContext(const std::string& key_id) {
+CdmProxyContext::GetD3D11DecryptContext(CdmProxy::KeyType key_type,
+                                        const std::string& key_id) {
   return base::nullopt;
 }
 #endif

media/cdm/cdm_proxy_context.h

@@ -2,8 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-#ifndef MEDIA_BASE_CDM_PROXY_CONTEXT_H_
-#define MEDIA_BASE_CDM_PROXY_CONTEXT_H_
+#ifndef MEDIA_CDM_CDM_PROXY_CONTEXT_H_
+#define MEDIA_CDM_CDM_PROXY_CONTEXT_H_
 
 #include <stdint.h>
 
@@ -14,6 +14,7 @@
 #include "base/optional.h"
 #include "build/build_config.h"
 #include "media/base/media_export.h"
+#include "media/cdm/cdm_proxy.h"
 
 #if defined(OS_WIN)
 #include <d3d11.h>
@@ -45,8 +46,10 @@ class MEDIA_EXPORT CdmProxyContext {
   // Returns D3D11DecryptContext on success. Returns nullopt otherwise. The
   // D3D11DecryptContext instance is only guaranteed to be valid before the
   // caller returns.
+  // |key_type| is the requesting key type.
   // |key_id| is the key ID of the media to decrypt.
   virtual base::Optional<D3D11DecryptContext> GetD3D11DecryptContext(
+      CdmProxy::KeyType key_type,
       const std::string& key_id) WARN_UNUSED_RESULT;
 #endif  // defined(OS_WIN)
 
@@ -61,4 +64,4 @@ class MEDIA_EXPORT CdmProxyContext {
 
 }  // namespace media
 
-#endif  // MEDIA_BASE_CDM_PROXY_CONTEXT_H_
+#endif  // MEDIA_CDM_CDM_PROXY_CONTEXT_H_

media/gpu/windows/d3d11_cdm_proxy.cc

@@ -15,7 +15,7 @@
 #include "base/win/object_watcher.h"
 #include "media/base/callback_registry.h"
 #include "media/base/cdm_context.h"
-#include "media/base/cdm_proxy_context.h"
+#include "media/cdm/cdm_proxy_context.h"
 #include "media/gpu/windows/d3d11_decryptor.h"
 
 namespace media {
@@ -62,20 +62,21 @@ class D3D11CdmProxyContext : public CdmProxyContext {
   ~D3D11CdmProxyContext() override = default;
 
   // The pointers are owned by the caller.
-  // TODO(rkuroiwa): Handle |key_type|.
   void SetKey(ID3D11CryptoSession* crypto_session,
               const std::vector<uint8_t>& key_id,
-              CdmProxy::KeyType /* key_type */,
+              CdmProxy::KeyType key_type,
               const std::vector<uint8_t>& key_blob) {
     std::string key_id_str(key_id.begin(), key_id.end());
     KeyInfo key_info(crypto_session, key_blob);
     // Note that this would overwrite an entry but it is completely valid, e.g.
     // updating the keyblob due to a configuration change.
-    key_info_map_[key_id_str] = std::move(key_info);
+    key_info_map_[key_id_str][key_type] = std::move(key_info);
   }
 
   void RemoveKey(ID3D11CryptoSession* crypto_session,
                  const std::vector<uint8_t>& key_id) {
+    // There's no need for a keytype for Remove() at the moment, because it's
+    // used for completely removing keys associated to |key_id|.
     std::string key_id_str(key_id.begin(), key_id.end());
     key_info_map_.erase(key_id_str);
   }
@@ -85,12 +86,18 @@ class D3D11CdmProxyContext : public CdmProxyContext {
 
   // CdmProxyContext implementation.
   base::Optional<D3D11DecryptContext> GetD3D11DecryptContext(
+      CdmProxy::KeyType key_type,
       const std::string& key_id) override {
-    auto key_info_it = key_info_map_.find(key_id);
-    if (key_info_it == key_info_map_.end())
+    auto key_id_find_it = key_info_map_.find(key_id);
+    if (key_id_find_it == key_info_map_.end())
       return base::nullopt;
 
-    auto& key_info = key_info_it->second;
+    auto& key_type_to_key_info = key_id_find_it->second;
+    auto key_type_find_it = key_type_to_key_info.find(key_type);
+    if (key_type_find_it == key_type_to_key_info.end())
+      return base::nullopt;
+
+    auto& key_info = key_type_find_it->second;
     D3D11DecryptContext context = {};
     context.crypto_session = key_info.crypto_session;
     context.key_blob = key_info.key_blob.data();
@@ -113,13 +120,13 @@ class D3D11CdmProxyContext : public CdmProxyContext {
     std::vector<uint8_t> key_blob;
   };
 
-  // Maps key ID to KeyInfo.
+  // Maps key ID -> key type -> KeyInfo.
   // The key ID's type is string, which is converted from |key_id| in
   // SetKey(). It's better to use string here rather than convert
   // vector<uint8_t> to string every time in GetD3D11DecryptContext() because
   // in most cases it would be called more often than SetKey() and RemoveKey()
   // combined.
-  std::map<std::string, KeyInfo> key_info_map_;
+  std::map<std::string, std::map<CdmProxy::KeyType, KeyInfo>> key_info_map_;
 
   const GUID key_info_guid_;
 

media/gpu/windows/d3d11_cdm_proxy_unittest.cc

@@ -13,7 +13,7 @@
 #include "base/power_monitor/power_monitor_source.h"
 #include "base/run_loop.h"
 #include "base/test/scoped_task_environment.h"
-#include "media/base/cdm_proxy_context.h"
+#include "media/cdm/cdm_proxy_context.h"
 #include "media/gpu/windows/d3d11_mocks.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
@@ -678,12 +678,39 @@ TEST_F(D3D11CdmProxyTest, GetCdmProxyContext) {
   ASSERT_TRUE(context->GetCdmProxyContext());
 }
 
+// No keys are set.
 TEST_F(D3D11CdmProxyTest, GetD3D11DecryptContextNoKey) {
   base::WeakPtr<CdmContext> context = proxy_->GetCdmContext();
   ASSERT_TRUE(context);
   CdmProxyContext* proxy_context = context->GetCdmProxyContext();
-  // The key ID doesn't matter.
-  auto decrypt_context = proxy_context->GetD3D11DecryptContext("");
+  auto decrypt_context =
+      proxy_context->GetD3D11DecryptContext(kTestKeyType, "");
+  EXPECT_FALSE(decrypt_context);
+}
+
+// A key is set but no keys for the key type requested.
+TEST_F(D3D11CdmProxyTest, GetD3D11DecryptContextNoKeyForKeyType) {
+  uint32_t crypto_session_id_from_initialize = 0;
+  EXPECT_CALL(callback_mock_,
+              InitializeCallback(CdmProxy::Status::kOk, kTestProtocol, _))
+      .WillOnce(SaveArg<2>(&crypto_session_id_from_initialize));
+  ASSERT_NO_FATAL_FAILURE(
+      Initialize(nullptr, base::BindOnce(&CallbackMock::InitializeCallback,
+                                         base::Unretained(&callback_mock_))));
+  ::testing::Mock::VerifyAndClearExpectations(&callback_mock_);
+
+  const std::vector<uint8_t> kAnyBlob = {
+      0x01, 0x4f, 0x83,
+  };
+
+  proxy_->SetKey(crypto_session_id_from_initialize, kAnyBlob,
+                 CdmProxy::KeyType::kDecryptAndDecode, kAnyBlob);
+
+  base::WeakPtr<CdmContext> context = proxy_->GetCdmContext();
+  CdmProxyContext* proxy_context = context->GetCdmProxyContext();
+  auto decrypt_context = proxy_context->GetD3D11DecryptContext(
+      CdmProxy::KeyType::kDecryptOnly,
+      std::string(kAnyBlob.begin(), kAnyBlob.end()));
   EXPECT_FALSE(decrypt_context);
 }
 
@@ -714,7 +741,8 @@ TEST_F(D3D11CdmProxyTest, SetKeyAndGetDecryptContext) {
                  kKeyBlob);
 
   std::string key_id_str(kKeyId.begin(), kKeyId.end());
-  auto decrypt_context = proxy_context->GetD3D11DecryptContext(key_id_str);
+  auto decrypt_context =
+      proxy_context->GetD3D11DecryptContext(kTestKeyType, key_id_str);
   ASSERT_TRUE(decrypt_context);
 
   EXPECT_TRUE(decrypt_context->crypto_session)
@@ -763,7 +791,8 @@ TEST_F(D3D11CdmProxyTest, ClearKeysAfterHardwareContentProtectionTeardown) {
   run_loop.Run();
 
   std::string key_id_str(kKeyId.begin(), kKeyId.end());
-  auto decrypt_context = proxy_context->GetD3D11DecryptContext(key_id_str);
+  auto decrypt_context =
+      proxy_context->GetD3D11DecryptContext(kTestKeyType, key_id_str);
   ASSERT_FALSE(decrypt_context);
 }
 
@@ -795,7 +824,8 @@ TEST_F(D3D11CdmProxyTest, RemoveKey) {
   proxy_->RemoveKey(crypto_session_id_from_initialize, kKeyId);
 
   std::string keyblob_str(kKeyId.begin(), kKeyId.end());
-  auto decrypt_context = proxy_context->GetD3D11DecryptContext(keyblob_str);
+  auto decrypt_context =
+      proxy_context->GetD3D11DecryptContext(kTestKeyType, keyblob_str);
   EXPECT_FALSE(decrypt_context);
 }
 

media/gpu/windows/d3d11_decryptor.cc

@@ -168,8 +168,8 @@ void D3D11Decryptor::Decrypt(StreamType stream_type,
     return;
   }
 
-  auto context =
-      cdm_proxy_context_->GetD3D11DecryptContext(decrypt_config->key_id());
+  auto context = cdm_proxy_context_->GetD3D11DecryptContext(
+      CdmProxy::KeyType::kDecryptOnly, decrypt_config->key_id());
   if (!context) {
     decrypt_cb.Run(kNoKey, nullptr);
     return;

media/gpu/windows/d3d11_decryptor.h

@@ -10,8 +10,8 @@
 #include "base/containers/span.h"
 #include "base/macros.h"
 #include "base/memory/weak_ptr.h"
-#include "media/base/cdm_proxy_context.h"
 #include "media/base/decryptor.h"
+#include "media/cdm/cdm_proxy_context.h"
 #include "media/gpu/media_gpu_export.h"
 #include "media/gpu/windows/d3d11_create_device_cb.h"
 

media/gpu/windows/d3d11_decryptor_unittest.cc

@@ -11,9 +11,9 @@
 #include "base/bind.h"
 #include "base/containers/span.h"
 #include "base/stl_util.h"
-#include "media/base/cdm_proxy_context.h"
 #include "media/base/decoder_buffer.h"
 #include "media/base/subsample_entry.h"
+#include "media/cdm/cdm_proxy_context.h"
 #include "media/gpu/windows/d3d11_mocks.h"
 
 using ::testing::_;
@@ -87,8 +87,9 @@ class CallbackMock {
 
 class CdmProxyContextMock : public CdmProxyContext {
  public:
-  MOCK_METHOD1(GetD3D11DecryptContext,
-               base::Optional<D3D11DecryptContext>(const std::string& key_id));
+  MOCK_METHOD2(GetD3D11DecryptContext,
+               base::Optional<D3D11DecryptContext>(CdmProxy::KeyType key_type,
+                                                   const std::string& key_id));
 };
 
 // Checks that BUFFER_DESC has these fields match.
@@ -181,7 +182,8 @@ class D3D11DecryptorTest : public ::testing::Test {
             DoAll(AddRefAndSetArgPointee<1>(video_context_mock_.Get()),
                   Return(S_OK)));
 
-    EXPECT_CALL(mock_proxy_, GetD3D11DecryptContext(kKeyId))
+    EXPECT_CALL(mock_proxy_,
+                GetD3D11DecryptContext(CdmProxy::KeyType::kDecryptOnly, kKeyId))
         .WillOnce(Return(*decrypt_context));
 
     // These return big enough size.
@@ -469,7 +471,8 @@ TEST_F(D3D11DecryptorTest, NoDecryptContext) {
   scoped_refptr<DecoderBuffer> encrypted_buffer = TestDecoderBuffer(
       kAnyInput, base::size(kAnyInput), kKeyId, kIv, {kAnyInputSubsample});
 
-  EXPECT_CALL(mock_proxy_, GetD3D11DecryptContext(kKeyId))
+  EXPECT_CALL(mock_proxy_,
+              GetD3D11DecryptContext(CdmProxy::KeyType::kDecryptOnly, kKeyId))
       .WillOnce(Return(base::nullopt));
 
   CallbackMock callbacks;
@@ -504,7 +507,8 @@ TEST_F(D3D11DecryptorTest, ReuseBuffers) {
   EXPECT_CALL(*crypto_session_mock.Get(), GetDevice(_))
       .Times(AtLeast(1))
       .WillRepeatedly(AddRefAndSetArgPointee<0>(device_mock_.Get()));
-  EXPECT_CALL(mock_proxy_, GetD3D11DecryptContext(kKeyId))
+  EXPECT_CALL(mock_proxy_,
+              GetD3D11DecryptContext(CdmProxy::KeyType::kDecryptOnly, kKeyId))
       .WillOnce(Return(decrypt_context));
 
   // Buffers should not be (re)initialized on the next call to decrypt because

media/gpu/windows/d3d11_h264_accelerator.cc

@@ -8,8 +8,8 @@
 
 #include "base/memory/ptr_util.h"
 #include "base/trace_event/trace_event.h"
-#include "media/base/cdm_proxy_context.h"
 #include "media/base/media_log.h"
+#include "media/cdm/cdm_proxy_context.h"
 #include "media/gpu/h264_decoder.h"
 #include "media/gpu/h264_dpb.h"
 #include "media/gpu/windows/d3d11_picture_buffer.h"
@@ -105,7 +105,7 @@ Status D3D11H264Accelerator::SubmitFrameMetadata(
   base::Optional<CdmProxyContext::D3D11DecryptContext> decrypt_context;
   if (is_encrypted) {
     decrypt_context = cdm_proxy_context_->GetD3D11DecryptContext(
-        pic->decrypt_config()->key_id());
+        CdmProxy::KeyType::kDecryptAndDecode, pic->decrypt_config()->key_id());
     if (!decrypt_context) {
       RecordFailure("Cannot find decrypt context for the frame.");
       return Status::kTryAgain;

media/mojo/services/mojo_cdm_proxy_unittest.cc

@@ -11,9 +11,9 @@
 #include "base/memory/weak_ptr.h"
 #include "base/run_loop.h"
 #include "base/test/test_message_loop.h"
-#include "media/base/cdm_proxy_context.h"
 #include "media/base/gmock_callback_support.h"
 #include "media/base/mock_filters.h"
+#include "media/cdm/cdm_proxy_context.h"
 #include "media/mojo/interfaces/cdm_proxy.mojom.h"
 #include "media/mojo/services/mojo_cdm_proxy.h"
 #include "media/mojo/services/mojo_cdm_proxy_service.h"