Code clean-up in AesDecryptor and test.

- Added GetKey() and SetKey() in AesDecryptor.
- Moved static functions out of AesDecryptorTest to make them file scope functions.

BUG=none
TEST=media_unittests

Review URL: https://chromiumcodereview.appspot.com/10822013

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

3 files changed, 142 insertions, 135 deletions

diff --git a/media/crypto/aes_decryptor.cc b/media/crypto/aes_decryptor.cc
index 552197b..e585800 100644
--- a/media/crypto/aes_decryptor.cc
+++ b/media/crypto/aes_decryptor.cc
@@ -4,6 +4,8 @@
 
 #include "media/crypto/aes_decryptor.h"
 
+#include <vector>
+
 #include "base/logging.h"
 #include "base/stl_util.h"
 #include "base/string_number_conversions.h"
@@ -84,7 +86,7 @@ static bool CheckData(const DecoderBuffer& input,
 
 enum ClearBytesBufferSel {
   kSrcContainsClearBytes,
-  kDstContainsClearBytes,
+  kDstContainsClearBytes
 };
 
 static void CopySubsamples(const std::vector<SubsampleEntry>& subsamples,
@@ -250,16 +252,7 @@ void AesDecryptor::AddKey(const std::string& key_system,
     return;
   }
 
-  {
-    base::AutoLock auto_lock(key_map_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] = decryption_key.release();
-  }
-
+  SetKey(key_id_string, decryption_key.Pass());
   client_->KeyAdded(key_system, session_id);
 }
 
@@ -272,17 +265,10 @@ void AesDecryptor::Decrypt(const scoped_refptr<DecoderBuffer>& encrypted,
   CHECK(encrypted->GetDecryptConfig());
   const std::string& key_id = encrypted->GetDecryptConfig()->key_id();
 
-  DecryptionKey* key = NULL;
-  {
-    base::AutoLock auto_lock(key_map_lock_);
-    KeyMap::const_iterator found = key_map_.find(key_id);
-    if (found != key_map_.end())
-      key = found->second;
-  }
-
+  DecryptionKey* key = GetKey(key_id);
   if (!key) {
-    // TODO(fgalligan): Fire a need_key event here and add a test.
-    DVLOG(1) << "Could not find a matching key for given key ID.";
+    // TODO(xhwang): Fire a need_key event here and add a test.
+    DVLOG(1) << "Could not find a matching key for the given key ID.";
     decrypt_cb.Run(kError, NULL);
     return;
   }
@@ -317,8 +303,28 @@ void AesDecryptor::Decrypt(const scoped_refptr<DecoderBuffer>& encrypted,
   decrypt_cb.Run(kSuccess, decrypted);
 }
 
-AesDecryptor::DecryptionKey::DecryptionKey(
-    const std::string& secret)
+void AesDecryptor::SetKey(const std::string& key_id,
+                          scoped_ptr<DecryptionKey> decryption_key) {
+  base::AutoLock auto_lock(key_map_lock_);
+  KeyMap::iterator found = key_map_.find(key_id);
+  if (found != key_map_.end()) {
+    delete found->second;
+    key_map_.erase(found);
+  }
+  key_map_[key_id] = decryption_key.release();
+}
+
+AesDecryptor::DecryptionKey* AesDecryptor::GetKey(
+    const std::string& key_id) const {
+  base::AutoLock auto_lock(key_map_lock_);
+  KeyMap::const_iterator found = key_map_.find(key_id);
+  if (found == key_map_.end())
+    return NULL;
+
+  return found->second;
+}
+
+AesDecryptor::DecryptionKey::DecryptionKey(const std::string& secret)
     : secret_(secret) {
 }
 
@@ -326,28 +332,25 @@ AesDecryptor::DecryptionKey::~DecryptionKey() {}
 
 bool AesDecryptor::DecryptionKey::Init() {
   CHECK(!secret_.empty());
-  decryption_key_.reset(
-      crypto::SymmetricKey::Import(crypto::SymmetricKey::AES, secret_));
-  if (!decryption_key_.get()) {
+  decryption_key_.reset(crypto::SymmetricKey::Import(
+      crypto::SymmetricKey::AES, secret_));
+  if (!decryption_key_.get())
     return false;
-  }
 
   std::string raw_key = DeriveKey(secret_,
                                   kWebmEncryptionSeed,
                                   secret_.length());
-  if (raw_key.empty()) {
+  if (raw_key.empty())
     return false;
-  }
-  webm_decryption_key_.reset(
-      crypto::SymmetricKey::Import(crypto::SymmetricKey::AES, raw_key));
-  if (!webm_decryption_key_.get()) {
+
+  webm_decryption_key_.reset(crypto::SymmetricKey::Import(
+      crypto::SymmetricKey::AES, raw_key));
+  if (!webm_decryption_key_.get())
     return false;
-  }
 
   hmac_key_ = DeriveKey(secret_, kWebmHmacSeed, kWebmSha1DigestSize);
-  if (hmac_key_.empty()) {
+  if (hmac_key_.empty())
     return false;
-  }
 
   return true;
 }
diff --git a/media/crypto/aes_decryptor.h b/media/crypto/aes_decryptor.h
index 86b258d..8c666ac 100644
--- a/media/crypto/aes_decryptor.h
+++ b/media/crypto/aes_decryptor.h
@@ -89,6 +89,13 @@ class MEDIA_EXPORT AesDecryptor : public Decryptor {
     DISALLOW_COPY_AND_ASSIGN(DecryptionKey);
   };
 
+  // Sets |key| for |key_id|. The AesDecryptor takes the ownership of the |key|.
+  void SetKey(const std::string& key_id, scoped_ptr<DecryptionKey> key);
+
+  // Gets a DecryptionKey associated with |key_id|. The AesDecryptor still owns
+  // the key. Returns NULL if no key is associated with |key_id|.
+  DecryptionKey* GetKey(const std::string& key_id) const;
+
   // KeyMap owns the DecryptionKey* and must delete them when they are
   // not needed any more.
   typedef base::hash_map<std::string, DecryptionKey*> KeyMap;
@@ -97,7 +104,7 @@ class MEDIA_EXPORT AesDecryptor : public Decryptor {
   // protect |key_map_|, the only member variable that is shared between
   // Decrypt() and other methods.
   KeyMap key_map_;  // Protected by the |key_map_lock_|.
-  base::Lock key_map_lock_;  // Protects the |key_map_|.
+  mutable base::Lock key_map_lock_;  // Protects the |key_map_|.
 
   // Make session ID unique per renderer by making it static.
   // TODO(xhwang): Make session ID more strictly defined if needed:
diff --git a/media/crypto/aes_decryptor_unittest.cc b/media/crypto/aes_decryptor_unittest.cc
index 7877f6e..12df91a 100644
--- a/media/crypto/aes_decryptor_unittest.cc
+++ b/media/crypto/aes_decryptor_unittest.cc
@@ -3,6 +3,7 @@
 // found in the LICENSE file.
 
 #include <string>
+#include <vector>
 
 #include "base/basictypes.h"
 #include "base/bind.h"
@@ -48,19 +49,19 @@ const WebmEncryptedData kWebmEncryptedFrames[] = {
     // key_id
     { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
       0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
-      0x10, 0x11, 0x12, 0x13,
-      }, 20,
+      0x10, 0x11, 0x12, 0x13
+    }, 20,
     // key
     { 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b,
-      0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23,
-      }, 16,
+      0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23
+    }, 16,
     // encrypted_data
     { 0xfb, 0xe7, 0x1d, 0xbb, 0x4c, 0x23, 0xce, 0xba,
       0xcc, 0xf8, 0xda, 0xc0, 0xff, 0xff, 0xff, 0xff,
       0xff, 0xff, 0xff, 0xff, 0x99, 0xaa, 0xff, 0xb7,
       0x74, 0x02, 0x4e, 0x1c, 0x75, 0x3d, 0xee, 0xcb,
-      0x64, 0xf7,
-      }, 34,
+      0x64, 0xf7
+    }, 34
   },
   {
     // plaintext
@@ -68,40 +69,40 @@ const WebmEncryptedData kWebmEncryptedFrames[] = {
     // key_id
     { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
       0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
-      0x10, 0x11, 0x12, 0x13,
-      }, 20,
+      0x10, 0x11, 0x12, 0x13
+    }, 20,
     // key
     { 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b,
-      0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23,
-      }, 16,
+      0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23
+    }, 16,
     // encrypted_data
     { 0x43, 0xe4, 0x78, 0x7a, 0x43, 0xe1, 0x49, 0xbb,
       0x44, 0x38, 0xdf, 0xfc, 0x00, 0x00, 0x00, 0x00,
       0x00, 0x00, 0x00, 0x00, 0xec, 0x8e, 0x87, 0x21,
       0xd3, 0xb9, 0x1c, 0x61, 0xf6, 0x5a, 0x60, 0xaa,
       0x07, 0x0e, 0x96, 0xd0, 0x54, 0x5d, 0x35, 0x9a,
-      0x4a, 0xd3,
-      }, 42,
+      0x4a, 0xd3
+    }, 42
   },
   {
     // plaintext
     "Original data.", 14,
     // key_id
     { 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0x2b,
-      0x2c, 0x2d, 0x2e, 0x2f, 0x30,
-      }, 13,
+      0x2c, 0x2d, 0x2e, 0x2f, 0x30
+    }, 13,
     // key
     { 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38,
-      0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x3e, 0x3f, 0x40,
-      }, 16,
+      0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x3e, 0x3f, 0x40
+    }, 16,
     // encrypted_data
     { 0xd9, 0x43, 0x30, 0xfd, 0x82, 0x77, 0x62, 0x04,
       0x08, 0xc2, 0x48, 0x89, 0x00, 0x00, 0x00, 0x00,
       0x00, 0x00, 0x00, 0x01, 0x48, 0x5e, 0x4a, 0x41,
       0x2a, 0x8b, 0xf4, 0xc6, 0x47, 0x54, 0x90, 0x34,
-      0xf4, 0x8b,
-      }, 34,
-  },
+      0xf4, 0x8b
+    }, 34
+  }
 };
 
 static const uint8 kWebmWrongKey[] = {
@@ -113,7 +114,7 @@ static const uint8 kWebmWrongSizedKey[] = { 0x20, 0x20 };
 // This is the encrypted data from frame 0 of |kWebmEncryptedFrames| except
 // byte 0 is changed from 0xfb to 0xfc. Bytes 0-11 of WebM encrypted data
 // contains the HMAC.
-static const unsigned char kWebmFrame0HmacDataChanged[] = {
+static const uint8 kWebmFrame0HmacDataChanged[] = {
   0xfc, 0xe7, 0x1d, 0xbb, 0x4c, 0x23, 0xce, 0xba,
   0xcc, 0xf8, 0xda, 0xc0, 0xff, 0xff, 0xff, 0xff,
   0xff, 0xff, 0xff, 0xff, 0x99, 0xaa, 0xff, 0xb7,
@@ -124,7 +125,7 @@ static const unsigned char kWebmFrame0HmacDataChanged[] = {
 // This is the encrypted data from frame 0 of |kWebmEncryptedFrames| except
 // byte 12 is changed from 0xff to 0x0f. Bytes 12-19 of WebM encrypted data
 // contains the IV.
-static const unsigned char kWebmFrame0IvDataChanged[] = {
+static const uint8 kWebmFrame0IvDataChanged[] = {
   0xfb, 0xe7, 0x1d, 0xbb, 0x4c, 0x23, 0xce, 0xba,
   0xcc, 0xf8, 0xda, 0xc0, 0x0f, 0xff, 0xff, 0xff,
   0xff, 0xff, 0xff, 0xff, 0x99, 0xaa, 0xff, 0xb7,
@@ -135,7 +136,7 @@ static const unsigned char kWebmFrame0IvDataChanged[] = {
 // This is the encrypted data from frame 0 of |kWebmEncryptedFrames| except
 // byte 33 is changed from 0xf7 to 0xf8. Bytes 20+ of WebM encrypted data
 // contains the encrypted frame.
-static const unsigned char kWebmFrame0FrameDataChanged[] = {
+static const uint8 kWebmFrame0FrameDataChanged[] = {
   0xfb, 0xe7, 0x1d, 0xbb, 0x4c, 0x23, 0xce, 0xba,
   0xcc, 0xf8, 0xda, 0xc0, 0xff, 0xff, 0xff, 0xff,
   0xff, 0xff, 0xff, 0xff, 0x99, 0xaa, 0xff, 0xb7,
@@ -181,83 +182,85 @@ static const uint8 kNoSubsampleData[] = {
 static const SubsampleEntry kSubsampleEntries[] = {
   { 2, 7 },
   { 3, 11 },
-  { 1, 0 },
+  { 1, 0 }
 };
 
+// Returns a 16 byte CTR counter block. The CTR counter block format is a
+// CTR IV appended with a CTR block counter. |iv| is a CTR IV. |iv_size| is
+// the size of |iv| in bytes.
+static std::string GenerateCounterBlock(const uint8* iv, int iv_size) {
+  const int kDecryptionKeySize = 16;
+  CHECK_GT(iv_size, 0);
+  CHECK_LE(iv_size, kDecryptionKeySize);
+  char counter_block_data[kDecryptionKeySize];
+
+  // Set the IV.
+  memcpy(counter_block_data, iv, iv_size);
+
+  // Set block counter to all 0's.
+  memset(counter_block_data + iv_size, 0, kDecryptionKeySize - iv_size);
+
+  return std::string(counter_block_data, kDecryptionKeySize);
+}
+
+// Creates a WebM encrypted buffer that the demuxer would pass to the
+// decryptor. |data| is the payload of a WebM encrypted Block. |key_id| is
+// initialization data from the WebM file. Every encrypted Block has
+// an HMAC and IV prepended to an encrypted frame. Current encrypted WebM
+// request for comments specification is here
+// http://wiki.webmproject.org/encryption/webm-encryption-rfc
+static scoped_refptr<DecoderBuffer> CreateWebMEncryptedBuffer(
+    const uint8* data, int data_size,
+    const uint8* key_id, int key_id_size) {
+  scoped_refptr<DecoderBuffer> encrypted_buffer = DecoderBuffer::CopyFrom(
+      data + kWebMHmacSize, data_size - kWebMHmacSize);
+  CHECK(encrypted_buffer);
+
+  uint64 network_iv;
+  memcpy(&network_iv, data + kWebMHmacSize, sizeof(network_iv));
+  const uint64 iv = base::NetToHost64(network_iv);
+  std::string webm_iv =
+      GenerateCounterBlock(reinterpret_cast<const uint8*>(&iv), sizeof(iv));
+  encrypted_buffer->SetDecryptConfig(
+      scoped_ptr<DecryptConfig>(new DecryptConfig(
+          std::string(reinterpret_cast<const char*>(key_id), key_id_size),
+          webm_iv,
+          std::string(reinterpret_cast<const char*>(data), kWebMHmacSize),
+          sizeof(iv),
+          std::vector<SubsampleEntry>())));
+  return encrypted_buffer;
+}
+
+static scoped_refptr<DecoderBuffer> CreateSubsampleEncryptedBuffer(
+    const uint8* data, int data_size,
+    const uint8* key_id, int key_id_size,
+    const uint8* iv, int iv_size,
+    int data_offset,
+    const std::vector<SubsampleEntry>& subsample_entries) {
+  scoped_refptr<DecoderBuffer> encrypted_buffer =
+      DecoderBuffer::CopyFrom(data, data_size);
+  CHECK(encrypted_buffer);
+  encrypted_buffer->SetDecryptConfig(
+      scoped_ptr<DecryptConfig>(new DecryptConfig(
+          std::string(reinterpret_cast<const char*>(key_id), key_id_size),
+          std::string(reinterpret_cast<const char*>(iv), iv_size),
+          std::string(),
+          data_offset,
+          subsample_entries)));
+  return encrypted_buffer;
+}
+
 class AesDecryptorTest : public testing::Test {
  public:
   AesDecryptorTest()
       : decryptor_(&client_),
         decrypt_cb_(base::Bind(&AesDecryptorTest::BufferDecrypted,
-                               base::Unretained(this))) {
+                               base::Unretained(this))),
+        subsample_entries_(kSubsampleEntries,
+                           kSubsampleEntries + arraysize(kSubsampleEntries)) {
   }
 
  protected:
-  // Returns a 16 byte CTR counter block. The CTR counter block format is a
-  // CTR IV appended with a CTR block counter. |iv| is a CTR IV. |iv_size| is
-  // the size of |iv| in bytes.
-  static std::string GenerateCounterBlock(const uint8* iv, int iv_size) {
-    const int kDecryptionKeySize = 16;
-    CHECK_GT(iv_size, 0);
-    CHECK_LE(iv_size, kDecryptionKeySize);
-    char counter_block_data[kDecryptionKeySize];
-
-    // Set the IV.
-    memcpy(counter_block_data, iv, iv_size);
-
-    // Set block counter to all 0's.
-    memset(counter_block_data + iv_size, 0, kDecryptionKeySize - iv_size);
-
-    return std::string(counter_block_data, kDecryptionKeySize);
-  }
-
-  // Creates a WebM encrypted buffer that the demuxer would pass to the
-  // decryptor. |data| is the payload of a WebM encrypted Block. |key_id| is
-  // initialization data from the WebM file. Every encrypted Block has
-  // an HMAC and IV prepended to an encrypted frame. Current encrypted WebM
-  // request for comments specification is here
-  // http://wiki.webmproject.org/encryption/webm-encryption-rfc
-  scoped_refptr<DecoderBuffer> CreateWebMEncryptedBuffer(
-      const uint8* data, int data_size,
-      const uint8* key_id, int key_id_size) {
-    scoped_refptr<DecoderBuffer> encrypted_buffer = DecoderBuffer::CopyFrom(
-        data + kWebMHmacSize, data_size - kWebMHmacSize);
-    CHECK(encrypted_buffer);
-
-    uint64 network_iv;
-    memcpy(&network_iv, data + kWebMHmacSize, sizeof(network_iv));
-    const uint64 iv = base::NetToHost64(network_iv);
-    std::string webm_iv =
-        GenerateCounterBlock(reinterpret_cast<const uint8*>(&iv), sizeof(iv));
-    encrypted_buffer->SetDecryptConfig(
-        scoped_ptr<DecryptConfig>(new DecryptConfig(
-            std::string(reinterpret_cast<const char*>(key_id), key_id_size),
-            webm_iv,
-            std::string(reinterpret_cast<const char*>(data), kWebMHmacSize),
-            sizeof(iv),
-            std::vector<SubsampleEntry>())));
-    return encrypted_buffer;
-  }
-
-  scoped_refptr<DecoderBuffer> CreateSubsampleEncryptedBuffer(
-      const uint8* data, int data_size,
-      const uint8* key_id, int key_id_size,
-      const uint8* iv, int iv_size,
-      int data_offset,
-      const std::vector<SubsampleEntry>& subsample_entries) {
-    scoped_refptr<DecoderBuffer> encrypted_buffer =
-        DecoderBuffer::CopyFrom(data, data_size);
-    CHECK(encrypted_buffer);
-    encrypted_buffer->SetDecryptConfig(
-        scoped_ptr<DecryptConfig>(new DecryptConfig(
-            std::string(reinterpret_cast<const char*>(key_id), key_id_size),
-            std::string(reinterpret_cast<const char*>(iv), iv_size),
-            std::string(),
-            data_offset,
-            subsample_entries)));
-    return encrypted_buffer;
-  }
-
   void GenerateKeyRequest(const uint8* key_id, int key_id_size) {
     EXPECT_CALL(client_, KeyMessageMock(kClearKeySystem, StrNe(std::string()),
                                         NotNull(), Gt(0), ""))
@@ -300,11 +303,11 @@ class AesDecryptorTest : public testing::Test {
     decryptor_.Decrypt(encrypted, decrypt_cb_);
   }
 
-  scoped_refptr<DecoderBuffer> encrypted_data_;
   MockDecryptorClient client_;
   AesDecryptor decryptor_;
   std::string session_id_string_;
   AesDecryptor::DecryptCB decrypt_cb_;
+  std::vector<SubsampleEntry> subsample_entries_;
 };
 
 TEST_F(AesDecryptorTest, NormalWebMDecryption) {
@@ -438,9 +441,7 @@ TEST_F(AesDecryptorTest, SubsampleDecryption) {
       kSubsampleKeyId, arraysize(kSubsampleKeyId),
       kSubsampleIv, arraysize(kSubsampleIv),
       0,
-      std::vector<SubsampleEntry>(
-          kSubsampleEntries,
-          kSubsampleEntries + arraysize(kSubsampleEntries)));
+      subsample_entries_);
   ASSERT_NO_FATAL_FAILURE(DecryptAndExpectToSucceed(
       encrypted_data, kSubsampleOriginalData, kSubsampleOriginalDataSize));
 }
@@ -457,9 +458,7 @@ TEST_F(AesDecryptorTest, SubsampleDecryptionWithOffset) {
       kSubsampleKeyId, arraysize(kSubsampleKeyId),
       kSubsampleIv, arraysize(kSubsampleIv),
       arraysize(kPaddedSubsampleData) - arraysize(kSubsampleData),
-      std::vector<SubsampleEntry>(
-          kSubsampleEntries,
-          kSubsampleEntries + arraysize(kSubsampleEntries)));
+      subsample_entries_);
   ASSERT_NO_FATAL_FAILURE(DecryptAndExpectToSucceed(
       encrypted_data, kSubsampleOriginalData, kSubsampleOriginalDataSize));
 }
@@ -483,9 +482,7 @@ TEST_F(AesDecryptorTest, IncorrectSubsampleSize) {
   GenerateKeyRequest(kSubsampleKeyId, arraysize(kSubsampleKeyId));
   AddKeyAndExpectToSucceed(kSubsampleKeyId, arraysize(kSubsampleKeyId),
                            kSubsampleKey, arraysize(kSubsampleKey));
-  std::vector<SubsampleEntry> entries(
-      kSubsampleEntries,
-      kSubsampleEntries + arraysize(kSubsampleEntries));
+  std::vector<SubsampleEntry> entries = subsample_entries_;
   entries[2].cypher_bytes += 1;
 
   scoped_refptr<DecoderBuffer> encrypted_data = CreateSubsampleEncryptedBuffer(