Add AddKey to AesDecryptor

BUG=117063
TEST=updated media_unittest:AesDecryptorTest.*

Review URL: http://codereview.chromium.org/9390024

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@126576 0039d316-1c4b-4281-b951-d872f2087c98

3 files changed, 127 insertions, 41 deletions

diff --git a/media/crypto/aes_decryptor.cc b/media/crypto/aes_decryptor.cc
index 8e83daf..8d56d76 100644
--- a/media/crypto/aes_decryptor.cc
+++ b/media/crypto/aes_decryptor.cc
@@ -4,9 +4,9 @@
 
 #include "media/crypto/aes_decryptor.h"
 
-#include <string>
-
 #include "base/logging.h"
+#include "base/stl_util.h"
+#include "base/string_piece.h"
 #include "crypto/encryptor.h"
 #include "crypto/symmetric_key.h"
 #include "media/base/buffers.h"
@@ -15,58 +15,94 @@
 
 namespace media {
 
-static const char* kInitialCounter = "0000000000000000";
+// TODO(xhwang): Get real IV from frames.
+static const char kInitialCounter[] = "0000000000000000";
 
-// Decrypts |input| using |raw_key|, which is the binary data for the decryption
-// key.
+// Decrypt |input| using |key|.
 // Return a scoped_refptr to a Buffer object with the decrypted data on success.
 // Return a scoped_refptr to NULL if the data could not be decrypted.
-// TODO(xhwang): Both the input and output are copied! Any performance concern?
 static scoped_refptr<Buffer> DecryptData(const Buffer& input,
-                                         const uint8* raw_key,
-                                         int raw_key_size) {
-  CHECK(raw_key && raw_key_size > 0);
+                                         crypto::SymmetricKey* key) {
   CHECK(input.GetDataSize());
-
-  scoped_ptr<crypto::SymmetricKey> key(crypto::SymmetricKey::Import(
-      crypto::SymmetricKey::AES,
-      std::string(reinterpret_cast<const char*>(raw_key), raw_key_size)));
-  if (!key.get()) {
-    DVLOG(1) << "Could not import key.";
-    return NULL;
-  }
+  CHECK(key);
 
   // Initialize encryption data.
   // The IV must be exactly as long as the cipher block size.
   crypto::Encryptor encryptor;
-  if (!encryptor.Init(key.get(), crypto::Encryptor::CBC, kInitialCounter)) {
+  if (!encryptor.Init(key, crypto::Encryptor::CBC, kInitialCounter)) {
     DVLOG(1) << "Could not initialize encryptor.";
     return NULL;
   }
 
-  std::string decrypt_text;
-  std::string encrypted_text(reinterpret_cast<const char*>(input.GetData()),
-                             input.GetDataSize());
-  if (!encryptor.Decrypt(encrypted_text, &decrypt_text)) {
+  std::string decrypted_text;
+  base::StringPiece encrypted_text(
+      reinterpret_cast<const char*>(input.GetData()),
+      input.GetDataSize());
+  if (!encryptor.Decrypt(encrypted_text, &decrypted_text)) {
     DVLOG(1) << "Could not decrypt data.";
     return NULL;
   }
 
+  // TODO(xhwang): Implement a string based Buffer implementation to avoid
+  // data copying.
   return DataBuffer::CopyFrom(
-      reinterpret_cast<const uint8*>(decrypt_text.data()),
-      decrypt_text.size());
+      reinterpret_cast<const uint8*>(decrypted_text.data()),
+      decrypted_text.size());
 }
 
 AesDecryptor::AesDecryptor() {}
 
+AesDecryptor::~AesDecryptor() {
+  STLDeleteValues(&key_map_);
+}
+
+void AesDecryptor::AddKey(const uint8* key_id, int key_id_size,
+                          const uint8* key, int key_size) {
+  CHECK(key_id && key);
+  CHECK_GT(key_id_size, 0);
+  CHECK_GT(key_size, 0);
+
+  std::string key_id_string(reinterpret_cast<const char*>(key_id), key_id_size);
+  std::string key_string(reinterpret_cast<const char*>(key) , key_size);
+
+  crypto::SymmetricKey* symmetric_key = crypto::SymmetricKey::Import(
+      crypto::SymmetricKey::AES, key_string);
+  if (!symmetric_key) {
+    DVLOG(1) << "Could not import key.";
+    return;
+  }
+
+  base::AutoLock auto_lock(lock_);
+  KeyMap::iterator found = key_map_.find(key_id_string);
+  if (found != key_map_.end()) {
+    delete found->second;
+    key_map_.erase(found);
+  }
+  key_map_[key_id_string] = symmetric_key;
+}
+
 scoped_refptr<Buffer> AesDecryptor::Decrypt(
     const scoped_refptr<Buffer>& encrypted) {
   CHECK(encrypted->GetDecryptConfig());
+  const uint8* key_id = encrypted->GetDecryptConfig()->key_id();
+  const int key_id_size = encrypted->GetDecryptConfig()->key_id_size();
+
+  // TODO(xhwang): Avoid always constructing a string with StringPiece?
+  std::string key_id_string(reinterpret_cast<const char*>(key_id), key_id_size);
+
+  crypto::SymmetricKey* key = NULL;
+  {
+    base::AutoLock auto_lock(lock_);
+    KeyMap::const_iterator found = key_map_.find(key_id_string);
+    if (found == key_map_.end()) {
+      DVLOG(1) << "Could not find a matching key for given key ID.";
+      return NULL;
+    }
+    key = found->second;
+  }
+
+  scoped_refptr<Buffer> decrypted = DecryptData(*encrypted, key);
 
-  // For now, the key is the key ID.
-  const uint8* key = encrypted->GetDecryptConfig()->key_id();
-  int key_size = encrypted->GetDecryptConfig()->key_id_size();
-  scoped_refptr<Buffer> decrypted = DecryptData(*encrypted, key, key_size);
   if (decrypted) {
     decrypted->SetTimestamp(encrypted->GetTimestamp());
     decrypted->SetDuration(encrypted->GetDuration());
diff --git a/media/crypto/aes_decryptor.h b/media/crypto/aes_decryptor.h
index 32d6802..212c9ac 100644
--- a/media/crypto/aes_decryptor.h
+++ b/media/crypto/aes_decryptor.h
@@ -5,10 +5,18 @@
 #ifndef MEDIA_CRYPTO_AES_DECRYPTOR_H_
 #define MEDIA_CRYPTO_AES_DECRYPTOR_H_
 
+#include <string>
+
 #include "base/basictypes.h"
+#include "base/hash_tables.h"
 #include "base/memory/ref_counted.h"
+#include "base/synchronization/lock.h"
 #include "media/base/media_export.h"
 
+namespace crypto {
+class SymmetricKey;
+}
+
 namespace media {
 
 class Buffer;
@@ -17,6 +25,14 @@ class Buffer;
 class MEDIA_EXPORT AesDecryptor {
  public:
   AesDecryptor();
+  ~AesDecryptor();
+
+  // Add a |key_id| and |key| pair to the key system. The key is not limited to
+  // a decryption key. It can be any data that the key system accepts, such as
+  // a license. If multiple calls of this function set different keys for the
+  // same |key_id|, the older key will be replaced by the newer key.
+  void AddKey(const uint8* key_id, int key_id_size,
+              const uint8* key, int key_size);
 
   // Decrypt |input| buffer. The |input| should not be NULL.
   // Return a Buffer that contains decrypted data if decryption succeeded.
@@ -24,6 +40,12 @@ class MEDIA_EXPORT AesDecryptor {
   scoped_refptr<Buffer> Decrypt(const scoped_refptr<Buffer>& input);
 
  private:
+  // KeyMap owns the crypto::SymmetricKey* and must delete them when they are
+  // not needed any more.
+  typedef base::hash_map<std::string, crypto::SymmetricKey*> KeyMap;
+  KeyMap key_map_;
+  base::Lock lock_;
+
   DISALLOW_COPY_AND_ASSIGN(AesDecryptor);
 };
 
diff --git a/media/crypto/aes_decryptor_unittest.cc b/media/crypto/aes_decryptor_unittest.cc
index 0143bd0..dceeb5a 100644
--- a/media/crypto/aes_decryptor_unittest.cc
+++ b/media/crypto/aes_decryptor_unittest.cc
@@ -11,7 +11,7 @@
 
 namespace media {
 
-// |kEncryptedDataHex| is encrypted from |kOriginalData| using |kRawKey|, whose
+// |kEncryptedData| is encrypted from |kOriginalData| using |kRightKey|, whose
 // length is |kKeySize|. Modifying any of these independently would fail the
 // test.
 static const char kOriginalData[] = "Original data.";
@@ -19,8 +19,11 @@ static const int kEncryptedDataSize = 16;
 static const unsigned char kEncryptedData[] =
     "\x82\x3A\x76\x92\xEC\x7F\xF8\x85\xEC\x23\x52\xFB\x19\xB1\xB9\x09";
 static const int kKeySize = 16;
-static const unsigned char kRawKey[] = "A wonderful key!";
+static const unsigned char kRightKey[] = "A wonderful key!";
 static const unsigned char kWrongKey[] = "I'm a wrong key.";
+static const int kKeyIdSize = 9;
+static const unsigned char kKeyId1[] = "Key ID 1.";
+static const unsigned char kKeyId2[] = "Key ID 2.";
 
 class AesDecryptorTest : public testing::Test {
  public:
@@ -29,9 +32,22 @@ class AesDecryptorTest : public testing::Test {
   }
 
  protected:
-  void SetKey(const uint8* key, int key_size) {
+  void SetKeyIdForEncryptedData(const uint8* key_id, int key_id_size) {
     encrypted_data_->SetDecryptConfig(
-        scoped_ptr<DecryptConfig>(new DecryptConfig(key, key_size)));
+        scoped_ptr<DecryptConfig>(new DecryptConfig(key_id, key_id_size)));
+  }
+
+  void DecryptAndExpectToSucceed() {
+    scoped_refptr<Buffer> decrypted = decryptor_.Decrypt(encrypted_data_);
+    ASSERT_TRUE(decrypted);
+    size_t data_length = sizeof(kOriginalData) - 1;
+    ASSERT_EQ(data_length, decrypted->GetDataSize());
+    EXPECT_EQ(0, memcmp(kOriginalData, decrypted->GetData(), data_length));
+  }
+
+  void DecryptAndExpectToFail() {
+    scoped_refptr<Buffer> decrypted = decryptor_.Decrypt(encrypted_data_);
+    EXPECT_FALSE(decrypted);
   }
 
   scoped_refptr<DataBuffer> encrypted_data_;
@@ -39,18 +55,30 @@ class AesDecryptorTest : public testing::Test {
 };
 
 TEST_F(AesDecryptorTest, NormalDecryption) {
-  SetKey(kRawKey, kKeySize);
-  scoped_refptr<Buffer> decrypted_data = decryptor_.Decrypt(encrypted_data_);
-  ASSERT_TRUE(decrypted_data.get());
-  size_t data_length = sizeof(kOriginalData) - 1;
-  ASSERT_EQ(data_length, decrypted_data->GetDataSize());
-  ASSERT_EQ(0, memcmp(kOriginalData, decrypted_data->GetData(), data_length));
+  decryptor_.AddKey(kKeyId1, kKeyIdSize, kRightKey, kKeySize);
+  SetKeyIdForEncryptedData(kKeyId1, kKeyIdSize);
+  ASSERT_NO_FATAL_FAILURE(DecryptAndExpectToSucceed());
 }
 
 TEST_F(AesDecryptorTest, WrongKey) {
-  SetKey(kWrongKey, kKeySize);
-  scoped_refptr<Buffer> decrypted_data = decryptor_.Decrypt(encrypted_data_);
-  EXPECT_FALSE(decrypted_data.get());
+  decryptor_.AddKey(kKeyId1, kKeyIdSize, kWrongKey, kKeySize);
+  SetKeyIdForEncryptedData(kKeyId1, kKeyIdSize);
+  ASSERT_NO_FATAL_FAILURE(DecryptAndExpectToFail());
+}
+
+TEST_F(AesDecryptorTest, MultipleKeys) {
+  decryptor_.AddKey(kKeyId1, kKeyIdSize, kRightKey, kKeySize);
+  decryptor_.AddKey(kKeyId2, kKeyIdSize, kWrongKey, kKeySize);
+  SetKeyIdForEncryptedData(kKeyId1, kKeyIdSize);
+  ASSERT_NO_FATAL_FAILURE(DecryptAndExpectToSucceed());
+}
+
+TEST_F(AesDecryptorTest, KeyReplacement) {
+  SetKeyIdForEncryptedData(kKeyId1, kKeyIdSize);
+  decryptor_.AddKey(kKeyId1, kKeyIdSize, kWrongKey, kKeySize);
+  ASSERT_NO_FATAL_FAILURE(DecryptAndExpectToFail());
+  decryptor_.AddKey(kKeyId1, kKeyIdSize, kRightKey, kKeySize);
+  ASSERT_NO_FATAL_FAILURE(DecryptAndExpectToSucceed());
 }
 
 }  // media