media: Implement RefreshTrustedInput() and SetLastKeyId()

These two functions are needed to notify the CDM about the last used
key IDs during suspend and resume.

Bug: 999747
Change-Id: I4f988784cc67d6e0699ec31ba2236de38c64615e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2331092
Commit-Queue: Xiaohan Wang <xhwang@chromium.org>
Reviewed-by: John Rummell <jrummell@chromium.org>
Cr-Commit-Position: refs/heads/master@{#796176}

media/base/win/mf_cdm_proxy.h

@@ -38,12 +38,12 @@ IMFCdmProxy : public IUnknown {
       _COM_Outptr_ IUnknown** object_out) = 0;
 
   // When the media Renderer is suspended, `MediaFoundationSourceWrapper`
-  // provides its last set of key IDs using `SetLastKeyIds()` when it is
+  // provides its last set of key IDs using `SetLastKeyId()` when it is
   // destructed. Then during resume, the new `MediaFoundationSourceWrapper`
-  // calls `RefreshTrustedInput()` to let the CDM use the key ID information to
+  // calls `RefreshTrustedInput()` to let the CDM use the key IDs information to
   // perform some optimization.
-  virtual HRESULT STDMETHODCALLTYPE SetLastKeyIds(GUID * key_ids,
-                                                  uint32_t key_ids_count) = 0;
+  virtual HRESULT STDMETHODCALLTYPE SetLastKeyId(_In_ uint32_t stream_id,
+                                                 _In_ REFGUID key_id) = 0;
   virtual HRESULT STDMETHODCALLTYPE RefreshTrustedInput() = 0;
 
   // Used by MediaFoundationProtectionManager to implement

media/cdm/win/media_foundation_cdm.cc

@@ -4,8 +4,14 @@
 
 #include "media/cdm/win/media_foundation_cdm.h"
 
+#include <stdlib.h>
+#include <vector>
+
 #include "base/bind.h"
 #include "base/logging.h"
+#include "base/win/scoped_co_mem.h"
+#include "base/win/scoped_propvariant.h"
+#include "base/win/win_util.h"
 #include "media/base/bind_to_current_loop.h"
 #include "media/base/cdm_promise.h"
 #include "media/base/win/mf_cdm_proxy.h"
@@ -18,25 +24,88 @@ namespace {
 
 using Microsoft::WRL::ClassicCom;
 using Microsoft::WRL::ComPtr;
-using Microsoft::WRL::MakeAndInitialize;
+using Microsoft::WRL::Make;
 using Microsoft::WRL::RuntimeClass;
 using Microsoft::WRL::RuntimeClassFlags;
 using Exception = CdmPromise::Exception;
 
+// GUID is little endian. The byte array in network order is big endian.
+std::vector<uint8_t> ByteArrayFromGUID(REFGUID guid) {
+  std::vector<uint8_t> byte_array(sizeof(GUID));
+  GUID* reversed_guid = reinterpret_cast<GUID*>(byte_array.data());
+  *reversed_guid = guid;
+  reversed_guid->Data1 = _byteswap_ulong(guid.Data1);
+  reversed_guid->Data2 = _byteswap_ushort(guid.Data2);
+  reversed_guid->Data3 = _byteswap_ushort(guid.Data3);
+  // Data4 is already a byte array so no need to byte swap.
+  return byte_array;
+}
+
+HRESULT CreatePolicySetEvent(ComPtr<IMFMediaEvent>& policy_set_event) {
+  base::win::ScopedPropVariant policy_set_prop;
+  PROPVARIANT* var_to_set = policy_set_prop.Receive();
+  var_to_set->vt = VT_UI4;
+  var_to_set->ulVal = 0;
+  RETURN_IF_FAILED(MFCreateMediaEvent(
+      MEPolicySet, GUID_NULL, S_OK, policy_set_prop.ptr(), &policy_set_event));
+  return S_OK;
+}
+
+// Notifies the Decryptor about the last key ID so the decryptor can prefetch
+// the corresponding key to reduce start-to-play time when resuming playback.
+// This is done by sending a MEContentProtectionMetadata event.
+HRESULT RefreshDecryptor(IMFTransform* decryptor,
+                         const GUID& protection_system_id,
+                         const GUID& last_key_id) {
+  // The MFT_MESSAGE_NOTIFY_START_OF_STREAM message is usually sent by the MF
+  // pipeline when starting playback. Here we send it out-of-band as it is a
+  // pre-requisite for getting the decryptor to process the
+  // MEContentProtectionMetadata event.
+  RETURN_IF_FAILED(
+      decryptor->ProcessMessage(MFT_MESSAGE_NOTIFY_START_OF_STREAM, 0));
+
+  // After receiving a MEContentProtectionMetadata event, the Decryptor
+  // requires that it is notified of a MEPolicySet event to continue decryption.
+  ComPtr<IMFMediaEvent> policy_set_event;
+  RETURN_IF_FAILED(CreatePolicySetEvent(policy_set_event));
+  RETURN_IF_FAILED(decryptor->ProcessMessage(
+      MFT_MESSAGE_NOTIFY_EVENT,
+      reinterpret_cast<ULONG_PTR>(policy_set_event.Get())));
+
+  // Prepare the MEContentProtectionMetadata event.
+  ComPtr<IMFMediaEvent> key_rotation_event;
+  RETURN_IF_FAILED(MFCreateMediaEvent(MEContentProtectionMetadata, GUID_NULL,
+                                      S_OK, nullptr, &key_rotation_event));
+
+  // MF_EVENT_STREAM_METADATA_SYSTEMID expects the system ID (GUID) to be in
+  // little endian order. So no need to call `ByteArrayFromGUID()`.
+  RETURN_IF_FAILED(key_rotation_event->SetBlob(
+      MF_EVENT_STREAM_METADATA_SYSTEMID,
+      reinterpret_cast<const uint8_t*>(&protection_system_id), sizeof(GUID)));
+
+  std::vector<uint8_t> last_key_id_byte_array = ByteArrayFromGUID(last_key_id);
+  RETURN_IF_FAILED(key_rotation_event->SetBlob(
+      MF_EVENT_STREAM_METADATA_CONTENT_KEYIDS, last_key_id_byte_array.data(),
+      last_key_id_byte_array.size()));
+
+  // The `dwInputStreamID` refers to a local stream ID of the Decryptor. Since
+  // Decryptors typically only support a single stream, always pass 0 here.
+  RETURN_IF_FAILED(
+      decryptor->ProcessEvent(/*dwInputStreamID=*/0, key_rotation_event.Get()));
+  return S_OK;
+}
+
 class CdmProxyImpl
     : public RuntimeClass<RuntimeClassFlags<ClassicCom>, IMFCdmProxy> {
  public:
-  CdmProxyImpl() = default;
+  explicit CdmProxyImpl(ComPtr<IMFContentDecryptionModule> mf_cdm)
+      : mf_cdm_(mf_cdm) {}
   ~CdmProxyImpl() override = default;
 
-  HRESULT RuntimeClassInitialize(ComPtr<IMFContentDecryptionModule> mf_cdm) {
-    mf_cdm_ = mf_cdm;
-    return S_OK;
-  }
-
   // IMFCdmProxy implementation
 
   STDMETHODIMP GetPMPServer(REFIID riid, LPVOID* object_result) override {
+    DVLOG_FUNC(1);
     ComPtr<IMFGetService> cdm_services;
     RETURN_IF_FAILED(mf_cdm_.As(&cdm_services));
     RETURN_IF_FAILED(cdm_services->GetService(
@@ -57,15 +126,16 @@ class CdmProxyImpl
       return S_OK;
     }
 
-    ComPtr<IMFTrustedInput> trusted_input;
-    RETURN_IF_FAILED(mf_cdm_->CreateTrustedInput(
-        content_init_data, content_init_data_size, &trusted_input));
+    if (!trusted_input_) {
+      RETURN_IF_FAILED(mf_cdm_->CreateTrustedInput(
+          content_init_data, content_init_data_size, &trusted_input_));
+    }
 
     // GetInputTrustAuthority takes IUnknown* as the output. Using other COM
     // interface will have a v-table mismatch issue.
     ComPtr<IUnknown> unknown;
     RETURN_IF_FAILED(
-        trusted_input->GetInputTrustAuthority(stream_id, riid, &unknown));
+        trusted_input_->GetInputTrustAuthority(stream_id, riid, &unknown));
 
     ComPtr<IMFInputTrustAuthority> input_trust_authority;
     RETURN_IF_FAILED(unknown.As(&input_trust_authority));
@@ -77,33 +147,77 @@ class CdmProxyImpl
     return S_OK;
   }
 
-  // TODO(xhwang): Implement this.
-  STDMETHODIMP SetLastKeyIds(GUID* key_ids, uint32_t key_ids_count) override {
-    NOTIMPLEMENTED();
+  STDMETHODIMP SetLastKeyId(uint32_t stream_id, REFGUID key_id) override {
+    DVLOG_FUNC(1);
+    last_key_ids_[stream_id] = key_id;
     return S_OK;
   }
 
-  // TODO(xhwang): Implement this.
   STDMETHODIMP RefreshTrustedInput() override {
-    NOTIMPLEMENTED();
+    DVLOG_FUNC(1);
+
+    // Refresh all decryptors of the last key IDs.
+    for (const auto& entry : input_trust_authorities_) {
+      const auto& stream_id = entry.first;
+      const auto& input_trust_authority = entry.second;
+
+      const auto& last_key_id = last_key_ids_[stream_id];
+      if (last_key_id == GUID_NULL)
+        continue;
+
+      ComPtr<IMFTransform> decryptor;
+      RETURN_IF_FAILED(
+          input_trust_authority->GetDecrypter(IID_PPV_ARGS(&decryptor)));
+      GUID protection_system_id;
+      RETURN_IF_FAILED(GetProtectionSystemId(&protection_system_id));
+      RETURN_IF_FAILED(
+          RefreshDecryptor(decryptor.Get(), protection_system_id, last_key_id));
+    }
+
+    input_trust_authorities_.clear();
+    last_key_ids_.clear();
     return S_OK;
   }
 
   STDMETHODIMP
   ProcessContentEnabler(IUnknown* request, IMFAsyncResult* result) override {
+    DVLOG_FUNC(1);
     ComPtr<IMFContentEnabler> content_enabler;
     RETURN_IF_FAILED(request->QueryInterface(IID_PPV_ARGS(&content_enabler)));
     return mf_cdm_->SetContentEnabler(content_enabler.Get(), result);
   }
 
  private:
+  HRESULT GetProtectionSystemId(GUID* protection_system_id) {
+    // Typically the CDM should only return one protection system ID. So just
+    // use the first one if available.
+    base::win::ScopedCoMem<GUID> protection_system_ids;
+    DWORD count = 0;
+    RETURN_IF_FAILED(
+        mf_cdm_->GetProtectionSystemIds(&protection_system_ids, &count));
+    if (count == 0)
+      return E_FAIL;
+
+    *protection_system_id = *protection_system_ids;
+    DVLOG(2) << __func__ << " protection_system_id="
+             << base::win::WStringFromGUID(*protection_system_id);
+
+    return S_OK;
+  }
+
   ComPtr<IMFContentDecryptionModule> mf_cdm_;
 
+  // Store IMFTrustedInput to avoid potential performance cost.
+  ComPtr<IMFTrustedInput> trusted_input_;
+
   // |stream_id| to IMFInputTrustAuthority (ITA) mapping. Serves two purposes:
   // 1. The same ITA should always be returned in GetInputTrustAuthority() for
   // the same |stream_id|.
   // 2. The ITA must keep alive for decryptors to work.
   std::map<uint32_t, ComPtr<IMFInputTrustAuthority>> input_trust_authorities_;
+
+  // |stream_id| to last used key ID mapping.
+  std::map<uint32_t, GUID> last_key_ids_;
 };
 
 }  // namespace
@@ -279,7 +393,7 @@ bool MediaFoundationCdm::GetMediaFoundationCdmProxy(
   DVLOG_FUNC(1);
 
   if (!cdm_proxy_)
-    MakeAndInitialize<CdmProxyImpl>(&cdm_proxy_, mf_cdm_);
+    cdm_proxy_ = Make<CdmProxyImpl>(mf_cdm_);
 
   BindToCurrentLoop(std::move(get_mf_cdm_proxy_cb)).Run(cdm_proxy_);
   return true;

media/renderers/win/media_foundation_renderer_unittest.cc

@@ -47,8 +47,8 @@ class MockMFCdmProxy
                                uint32_t content_init_data_size,
                                REFIID riid,
                                IUnknown** object_result));
-  MOCK_STDCALL_METHOD2(SetLastKeyIds,
-                       HRESULT(GUID* key_ids, uint32_t key_ids_count));
+  MOCK_STDCALL_METHOD2(SetLastKeyId,
+                       HRESULT(uint32_t stream_id, REFGUID key_id));
   MOCK_STDCALL_METHOD0(RefreshTrustedInput, HRESULT());
   MOCK_STDCALL_METHOD2(ProcessContentEnabler,
                        HRESULT(IUnknown* request, IMFAsyncResult* result));

media/renderers/win/media_foundation_source_wrapper.cc

@@ -22,14 +22,12 @@ MediaFoundationSourceWrapper::~MediaFoundationSourceWrapper() {
     return;
 
   // Notify |cdm_proxy_| of last Key IDs.
-  std::vector<GUID> key_ids(StreamCount());
   for (uint32_t stream_id = 0; stream_id < StreamCount(); stream_id++) {
-    key_ids[stream_id] = media_streams_[stream_id]->GetLastKeyId();
+    HRESULT hr = cdm_proxy_->SetLastKeyId(
+        stream_id, media_streams_[stream_id]->GetLastKeyId());
+    DLOG_IF(ERROR, FAILED(hr))
+        << "Failed to notify CDM proxy of last Key IDs: " << PrintHr(hr);
   }
-
-  HRESULT hr = cdm_proxy_->SetLastKeyIds(key_ids.data(), key_ids.size());
-  DLOG_IF(ERROR, FAILED(hr))
-      << "Failed to notify CDM proxy of last Key IDs: " << PrintHr(hr);
 }
 
 HRESULT MediaFoundationSourceWrapper::RuntimeClassInitialize(