Refactor EncryptorSource to prepare for live support

Remove Initialize() and Add AddKey().

Also remove fixed_encryptor_source.* and move the functionality to
encryptor_source.*.

Change-Id: I4fb61013177874a8b81854f10b2deda83accc683

media/base/encryptor_source.cc

@@ -6,12 +6,11 @@
 
 #include "media/base/encryptor_source.h"
 
+#include "base/strings/string_number_conversions.h"
 #include "media/base/aes_encryptor.h"
+#include "media/base/buffer_writer.h"
 
 namespace {
-// Generate 64bit IV by default.
-const size_t kDefaultIvSize = 8u;
-
 const uint8 kWidevineSystemId[] = {0xed, 0xef, 0x8b, 0xa9, 0x79, 0xd6,
                                    0x4a, 0xce, 0xa3, 0xc8, 0x27, 0xdc,
                                    0xd5, 0x1d, 0x21, 0xed};
@@ -19,23 +18,103 @@ const uint8 kWidevineSystemId[] = {0xed, 0xef, 0x8b, 0xa9, 0x79, 0xd6,
 
 namespace media {
 
-EncryptorSource::EncryptorSource()
+EncryptionKey::EncryptionKey() {}
-    : iv_size_(kDefaultIvSize),
+EncryptionKey::~EncryptionKey() {}
-      key_system_id_(kWidevineSystemId,
-                     kWidevineSystemId + arraysize(kWidevineSystemId)) {}
 
 EncryptorSource::~EncryptorSource() {}
 
-scoped_ptr<AesCtrEncryptor> EncryptorSource::CreateEncryptor() {
+Status EncryptorSource::GetKey(TrackType track_type, EncryptionKey* key) {
-  scoped_ptr<AesCtrEncryptor> encryptor(new AesCtrEncryptor());
+  DCHECK(key);
-  const bool initialized =
+  DCHECK(encryption_key_);
-      iv_.empty() ? encryptor->InitializeWithRandomIv(key_, iv_size_)
+  *key = *encryption_key_;
-                  : encryptor->InitializeWithIv(key_, iv_);
+  return Status::OK;
-  if (!initialized) {
-    LOG(ERROR) << "Failed to the initialize encryptor.";
-    return scoped_ptr<AesCtrEncryptor>();
 }
-  return encryptor.Pass();
+
+scoped_ptr<EncryptorSource> EncryptorSource::CreateFromHexStrings(
+    const std::string& key_id_hex,
+    const std::string& key_hex,
+    const std::string& pssh_data_hex,
+    const std::string& iv_hex) {
+  scoped_ptr<EncryptionKey> encryption_key(new EncryptionKey());
+
+  if (!base::HexStringToBytes(key_id_hex, &encryption_key->key_id)) {
+    LOG(ERROR) << "Cannot parse key_id_hex " << key_id_hex;
+    return scoped_ptr<EncryptorSource>();
+  }
+
+  if (!base::HexStringToBytes(key_hex, &encryption_key->key)) {
+    LOG(ERROR) << "Cannot parse key_hex " << key_hex;
+    return scoped_ptr<EncryptorSource>();
+  }
+
+  std::vector<uint8> pssh_data;
+  if (!base::HexStringToBytes(pssh_data_hex, &pssh_data)) {
+    LOG(ERROR) << "Cannot parse pssh_hex " << pssh_data_hex;
+    return scoped_ptr<EncryptorSource>();
+  }
+
+  if (!iv_hex.empty()) {
+    if (!base::HexStringToBytes(iv_hex, &encryption_key->iv)) {
+      LOG(ERROR) << "Cannot parse iv_hex " << iv_hex;
+      return scoped_ptr<EncryptorSource>();
+    }
+  }
+
+  encryption_key->pssh = PsshBoxFromPsshData(pssh_data);
+  return scoped_ptr<EncryptorSource>(
+      new EncryptorSource(encryption_key.Pass()));
+}
+
+EncryptorSource::TrackType EncryptorSource::GetTrackTypeFromString(
+    const std::string& track_type_string) {
+  if (track_type_string == "SD")
+    return TRACK_TYPE_SD;
+  if (track_type_string == "HD")
+    return TRACK_TYPE_HD;
+  if (track_type_string == "AUDIO")
+    return TRACK_TYPE_AUDIO;
+  LOG(WARNING) << "Unexpected track type: " << track_type_string;
+  return TRACK_TYPE_UNKNOWN;
+}
+
+std::string EncryptorSource::TrackTypeToString(TrackType track_type) {
+  switch (track_type) {
+    case TRACK_TYPE_SD:
+      return "SD";
+    case TRACK_TYPE_HD:
+      return "HD";
+    case TRACK_TYPE_AUDIO:
+      return "AUDIO";
+    default:
+      NOTIMPLEMENTED() << "Unknown track type: " << track_type;
+      return "UNKNOWN";
+  }
+}
+
+std::vector<uint8> EncryptorSource::PsshBoxFromPsshData(
+    const std::vector<uint8>& pssh_data) {
+  const uint8 kPsshFourCC[] = {'p', 's', 's', 'h'};
+  const uint32 kVersionAndFlags = 0;
+
+  const uint32 pssh_data_size = pssh_data.size();
+  const uint32 total_size =
+      sizeof(total_size) + sizeof(kPsshFourCC) + sizeof(kVersionAndFlags) +
+      sizeof(kWidevineSystemId) + sizeof(pssh_data_size) + pssh_data_size;
+
+  BufferWriter writer;
+  writer.AppendInt(total_size);
+  writer.AppendArray(kPsshFourCC, sizeof(kPsshFourCC));
+  writer.AppendInt(kVersionAndFlags);
+  writer.AppendArray(kWidevineSystemId, sizeof(kWidevineSystemId));
+  writer.AppendInt(pssh_data_size);
+  writer.AppendVector(pssh_data);
+  return std::vector<uint8>(writer.Buffer(), writer.Buffer() + writer.Size());
+}
+
+EncryptorSource::EncryptorSource() {}
+EncryptorSource::EncryptorSource(scoped_ptr<EncryptionKey> encryption_key)
+    : encryption_key_(encryption_key.Pass()) {
+  DCHECK(encryption_key_);
 }
 
 }  // namespace media

media/base/encryptor_source.h

@@ -10,56 +10,73 @@
 #include <vector>
 
 #include "base/memory/scoped_ptr.h"
-#include "media/base/aes_encryptor.h"
 #include "media/base/status.h"
 
 namespace media {
 
-class AesCtrEncryptor;
+struct EncryptionKey {
+  EncryptionKey();
+  ~EncryptionKey();
+
+  std::vector<uint8> key_id;
+  std::vector<uint8> key;
+  std::vector<uint8> pssh;
+  std::vector<uint8> iv;
+};
 
 /// EncryptorSource is responsible for encryption key acquisition.
 class EncryptorSource {
  public:
-  EncryptorSource();
+  enum TrackType {
+    TRACK_TYPE_UNKNOWN = 0,
+    TRACK_TYPE_SD = 1,
+    TRACK_TYPE_HD = 2,
+    TRACK_TYPE_AUDIO = 3,
+    NUM_VALID_TRACK_TYPES = 3
+  };
+
   virtual ~EncryptorSource();
 
-  /// Initialize the encryptor source. Calling other public methods of this
+  /// Get encryption key of the specified track type.
-  /// class without this method returning OK results in an undefined behavior.
+  /// @return OK on success, an error status otherwise.
-  virtual Status Initialize() = 0;
+  virtual Status GetKey(TrackType track_type, EncryptionKey* key);
 
-  /// Create an encryptor from this encryptor source. The encryptor will be
+  /// Create EncryptorSource object from hex strings.
-  /// initialized with a random IV of the default size by default. The behavior
+  /// @param key_id_hex is the key id in hex string.
-  /// can be adjusted using set_iv_size or set_iv (exclusive).
+  /// @param key_hex is the key in hex string.
-  scoped_ptr<AesCtrEncryptor> CreateEncryptor();
+  /// @param pssh_data_hex is the pssh_data in hex string.
+  /// @param iv_hex is the IV in hex string. If not specified, a randomly
+  ///        generated IV with the default length will be used.
+  /// Note: GetKey on the created key source will always return the same key
+  ///       for all track types.
+  static scoped_ptr<EncryptorSource> CreateFromHexStrings(
+      const std::string& key_id_hex,
+      const std::string& key_hex,
+      const std::string& pssh_data_hex,
+      const std::string& iv_hex);
 
-  const std::vector<uint8>& key_id() const { return key_id_; }
+  /// Convert string representation of track type to enum representation.
-  const std::vector<uint8>& key() const { return key_; }
+  static TrackType GetTrackTypeFromString(const std::string& track_type_string);
-  const std::vector<uint8>& pssh() const { return pssh_; }
-  const std::vector<uint8>& key_system_id() const { return key_system_id_; }
-  size_t iv_size() const { return iv_.empty() ? iv_size_ : iv_.size(); }
 
-  /// Set IV size. The encryptor will be initialized with a random IV of the
+  /// Convert TrackType to string.
-  /// specified size. Mutually exclusive with set_iv.
+  static std::string TrackTypeToString(TrackType track_type);
-  void set_iv_size(size_t iv_size) { iv_size_ = iv_size; }
-  /// Set IV. The encryptor will be initialized with the specified IV.
-  /// Mutually exclusive with set_iv_size.
-  void set_iv(std::vector<uint8>& iv) { iv_ = iv; }
 
  protected:
-  void set_key_id(const std::vector<uint8>& key_id) { key_id_ = key_id; }
+  EncryptorSource();
-  void set_key(const std::vector<uint8>& key) { key_ = key; }
+
-  void set_pssh(const std::vector<uint8>& pssh) { pssh_ = pssh; }
+  /// @return the raw bytes of the pssh box with system ID and box header
+  ///         included.
+  static std::vector<uint8> PsshBoxFromPsshData(
+      const std::vector<uint8>& pssh_data);
 
  private:
-  std::vector<uint8> key_id_;
+  explicit EncryptorSource(scoped_ptr<EncryptionKey> encryption_key);
-  std::vector<uint8> key_;
+
-  std::vector<uint8> pssh_;
+  scoped_ptr<EncryptionKey> encryption_key_;
-  size_t iv_size_;
-  std::vector<uint8> iv_;
-  const std::vector<uint8> key_system_id_;
 
   DISALLOW_COPY_AND_ASSIGN(EncryptorSource);
 };
-}
+
+}  // namespace media
 
 #endif  // MEDIA_BASE_ENCRYPTOR_SOURCE_H_

media/base/media_base.gyp

@@ -76,8 +76,6 @@
         'decryptor_source.h',
         'encryptor_source.cc',
         'encryptor_source.h',
-        'fixed_encryptor_source.cc',
-        'fixed_encryptor_source.h',
         'limits.h',
         'media_parser.h',
         'media_sample.cc',

media/base/widevine_encryptor_source.cc

@@ -9,6 +9,7 @@
 #include "base/base64.h"
 #include "base/json/json_reader.h"
 #include "base/json/json_writer.h"
+#include "base/stl_util.h"
 #include "base/time/time.h"
 #include "base/threading/platform_thread.h"
 #include "base/values.h"
@@ -64,9 +65,9 @@ bool GetKeyAndKeyId(const base::DictionaryValue& track_dict,
   return true;
 }
 
-bool GetPssh(const base::DictionaryValue& track_dict,
+bool GetPsshData(const base::DictionaryValue& track_dict,
-             std::vector<uint8>* pssh) {
+                 std::vector<uint8>* pssh_data) {
-  DCHECK(pssh);
+  DCHECK(pssh_data);
 
   const base::ListValue* pssh_list;
   RCHECK(track_dict.GetList("pssh", &pssh_list));
@@ -82,11 +83,11 @@ bool GetPssh(const base::DictionaryValue& track_dict,
     LOG(ERROR) << "Expecting drm_type 'WIDEVINE', get '" << drm_type << "'.";
     return false;
   }
-  std::string pssh_base64_string;
+  std::string pssh_data_base64_string;
-  RCHECK(pssh_dict->GetString("data", &pssh_base64_string));
+  RCHECK(pssh_dict->GetString("data", &pssh_data_base64_string));
 
-  VLOG(2) << "Pssh:" << pssh_base64_string;
+  VLOG(2) << "Pssh Data:" << pssh_data_base64_string;
-  RCHECK(Base64StringToBytes(pssh_base64_string, pssh));
+  RCHECK(Base64StringToBytes(pssh_data_base64_string, pssh_data));
   return true;
 }
 
@@ -97,18 +98,47 @@ namespace media {
 WidevineEncryptorSource::WidevineEncryptorSource(
     const std::string& server_url,
     const std::string& content_id,
-    TrackType track_type,
     scoped_ptr<RequestSigner> signer)
     : http_fetcher_(new SimpleHttpFetcher()),
       server_url_(server_url),
       content_id_(content_id),
-      track_type_(track_type),
+      signer_(signer.Pass()),
-      signer_(signer.Pass()) {
+      key_fetched_(false) {
   DCHECK(signer_);
 }
-WidevineEncryptorSource::~WidevineEncryptorSource() {}
+WidevineEncryptorSource::~WidevineEncryptorSource() {
+  STLDeleteValues(&encryption_key_map_);
+}
 
-Status WidevineEncryptorSource::Initialize() {
+Status WidevineEncryptorSource::GetKey(TrackType track_type,
+                                       EncryptionKey* key) {
+  DCHECK(track_type == TRACK_TYPE_SD || track_type == TRACK_TYPE_HD ||
+         track_type == TRACK_TYPE_AUDIO);
+  Status status;
+  if (!key_fetched_) {
+    base::AutoLock auto_lock(lock_);
+    if (!key_fetched_) {
+      status = FetchKeys();
+      if (status.ok())
+        key_fetched_ = true;
+    }
+  }
+  if (!status.ok())
+    return status;
+  if (encryption_key_map_.find(track_type) == encryption_key_map_.end()) {
+    return Status(error::INTERNAL_ERROR,
+                  "Cannot find key of type " + TrackTypeToString(track_type));
+  }
+  *key = *encryption_key_map_[track_type];
+  return Status::OK;
+}
+
+void WidevineEncryptorSource::set_http_fetcher(
+    scoped_ptr<HttpFetcher> http_fetcher) {
+  http_fetcher_ = http_fetcher.Pass();
+}
+
+Status WidevineEncryptorSource::FetchKeys() {
   std::string request;
   FillRequest(content_id_, &request);
 
@@ -156,24 +186,6 @@ Status WidevineEncryptorSource::Initialize() {
                 "Failed to recover from server internal error.");
 }
 
-WidevineEncryptorSource::TrackType
-WidevineEncryptorSource::GetTrackTypeFromString(
-    const std::string& track_type_string) {
-  if (track_type_string == "SD")
-    return TRACK_TYPE_SD;
-  if (track_type_string == "HD")
-    return TRACK_TYPE_HD;
-  if (track_type_string == "AUDIO")
-    return TRACK_TYPE_AUDIO;
-  LOG(WARNING) << "Unexpected track type: " << track_type_string;
-  return TRACK_TYPE_UNKNOWN;
-}
-
-void WidevineEncryptorSource::set_http_fetcher(
-    scoped_ptr<HttpFetcher> http_fetcher) {
-  http_fetcher_ = http_fetcher.Pass();
-}
-
 void WidevineEncryptorSource::FillRequest(const std::string& content_id,
                                           std::string* request) {
   DCHECK(request);
@@ -252,11 +264,6 @@ bool WidevineEncryptorSource::DecodeResponse(const std::string& raw_response,
   return true;
 }
 
-bool WidevineEncryptorSource::IsExpectedTrackType(
-    const std::string& track_type_string) {
-  return track_type_ == GetTrackTypeFromString(track_type_string);
-}
-
 bool WidevineEncryptorSource::ExtractEncryptionKey(const std::string& response,
                                                    bool* transient_error) {
   DCHECK(transient_error);
@@ -281,33 +288,28 @@ bool WidevineEncryptorSource::ExtractEncryptionKey(const std::string& response,
 
   const base::ListValue* tracks;
   RCHECK(license_dict->GetList("tracks", &tracks));
+  RCHECK(tracks->GetSize() >= NUM_VALID_TRACK_TYPES);
 
-  for (base::ListValue::const_iterator it = tracks->begin();
+  for (size_t i = 0; i < tracks->GetSize(); ++i) {
-       it != tracks->end();
-       ++it) {
     const base::DictionaryValue* track_dict;
-    RCHECK((*it)->GetAsDictionary(&track_dict));
+    RCHECK(tracks->GetDictionary(i, &track_dict));
 
-    std::string track_type;
+    std::string track_type_str;
-    RCHECK(track_dict->GetString("type", &track_type));
+    RCHECK(track_dict->GetString("type", &track_type_str));
-    if (!IsExpectedTrackType(track_type))
+    TrackType track_type = GetTrackTypeFromString(track_type_str);
-      continue;
+    DCHECK_NE(TRACK_TYPE_UNKNOWN, track_type);
+    RCHECK(encryption_key_map_.find(track_type) == encryption_key_map_.end());
 
-    std::vector<uint8> key_id;
+    scoped_ptr<EncryptionKey> encryption_key(new EncryptionKey());
-    std::vector<uint8> key;
+    std::vector<uint8> pssh_data;
-    std::vector<uint8> pssh;
+    if (!GetKeyAndKeyId(*track_dict, &encryption_key->key,
-    if (!GetKeyAndKeyId(*track_dict, &key, &key_id) ||
+                        &encryption_key->key_id) ||
-        !GetPssh(*track_dict, &pssh))
+        !GetPsshData(*track_dict, &pssh_data))
       return false;
-
+    encryption_key->pssh = PsshBoxFromPsshData(pssh_data);
-    set_key_id(key_id);
+    encryption_key_map_[track_type] = encryption_key.release();
-    set_key(key);
-    set_pssh(pssh);
-    return true;
   }
-  LOG(ERROR) << "Cannot find key of type " << track_type_ << " from '"
+  return true;
-             << response << "'.";
-  return false;
 }
 
 }  // namespace media

media/base/widevine_encryptor_source.h

@@ -7,8 +7,11 @@
 #ifndef MEDIA_BASE_WIDEVINE_ENCRYPTOR_SOURCE_H_
 #define MEDIA_BASE_WIDEVINE_ENCRYPTOR_SOURCE_H_
 
+#include <map>
+
 #include "base/basictypes.h"
 #include "base/memory/scoped_ptr.h"
+#include "base/synchronization/lock.h"
 #include "media/base/encryptor_source.h"
 
 namespace media {
@@ -19,34 +22,25 @@ class RequestSigner;
 /// Encryptor source which talks to the Widevine encryption service.
 class WidevineEncryptorSource : public EncryptorSource {
  public:
-  enum TrackType {
-    TRACK_TYPE_UNKNOWN = 0,
-    TRACK_TYPE_SD,
-    TRACK_TYPE_HD,
-    TRACK_TYPE_AUDIO
-  };
-
   /// @param server_url is the Widevine common encryption server url.
   /// @param content_id the unique id identify the content to be encrypted.
-  /// @param track_type is the content type, can be AUDIO, SD or HD.
   /// @param signer must not be NULL.
   WidevineEncryptorSource(const std::string& server_url,
                           const std::string& content_id,
-                          TrackType track_type,
                           scoped_ptr<RequestSigner> signer);
   virtual ~WidevineEncryptorSource();
 
   /// EncryptorSource implementation override.
-  virtual Status Initialize() OVERRIDE;
+  virtual Status GetKey(TrackType track_type, EncryptionKey* key) OVERRIDE;
 
   /// Inject an @b HttpFetcher object, mainly used for testing.
   /// @param http_fetcher points to the @b HttpFetcher object to be injected.
   void set_http_fetcher(scoped_ptr<HttpFetcher> http_fetcher);
 
-  static WidevineEncryptorSource::TrackType GetTrackTypeFromString(
-      const std::string& track_type_string);
-
  private:
+  // Fetch keys from server.
+  Status FetchKeys();
+
   // Fill |request| with necessary fields for Widevine encryption request.
   // |request| should not be NULL.
   void FillRequest(const std::string& content_id, std::string* request);
@@ -56,7 +50,6 @@ class WidevineEncryptorSource : public EncryptorSource {
   // Decode |response| from JSON formatted |raw_response|.
   // |response| should not be NULL.
   bool DecodeResponse(const std::string& raw_response, std::string* response);
-  bool IsExpectedTrackType(const std::string& track_type_string);
   // Extract encryption key from |response|, which is expected to be properly
   // formatted. |transient_error| will be set to true if it fails and the
   // failure is because of a transient error from the server. |transient_error|
@@ -70,9 +63,12 @@ class WidevineEncryptorSource : public EncryptorSource {
   scoped_ptr<HttpFetcher> http_fetcher_;
   std::string server_url_;
   std::string content_id_;
-  TrackType track_type_;
   scoped_ptr<RequestSigner> signer_;
 
+  mutable base::Lock lock_;
+  bool key_fetched_;  // Protected by lock_;
+  std::map<TrackType, EncryptionKey*> encryption_key_map_;
+
   DISALLOW_COPY_AND_ASSIGN(WidevineEncryptorSource);
 };
 

media/base/widevine_encryptor_source_unittest.cc

@@ -17,16 +17,11 @@
 namespace {
 const char kServerUrl[] = "http://www.foo.com/getcontentkey";
 const char kContentId[] = "ContentFoo";
-const char kTrackType[] = "SD";
 const char kSignerName[] = "SignerFoo";
 
 const char kMockSignature[] = "MockSignature";
 
-const char kMockKeyId[] = "MockKeyId";
+// The license service may return an error indicating a transient error has
-const char kMockKey[] = "MockKey";
-const char kMockPsshData[] = "MockPsshData";
-
-// The lisence service may return an error indicating a transient error has
 // just happened in the server, or other types of errors.
 // WidevineEncryptorSource will perform a number of retries on transient errors;
 // WidevineEncryptorSource does not know about other errors and retries are not
@@ -39,11 +34,11 @@ const char kExpectedRequestMessageFormat[] =
     "\"tracks\":[{\"type\":\"SD\"},{\"type\":\"HD\"},{\"type\":\"AUDIO\"}]}";
 const char kExpectedSignedMessageFormat[] =
     "{\"request\":\"%s\",\"signature\":\"%s\",\"signer\":\"%s\"}";
-const char kLicenseOkResponseFormat[] =
+const char kTrackFormat[] =
-    "{\"status\":\"OK\",\"tracks\":[{\"type\":\"%s\",\"key_id\":\"%s\",\"key\":"
+    "{\"type\":\"%s\",\"key_id\":\"%s\",\"key\":"
-    "\"%s\",\"pssh\":[{\"drm_type\":\"WIDEVINE\",\"data\":\"%s\"}]}]}";
+    "\"%s\",\"pssh\":[{\"drm_type\":\"WIDEVINE\",\"data\":\"%s\"}]}";
-const char kLicenseErrorResponseFormat[] =
+const char kLicenseResponseFormat[] =
-    "{\"status\":\"%s\",\"drm\":[],\"tracks\":[]}";
+    "{\"status\":\"%s\",\"tracks\":[%s]}";
 const char kHttpResponseFormat[] = "{\"response\":\"%s\"}";
 
 std::string Base64Encode(const std::string& input) {
@@ -55,6 +50,41 @@ std::string Base64Encode(const std::string& input) {
 std::string ToString(const std::vector<uint8> v) {
   return std::string(v.begin(), v.end());
 }
+
+std::string GetMockKeyId(const std::string& track_type) {
+  return "MockKeyId" + track_type;
+}
+
+std::string GetMockKey(const std::string& track_type) {
+  return "MockKey" + track_type;
+}
+
+std::string GetMockPsshData(const std::string& track_type) {
+  return "MockPsshData" + track_type;
+}
+
+std::string GenerateMockLicenseResponse() {
+  const std::string kTrackTypes[] = {"SD", "HD", "AUDIO"};
+  std::string tracks;
+  for (size_t i = 0; i < 3; ++i) {
+    if (!tracks.empty())
+      tracks += ",";
+    tracks += base::StringPrintf(
+        kTrackFormat,
+        kTrackTypes[i].c_str(),
+        Base64Encode(GetMockKeyId(kTrackTypes[i])).c_str(),
+        Base64Encode(GetMockKey(kTrackTypes[i])).c_str(),
+        Base64Encode(GetMockPsshData(kTrackTypes[i])).c_str());
+  }
+  return base::StringPrintf(kLicenseResponseFormat, "OK", tracks.c_str());
+}
+
+std::string GetPsshDataFromPsshBox(const std::string& pssh_box) {
+  const size_t kPsshDataOffset = 32u;
+  DCHECK_LT(kPsshDataOffset, pssh_box.size());
+  return pssh_box.substr(kPsshDataOffset);
+}
+
 }  // namespace
 
 using ::testing::_;
@@ -103,7 +133,6 @@ class WidevineEncryptorSourceTest : public ::testing::Test {
     widevine_encryptor_source_.reset(new WidevineEncryptorSource(
         kServerUrl,
         kContentId,
-        WidevineEncryptorSource::GetTrackTypeFromString(kTrackType),
         mock_request_signer_.PassAs<RequestSigner>()));
     widevine_encryptor_source_->set_http_fetcher(
         mock_http_fetcher_.PassAs<HttpFetcher>());
@@ -118,14 +147,14 @@ class WidevineEncryptorSourceTest : public ::testing::Test {
 };
 
 TEST_F(WidevineEncryptorSourceTest, GetTrackTypeFromString) {
-  EXPECT_EQ(WidevineEncryptorSource::TRACK_TYPE_SD,
+  EXPECT_EQ(EncryptorSource::TRACK_TYPE_SD,
-            WidevineEncryptorSource::GetTrackTypeFromString("SD"));
+            EncryptorSource::GetTrackTypeFromString("SD"));
-  EXPECT_EQ(WidevineEncryptorSource::TRACK_TYPE_HD,
+  EXPECT_EQ(EncryptorSource::TRACK_TYPE_HD,
-            WidevineEncryptorSource::GetTrackTypeFromString("HD"));
+            EncryptorSource::GetTrackTypeFromString("HD"));
-  EXPECT_EQ(WidevineEncryptorSource::TRACK_TYPE_AUDIO,
+  EXPECT_EQ(EncryptorSource::TRACK_TYPE_AUDIO,
-            WidevineEncryptorSource::GetTrackTypeFromString("AUDIO"));
+            EncryptorSource::GetTrackTypeFromString("AUDIO"));
-  EXPECT_EQ(WidevineEncryptorSource::TRACK_TYPE_UNKNOWN,
+  EXPECT_EQ(EncryptorSource::TRACK_TYPE_UNKNOWN,
-            WidevineEncryptorSource::GetTrackTypeFromString("FOO"));
+            EncryptorSource::GetTrackTypeFromString("FOO"));
 }
 
 TEST_F(WidevineEncryptorSourceTest, GeneratureSignatureFailure) {
@@ -133,8 +162,10 @@ TEST_F(WidevineEncryptorSourceTest, GeneratureSignatureFailure) {
       .WillOnce(Return(false));
 
   CreateWidevineEncryptorSource();
+  EncryptionKey encryption_key;
   ASSERT_EQ(Status(error::INTERNAL_ERROR, "Signature generation failed."),
-            widevine_encryptor_source_->Initialize());
+            widevine_encryptor_source_->GetKey(EncryptorSource::TRACK_TYPE_SD,
+                                               &encryption_key));
 }
 
 // Check whether expected request message and post data was generated and
@@ -155,30 +186,35 @@ TEST_F(WidevineEncryptorSourceTest, HttpPostFailure) {
       .WillOnce(Return(kMockStatus));
 
   CreateWidevineEncryptorSource();
-  ASSERT_EQ(kMockStatus, widevine_encryptor_source_->Initialize());
+  EncryptionKey encryption_key;
+  ASSERT_EQ(kMockStatus,
+            widevine_encryptor_source_->GetKey(EncryptorSource::TRACK_TYPE_SD,
+                                               &encryption_key));
 }
 
 TEST_F(WidevineEncryptorSourceTest, LicenseStatusOK) {
   EXPECT_CALL(*mock_request_signer_, GenerateSignature(_, _))
       .WillOnce(Return(true));
 
-  std::string mock_license_status =
-      base::StringPrintf(kLicenseOkResponseFormat,
-                         kTrackType,
-                         Base64Encode(kMockKeyId).c_str(),
-                         Base64Encode(kMockKey).c_str(),
-                         Base64Encode(kMockPsshData).c_str());
   std::string expected_response = base::StringPrintf(
-      kHttpResponseFormat, Base64Encode(mock_license_status).c_str());
+      kHttpResponseFormat, Base64Encode(GenerateMockLicenseResponse()).c_str());
 
   EXPECT_CALL(*mock_http_fetcher_, Post(_, _, _))
       .WillOnce(DoAll(SetArgPointee<2>(expected_response), Return(Status::OK)));
 
   CreateWidevineEncryptorSource();
-  ASSERT_OK(widevine_encryptor_source_->Initialize());
+
-  EXPECT_EQ(kMockKeyId, ToString(widevine_encryptor_source_->key_id()));
+  EncryptionKey encryption_key;
-  EXPECT_EQ(kMockKey, ToString(widevine_encryptor_source_->key()));
+  const std::string kTrackTypes[] = {"SD", "HD", "AUDIO"};
-  EXPECT_EQ(kMockPsshData, ToString(widevine_encryptor_source_->pssh()));
+  for (size_t i = 0; i < 3; ++i) {
+    ASSERT_OK(widevine_encryptor_source_->GetKey(
+        EncryptorSource::GetTrackTypeFromString(kTrackTypes[i]),
+        &encryption_key));
+    EXPECT_EQ(GetMockKeyId(kTrackTypes[i]), ToString(encryption_key.key_id));
+    EXPECT_EQ(GetMockKey(kTrackTypes[i]), ToString(encryption_key.key));
+    EXPECT_EQ(GetMockPsshData(kTrackTypes[i]),
+              GetPsshDataFromPsshBox(ToString(encryption_key.pssh)));
+  }
 }
 
 TEST_F(WidevineEncryptorSourceTest, RetryOnTransientError) {
@@ -186,18 +222,12 @@ TEST_F(WidevineEncryptorSourceTest, RetryOnTransientError) {
       .WillOnce(Return(true));
 
   std::string mock_license_status = base::StringPrintf(
-      kLicenseErrorResponseFormat, kLicenseStatusTransientError);
+      kLicenseResponseFormat, kLicenseStatusTransientError, "");
   std::string expected_response = base::StringPrintf(
       kHttpResponseFormat, Base64Encode(mock_license_status).c_str());
 
-  std::string mock_retried_license_status =
-      base::StringPrintf(kLicenseOkResponseFormat,
-                         kTrackType,
-                         Base64Encode(kMockKeyId).c_str(),
-                         Base64Encode(kMockKey).c_str(),
-                         Base64Encode(kMockPsshData).c_str());
   std::string expected_retried_response = base::StringPrintf(
-      kHttpResponseFormat, Base64Encode(mock_retried_license_status).c_str());
+      kHttpResponseFormat, Base64Encode(GenerateMockLicenseResponse()).c_str());
 
   // Retry is expected on transient error.
   EXPECT_CALL(*mock_http_fetcher_, Post(_, _, _))
@@ -206,10 +236,13 @@ TEST_F(WidevineEncryptorSourceTest, RetryOnTransientError) {
                       Return(Status::OK)));
 
   CreateWidevineEncryptorSource();
-  ASSERT_OK(widevine_encryptor_source_->Initialize());
+  EncryptionKey encryption_key;
-  EXPECT_EQ(kMockKeyId, ToString(widevine_encryptor_source_->key_id()));
+  ASSERT_OK(widevine_encryptor_source_->GetKey(EncryptorSource::TRACK_TYPE_SD,
-  EXPECT_EQ(kMockKey, ToString(widevine_encryptor_source_->key()));
+                                               &encryption_key));
-  EXPECT_EQ(kMockPsshData, ToString(widevine_encryptor_source_->pssh()));
+  EXPECT_EQ(GetMockKeyId("SD"), ToString(encryption_key.key_id));
+  EXPECT_EQ(GetMockKey("SD"), ToString(encryption_key.key));
+  EXPECT_EQ(GetMockPsshData("SD"),
+            GetPsshDataFromPsshBox(ToString(encryption_key.pssh)));
 }
 
 TEST_F(WidevineEncryptorSourceTest, NoRetryOnUnknownError) {
@@ -217,7 +250,7 @@ TEST_F(WidevineEncryptorSourceTest, NoRetryOnUnknownError) {
       .WillOnce(Return(true));
 
   std::string mock_license_status = base::StringPrintf(
-      kLicenseErrorResponseFormat, kLicenseStatusUnknownError);
+      kLicenseResponseFormat, kLicenseStatusUnknownError, "");
   std::string mock_response = base::StringPrintf(
       kHttpResponseFormat, Base64Encode(mock_license_status).c_str());
 
@@ -225,8 +258,10 @@ TEST_F(WidevineEncryptorSourceTest, NoRetryOnUnknownError) {
       .WillOnce(DoAll(SetArgPointee<2>(mock_response), Return(Status::OK)));
 
   CreateWidevineEncryptorSource();
+  EncryptionKey encryption_key;
   ASSERT_EQ(error::SERVER_ERROR,
-            widevine_encryptor_source_->Initialize().error_code());
+            widevine_encryptor_source_->GetKey(EncryptorSource::TRACK_TYPE_SD,
+                                               &encryption_key).error_code());
 }
 
 }  // namespace media