SAMPLE-AES encryption in PesPacketGenerator

- Allow setting a key on PesPacketGenerator.
- Add SAMPLE AES mode to AesPatternCryptor.

Change-Id: Iace2381b4cd3bbce63cd5bdb4cdf3a7cea47537a

packager/media/base/aes_pattern_cryptor.cc

@@ -15,11 +15,13 @@ namespace media {
 
 AesPatternCryptor::AesPatternCryptor(uint8_t crypt_byte_block,
                                      uint8_t skip_byte_block,
+                                     PatternEncryptionMode encryption_mode,
                                      ConstantIvFlag constant_iv_flag,
                                      scoped_ptr<AesCryptor> cryptor)
     : AesCryptor(constant_iv_flag),
       crypt_byte_block_(crypt_byte_block),
       skip_byte_block_(skip_byte_block),
+      encryption_mode_(encryption_mode),
       cryptor_(cryptor.Pass()) {
   DCHECK(cryptor_);
 }
@@ -45,7 +47,7 @@ bool AesPatternCryptor::CryptInternal(const uint8_t* text,
 
   while (text_size > 0) {
     const size_t crypt_byte_size = crypt_byte_block_ * AES_BLOCK_SIZE;
-    if (text_size >= crypt_byte_size) {
+    if (NeedEncrypt(text_size, crypt_byte_size)) {
       if (!cryptor_->Crypt(text, crypt_byte_size, crypt_text))
         return false;
     } else {
@@ -71,5 +73,12 @@ void AesPatternCryptor::SetIvInternal() {
   CHECK(cryptor_->SetIv(iv()));
 }
 
+bool AesPatternCryptor::NeedEncrypt(size_t input_size,
+                                    size_t target_data_size) {
+  if (encryption_mode_ == kSkipIfCryptByteBlockRemaining)
+    return input_size > target_data_size;
+  return input_size >= target_data_size;
+}
+
 }  // namespace media
 }  // namespace edash_packager

packager/media/base/aes_pattern_cryptor.h

@@ -15,10 +15,31 @@ namespace media {
 /// Implements pattern-based encryption/decryption.
 class AesPatternCryptor : public AesCryptor {
  public:
+  /// Enumerator for controling encrytion/decryption mode for the last
+  /// encryption/decrytion block(s).
+  enum PatternEncryptionMode {
+    /// Use kEncryptIfCryptByteBlockRemaining if the last blocks is exactly the
+    /// same as the number of remaining bytes. IOW
+    /// if (remaining_bytes == encryption_block_bytes)
+    ///   encrypt(remaining_data)
+    kEncryptIfCryptByteBlockRemaining,
+    /// Use kSkipIfCryptByteBlockRemaining to not encrypt/decrypt the last
+    /// clocks if it is exactly the same as the number of remaining bytes.
+    /// if (remaining_bytes > encryption_block_bytes) {
+    ///   // Since this is the last blocks, this is effectively the same
+    ///   // condition as remaining_bytes != encryption_block_bytes.
+    ///   encrypt().
+    /// }
+    /// Use this mode for HLS SAMPLE-AES.
+    kSkipIfCryptByteBlockRemaining,
+  };
+
   /// @param crypt_byte_block indicates number of encrypted blocks (16-byte) in
   ///        pattern based encryption.
   /// @param skip_byte_block indicates number of unencrypted blocks (16-byte)
   ///        in pattern based encryption.
+  /// @param encryption_mode is used to determine the behavior for the last
+  ///        block.
   /// @param constant_iv_flag indicates whether a constant iv is used,
   ///        kUseConstantIv means that the same iv is used for all Crypt calls
   ///        until iv is changed via SetIv; otherwise, iv can be incremented
@@ -29,6 +50,7 @@ class AesPatternCryptor : public AesCryptor {
   ///        encryption/decryption.
   AesPatternCryptor(uint8_t crypt_byte_block,
                     uint8_t skip_byte_block,
+                    PatternEncryptionMode encryption_mode,
                     ConstantIvFlag constant_iv_flag,
                     scoped_ptr<AesCryptor> cryptor);
   ~AesPatternCryptor() override;
@@ -46,8 +68,11 @@ class AesPatternCryptor : public AesCryptor {
                      size_t* crypt_text_size) override;
   void SetIvInternal() override;
 
+  bool NeedEncrypt(size_t input_size, size_t target_data_size);
+
   const uint8_t crypt_byte_block_;
   const uint8_t skip_byte_block_;
+  const PatternEncryptionMode encryption_mode_;
   scoped_ptr<AesCryptor> cryptor_;
 
   DISALLOW_COPY_AND_ASSIGN(AesPatternCryptor);

packager/media/base/aes_pattern_cryptor_unittest.cc

@@ -43,6 +43,7 @@ class AesPatternCryptorTest : public ::testing::Test {
       : mock_cryptor_(new MockAesCryptor),
         pattern_cryptor_(kCryptByteBlock,
                          kSkipByteBlock,
+                         AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
                          AesCryptor::kDontUseConstantIv,
                          scoped_ptr<MockAesCryptor>(mock_cryptor_)) {}
 
@@ -141,9 +142,11 @@ INSTANTIATE_TEST_CASE_P(PatternTestCases,
 
 TEST(AesPatternCryptorConstIvTest, UseConstantIv) {
   MockAesCryptor* mock_cryptor = new MockAesCryptor;
-  AesPatternCryptor pattern_cryptor(kCryptByteBlock, kSkipByteBlock,
+  AesPatternCryptor pattern_cryptor(
-                                    AesPatternCryptor::kUseConstantIv,
+      kCryptByteBlock, kSkipByteBlock,
-                                    scoped_ptr<MockAesCryptor>(mock_cryptor));
+      AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
+      AesPatternCryptor::kUseConstantIv,
+      scoped_ptr<MockAesCryptor>(mock_cryptor));
 
   std::vector<uint8_t> iv(8, 'i');
   // SetIv will be called twice:
@@ -158,9 +161,11 @@ TEST(AesPatternCryptorConstIvTest, UseConstantIv) {
 
 TEST(AesPatternCryptorConstIvTest, DontUseConstantIv) {
   MockAesCryptor* mock_cryptor = new MockAesCryptor;
-  AesPatternCryptor pattern_cryptor(kCryptByteBlock, kSkipByteBlock,
+  AesPatternCryptor pattern_cryptor(
-                                    AesPatternCryptor::kDontUseConstantIv,
+      kCryptByteBlock, kSkipByteBlock,
-                                    scoped_ptr<MockAesCryptor>(mock_cryptor));
+      AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
+      AesPatternCryptor::kDontUseConstantIv,
+      scoped_ptr<MockAesCryptor>(mock_cryptor));
 
   std::vector<uint8_t> iv(8, 'i');
   // SetIv will be called only once by AesPatternCryptor::SetIv.
@@ -171,5 +176,48 @@ TEST(AesPatternCryptorConstIvTest, DontUseConstantIv) {
   ASSERT_TRUE(pattern_cryptor.Crypt("010203", &crypt_text));
 }
 
+TEST(SampleAesPatternCryptor, 16Bytes) {
+  MockAesCryptor* mock_cryptor = new MockAesCryptor();
+  EXPECT_CALL(*mock_cryptor, CryptInternal(_, _, _, _)).Times(0);
+
+  const uint8_t kSampleAesEncryptedBlock = 1;
+  const uint8_t kSampleAesClearBlock = 9;
+  AesPatternCryptor pattern_cryptor(
+      kSampleAesEncryptedBlock, kSampleAesClearBlock,
+      AesPatternCryptor::kSkipIfCryptByteBlockRemaining,
+      AesPatternCryptor::kUseConstantIv,
+      scoped_ptr<MockAesCryptor>(mock_cryptor));
+
+  std::vector<uint8_t> iv(8, 'i');
+  // SetIv will be called only once by AesPatternCryptor::SetIv.
+  EXPECT_TRUE(pattern_cryptor.SetIv(iv));
+
+  std::string crypt_text;
+  // Exactly 16 bytes, mock's Crypt should not be called.
+  ASSERT_TRUE(pattern_cryptor.Crypt("0123456789abcdef", &crypt_text));
+}
+
+TEST(SampleAesPatternCryptor, MoreThan16Bytes) {
+  MockAesCryptor* mock_cryptor = new MockAesCryptor();
+  EXPECT_CALL(*mock_cryptor, CryptInternal(_, 16u, _, _))
+      .WillOnce(Return(true));
+
+  const uint8_t kSampleAesEncryptedBlock = 1;
+  const uint8_t kSampleAesClearBlock = 9;
+  AesPatternCryptor pattern_cryptor(
+      kSampleAesEncryptedBlock, kSampleAesClearBlock,
+      AesPatternCryptor::kSkipIfCryptByteBlockRemaining,
+      AesPatternCryptor::kUseConstantIv,
+      scoped_ptr<MockAesCryptor>(mock_cryptor));
+
+  std::vector<uint8_t> iv(8, 'i');
+  // SetIv will be called only once by AesPatternCryptor::SetIv.
+  EXPECT_TRUE(pattern_cryptor.SetIv(iv));
+
+  std::string crypt_text;
+  // More than 16 bytes so mock's CryptInternal should be called.
+  ASSERT_TRUE(pattern_cryptor.Crypt("0123456789abcdef012", &crypt_text));
+}
+
 }  // namespace media
 }  // namespace edash_packager

packager/media/base/decryptor_source.cc

@@ -51,13 +51,17 @@ bool DecryptorSource::DecryptSampleBuffer(const DecryptConfig* decrypt_config,
       case FOURCC_cens:
         aes_decryptor.reset(new AesPatternCryptor(
             decrypt_config->crypt_byte_block(),
-            decrypt_config->skip_byte_block(), AesCryptor::kDontUseConstantIv,
+            decrypt_config->skip_byte_block(),
-            scoped_ptr<AesCryptor>(new AesCtrDecryptor)));
+            AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
+            AesCryptor::kDontUseConstantIv,
+            scoped_ptr<AesCryptor>(new AesCtrDecryptor())));
         break;
       case FOURCC_cbcs:
         aes_decryptor.reset(new AesPatternCryptor(
             decrypt_config->crypt_byte_block(),
-            decrypt_config->skip_byte_block(), AesCryptor::kUseConstantIv,
+            decrypt_config->skip_byte_block(),
+            AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
+            AesCryptor::kUseConstantIv,
             scoped_ptr<AesCryptor>(new AesCbcDecryptor(kNoPadding))));
         break;
       default:

packager/media/filters/nal_unit_to_byte_stream_converter.cc

@@ -27,10 +27,30 @@ const uint8_t kEmulationPreventionByte = 0x03;
 
 const uint8_t kAccessUnitDelimiterRbspAnyPrimaryPicType = 0xF0;
 
-// Inserts emulation byte where necessary.
+void AppendNalu(const Nalu& nalu,
-void EscapeRawByteSequencePayload(const uint8_t* input,
+                int nalu_length_size,
-                                  size_t input_size,
+                bool escape_data,
-                                  BufferWriter* output_writer) {
+                BufferWriter* buffer_writer) {
+  if (escape_data) {
+    EscapeNalByteSequence(nalu.data(), nalu.header_size() + nalu.payload_size(),
+                          buffer_writer);
+  } else {
+    buffer_writer->AppendArray(nalu.data(),
+                               nalu.header_size() + nalu.payload_size());
+  }
+}
+
+void AddAccessUnitDelimiter(BufferWriter* buffer_writer) {
+  buffer_writer->AppendInt(static_cast<uint8_t>(Nalu::H264_AUD));
+  // For now, primary_pic_type is 7 which is "anything".
+  buffer_writer->AppendInt(kAccessUnitDelimiterRbspAnyPrimaryPicType);
+}
+
+}  // namespace
+
+void EscapeNalByteSequence(const uint8_t* input,
+                           size_t input_size,
+                           BufferWriter* output_writer) {
   // Keep track of consecutive zeros that it has seen (not including the current
   // byte), so that the algorithm doesn't need to go back to check the same
   // bytes.
@@ -43,6 +63,7 @@ void EscapeRawByteSequencePayload(const uint8_t* input,
         // Must be escaped.
         output_writer->AppendInt(kEmulationPreventionByte);
       }
+
       output_writer->AppendInt(input[i]);
       // Note that input[i] can be 0.
       // 00 00 00 00 00 00 should become
@@ -64,27 +85,6 @@ void EscapeRawByteSequencePayload(const uint8_t* input,
   }
 }
 
-void AppendNalu(const Nalu& nalu,
-                int nalu_length_size,
-                bool escape_data,
-                BufferWriter* buffer_writer) {
-  if (escape_data) {
-    EscapeRawByteSequencePayload(
-        nalu.data(), nalu.header_size() + nalu.payload_size(), buffer_writer);
-  } else {
-    buffer_writer->AppendArray(nalu.data(),
-                               nalu.header_size() + nalu.payload_size());
-  }
-}
-
-void AddAccessUnitDelimiter(BufferWriter* buffer_writer) {
-  buffer_writer->AppendInt(static_cast<uint8_t>(Nalu::H264_AUD));
-  // For now, primary_pic_type is 7 which is "anything".
-  buffer_writer->AppendInt(kAccessUnitDelimiterRbspAnyPrimaryPicType);
-}
-
-}  // namespace
-
 NalUnitToByteStreamConverter::NalUnitToByteStreamConverter()
     : nalu_length_size_(0), escape_data_(false) {}
 NalUnitToByteStreamConverter::~NalUnitToByteStreamConverter() {}

packager/media/filters/nal_unit_to_byte_stream_converter.h

@@ -16,8 +16,20 @@
 namespace edash_packager {
 namespace media {
 
+class BufferWriter;
 class VideoStreamInfo;
 
+/// Inserts emulation prevention byte (0x03) where necessary.
+/// It is safe to call this again on the output, i.e. it is OK to "re-escape".
+/// This cannot do in-place escaping.
+/// @param input is the data to be escaped. This may not be the same (internal)
+///        buffer as @a output.
+/// @param input_size is the size of the input.
+/// @param output is the escaped data.
+void EscapeNalByteSequence(const uint8_t* input,
+                           size_t input_size,
+                           BufferWriter* output);
+
 // Methods are virtual for mocking.
 class NalUnitToByteStreamConverter {
  public:

packager/media/filters/nalu_reader.h

@@ -92,8 +92,12 @@ class Nalu {
                   const uint8_t* data,
                   uint64_t size) WARN_UNUSED_RESULT;
 
+  /// This is the pointer to the Nalu data, pointing to the header.
   const uint8_t* data() const { return data_; }
+
+  /// The size of the header, e.g. 1 for H.264.
   uint64_t header_size() const { return header_size_; }
+  /// Size of this Nalu minus header_size().
   uint64_t payload_size() const { return payload_size_; }
 
   // H.264 Specific:
@@ -103,6 +107,8 @@ class Nalu {
   int nuh_layer_id() const { return nuh_layer_id_; }
   int nuh_temporal_id() const { return nuh_temporal_id_; }
 
+  /// H264NaluType and H265NaluType enums may be used to compare against the
+  /// return value.
   int type() const { return type_; }
   bool is_video_slice() const { return is_video_slice_; }
   bool can_start_access_unit() const { return can_start_access_unit_; }

packager/media/formats/mp2t/pes_packet_generator.cc

@@ -6,10 +6,17 @@
 
 #include "packager/media/formats/mp2t/pes_packet_generator.h"
 
+#include <algorithm>
+#include <cstring>
+
+#include "packager/media/base/aes_encryptor.h"
+#include "packager/media/base/aes_pattern_cryptor.h"
 #include "packager/media/base/audio_stream_info.h"
+#include "packager/media/base/buffer_writer.h"
 #include "packager/media/base/media_sample.h"
 #include "packager/media/base/video_stream_info.h"
 #include "packager/media/filters/nal_unit_to_byte_stream_converter.h"
+#include "packager/media/filters/nalu_reader.h"
 #include "packager/media/formats/mp2t/pes_packet.h"
 #include "packager/media/formats/mp4/aac_audio_specific_config.h"
 
@@ -19,11 +26,68 @@ namespace mp2t {
 
 namespace {
 const bool kEscapeData = true;
-const uint8_t kVideoStreamId = 0xe0;
+const uint8_t kVideoStreamId = 0xE0;
-const uint8_t kAudioStreamId = 0xc0;
+const uint8_t kAudioStreamId = 0xC0;
+const double kTsTimescale = 90000.0;
+
+// |target_data| is input as well as output. On success |target_data| contains
+// the encrypted sample. The input data should be Nal unit byte stream.
+// This function constructs encrypted sample in |encrypted_sample_data| then
+// swap with target_data on success.
+bool EncryptH264Sample(AesCryptor* encryptor,
+                       std::vector<uint8_t>* target_data) {
+  BufferWriter encrypted_sample_data(target_data->size() * 1.5);
+
+  const int kLeadingClearBytesSize = 32;
+  // Any Nalu smaller than 48 bytes shall not be encrypted.
+  const uint64_t kSmallNalUnitSize = 48;
+
+  NaluReader nalu_reader(Nalu::kH264, 0, target_data->data(),
+                         target_data->size());
+  NaluReader::Result result;
+  Nalu nalu;
+  while ((result = nalu_reader.Advance(&nalu)) == NaluReader::Result::kOk) {
+    encrypted_sample_data.AppendInt(static_cast<uint32_t>(0x00000001));
+    const uint64_t nalu_total_size = nalu.header_size() + nalu.payload_size();
+    if (nalu.type() != Nalu::H264NaluType::H264_NonIDRSlice &&
+        nalu.type() != Nalu::H264NaluType::H264_IDRSlice) {
+      VLOG(3) << "Found Nalu type: " << nalu.type() << " skipping encryption.";
+      encrypted_sample_data.AppendArray(nalu.data(), nalu_total_size);
+      continue;
+    }
+
+    if (nalu_total_size <= kSmallNalUnitSize) {
+      encrypted_sample_data.AppendArray(nalu.data(), nalu_total_size);
+      continue;
+    }
+
+    const uint8_t* current = nalu.data() + kLeadingClearBytesSize;
+    const uint64_t bytes_remaining = nalu_total_size - kLeadingClearBytesSize;
+
+    if (!encryptor->Crypt(current, bytes_remaining,
+                          const_cast<uint8_t*>(current))) {
+      return false;
+    }
+    EscapeNalByteSequence(nalu.data(), nalu_total_size, &encrypted_sample_data);
+  }
+
+  encrypted_sample_data.SwapBuffer(target_data);
+  return true;
+}
+
+bool EncryptAacSample(AesCryptor* encryptor,
+                      std::vector<uint8_t>* target_data) {
+  const int kUnencryptedLeaderSize = 16;
+  if (target_data->size() <= kUnencryptedLeaderSize)
+    return true;
+  uint8_t* data_ptr = target_data->data() + kUnencryptedLeaderSize;
+  return encryptor->Crypt(
+      data_ptr, target_data->size() - kUnencryptedLeaderSize, data_ptr);
+}
 }  // namespace
 
-PesPacketGenerator::PesPacketGenerator() : pes_packets_deleter_(&pes_packets_) {}
+PesPacketGenerator::PesPacketGenerator()
+    : pes_packets_deleter_(&pes_packets_) {}
 PesPacketGenerator::~PesPacketGenerator() {}
 
 bool PesPacketGenerator::Initialize(const StreamInfo& stream_info) {
@@ -38,7 +102,7 @@ bool PesPacketGenerator::Initialize(const StreamInfo& stream_info) {
                        << " is not supported.";
       return false;
     }
-    timescale_scale_ = 90000.0 / video_stream_info.time_scale();
+    timescale_scale_ = kTsTimescale / video_stream_info.time_scale();
     converter_.reset(new NalUnitToByteStreamConverter());
     return converter_->Initialize(video_stream_info.extra_data().data(),
                                   video_stream_info.extra_data().size(),
@@ -51,7 +115,7 @@ bool PesPacketGenerator::Initialize(const StreamInfo& stream_info) {
                        << " is not supported yet.";
       return false;
     }
-    timescale_scale_ = 90000.0 / audio_stream_info.time_scale();
+    timescale_scale_ = kTsTimescale / audio_stream_info.time_scale();
     adts_converter_.reset(new mp4::AACAudioSpecificConfig());
     return adts_converter_->Parse(audio_stream_info.extra_data());
   }
@@ -68,12 +132,21 @@ bool PesPacketGenerator::PushSample(scoped_refptr<MediaSample> sample) {
   current_processing_pes_->set_dts(timescale_scale_ * sample->dts());
   if (stream_type_ == kStreamVideo) {
     DCHECK(converter_);
+    std::vector<uint8_t> byte_stream;
     if (!converter_->ConvertUnitToByteStream(
             sample->data(), sample->data_size(), sample->is_key_frame(),
-            current_processing_pes_->mutable_data())) {
+            &byte_stream)) {
       LOG(ERROR) << "Failed to convert sample to byte stream.";
       return false;
     }
+
+    if (encryptor_) {
+      if (!EncryptH264Sample(encryptor_.get(), &byte_stream)) {
+        LOG(ERROR) << "Failed to encrypt byte stream.";
+        return false;
+      }
+    }
+    current_processing_pes_->mutable_data()->swap(byte_stream);
     current_processing_pes_->set_stream_id(kVideoStreamId);
     pes_packets_.push_back(current_processing_pes_.release());
     return true;
@@ -83,6 +156,18 @@ bool PesPacketGenerator::PushSample(scoped_refptr<MediaSample> sample) {
 
   std::vector<uint8_t> aac_frame(sample->data(),
                                  sample->data() + sample->data_size());
+
+  if (encryptor_) {
+    if (!EncryptAacSample(encryptor_.get(), &aac_frame)) {
+      LOG(ERROR) << "Failed to encrypt ADTS AAC.";
+      return false;
+    }
+  }
+
+  // TODO(rkuroiwa): ConvertToADTS() makes another copy of aac_frame internally.
+  // Optimize copying in this function, possibly by adding a method on
+  // AACAudioSpecificConfig that takes {pointer, length} pair and returns a
+  // vector that has the ADTS header.
   if (!adts_converter_->ConvertToADTS(&aac_frame))
     return false;
 
@@ -94,6 +179,30 @@ bool PesPacketGenerator::PushSample(scoped_refptr<MediaSample> sample) {
   return true;
 }
 
+bool PesPacketGenerator::SetEncryptionKey(
+    scoped_ptr<EncryptionKey> encryption_key) {
+  if (stream_type_ == kStreamVideo) {
+    scoped_ptr<AesCbcEncryptor> cbc(
+        new AesCbcEncryptor(CbcPaddingScheme::kNoPadding));
+
+    const uint8_t kEncryptedBlocks = 1;
+    const uint8_t kClearBlocks = 9;
+    encryptor_.reset(new AesPatternCryptor(
+        kEncryptedBlocks, kClearBlocks,
+        AesPatternCryptor::kSkipIfCryptByteBlockRemaining,
+        AesCryptor::ConstantIvFlag::kUseConstantIv, cbc.Pass()));
+  } else if (stream_type_ == kStreamAudio) {
+    encryptor_.reset(
+        new AesCbcEncryptor(CbcPaddingScheme::kNoPadding,
+                            AesCryptor::ConstantIvFlag::kUseConstantIv));
+  } else {
+    LOG(ERROR) << "Cannot encrypt stream type: " << stream_type_;
+    return false;
+  }
+
+  return encryptor_->InitializeWithIv(encryption_key->key, encryption_key->iv);
+}
+
 size_t PesPacketGenerator::NumberOfReadyPesPackets() {
   return pes_packets_.size();
 }

packager/media/formats/mp2t/pes_packet_generator.h

@@ -11,6 +11,8 @@
 
 #include "packager/base/memory/scoped_ptr.h"
 #include "packager/base/stl_util.h"
+#include "packager/media/base/aes_cryptor.h"
+#include "packager/media/base/key_source.h"
 #include "packager/media/base/media_sample.h"
 #include "packager/media/base/stream_info.h"
 
@@ -48,6 +50,12 @@ class PesPacketGenerator {
   /// @return true on success, false otherwise.
   virtual bool PushSample(scoped_refptr<MediaSample> sample);
 
+  /// Sets the encryption key for encrypting samples.
+  /// @param encryption_key is the key that will be used to encrypt further
+  ///        samples.
+  /// @return true on success, false otherwise.
+  virtual bool SetEncryptionKey(scoped_ptr<EncryptionKey> encryption_key);
+
   /// @return The number of PES packets that are ready to be consumed.
   virtual size_t NumberOfReadyPesPackets();
 
@@ -80,6 +88,9 @@ class PesPacketGenerator {
   std::list<PesPacket*> pes_packets_;
   STLElementDeleter<decltype(pes_packets_)> pes_packets_deleter_;
 
+  // Current encryption key.
+  scoped_ptr<AesCryptor> encryptor_;
+
   DISALLOW_COPY_AND_ASSIGN(PesPacketGenerator);
 };
 

packager/media/formats/mp2t/pes_packet_generator_unittest.cc

@@ -133,6 +133,91 @@ class PesPacketGeneratorTest : public ::testing::Test {
     generator_.adts_converter_ = mock.Pass();
   }
 
+  void H264EncryptionTest(const uint8_t* input,
+                          size_t input_size,
+                          const uint8_t* expected_output,
+                          size_t expected_output_size) {
+    scoped_refptr<VideoStreamInfo> stream_info(
+        CreateVideoStreamInfo(kH264VideoCodec));
+    EXPECT_TRUE(generator_.Initialize(*stream_info));
+    EXPECT_EQ(0u, generator_.NumberOfReadyPesPackets());
+
+    scoped_refptr<MediaSample> sample =
+        MediaSample::CopyFrom(input, input_size, kIsKeyFrame);
+    const uint32_t kPts = 12345;
+    const uint32_t kDts = 12300;
+    sample->set_pts(kPts);
+    sample->set_dts(kDts);
+
+    scoped_ptr<MockNalUnitToByteStreamConverter> mock(
+        new MockNalUnitToByteStreamConverter());
+
+    // Returning only the input data so that it doesn't have all the unnecessary
+    // NALUs to test encryption.
+    std::vector<uint8_t> clear_data(input, input + input_size);
+    EXPECT_CALL(*mock, ConvertUnitToByteStream(_, input_size, kIsKeyFrame, _))
+        .WillOnce(DoAll(SetArgPointee<3>(clear_data), Return(true)));
+
+    UseMockNalUnitToByteStreamConverter(mock.Pass());
+
+    const std::vector<uint8_t> all_zero(16, 0);
+    scoped_ptr<EncryptionKey> encryption_key(new EncryptionKey());
+    encryption_key->key = all_zero;
+    encryption_key->iv = all_zero;
+    EXPECT_TRUE(generator_.SetEncryptionKey(encryption_key.Pass()));
+
+    EXPECT_TRUE(generator_.PushSample(sample));
+    EXPECT_EQ(1u, generator_.NumberOfReadyPesPackets());
+    scoped_ptr<PesPacket> pes_packet = generator_.GetNextPesPacket();
+    ASSERT_TRUE(pes_packet);
+
+    std::vector<uint8_t> expected(expected_output,
+                                  expected_output + expected_output_size);
+    ASSERT_EQ(expected.size(), pes_packet->data().size());
+    for (size_t i = 0; i < expected.size(); ++i) {
+      EXPECT_EQ(expected[i], pes_packet->data()[i]) << " mismatch at " << i;
+    }
+    //EXPECT_EQ(expected, pes_packet->data());
+  }
+
+  // The input data should be the size of an aac frame, i.e. should not be the
+  // size of an ADTS frame.
+  void AacEncryptionTest(const uint8_t* input,
+                         size_t input_size,
+                         const uint8_t* expected_output,
+                         size_t expected_output_size) {
+    scoped_refptr<AudioStreamInfo> stream_info(
+        CreateAudioStreamInfo(kAacAudioCodec));
+    EXPECT_TRUE(generator_.Initialize(*stream_info));
+    EXPECT_EQ(0u, generator_.NumberOfReadyPesPackets());
+
+    // For aac, the input from MediaSample is used. Then ADTS header is added,
+    // so EXPECT_CALL does not return the |input| data.
+    scoped_refptr<MediaSample> sample = MediaSample::CopyFrom(
+        input, input_size, kIsKeyFrame);
+
+    scoped_ptr<MockAACAudioSpecificConfig> mock(
+        new MockAACAudioSpecificConfig());
+    EXPECT_CALL(*mock, ConvertToADTS(_)).WillOnce(Return(true));
+
+    UseMockAACAudioSpecificConfig(mock.Pass());
+
+    const std::vector<uint8_t> all_zero(16, 0);
+    scoped_ptr<EncryptionKey> encryption_key(new EncryptionKey());
+    encryption_key->key = all_zero;
+    encryption_key->iv = all_zero;
+    EXPECT_TRUE(generator_.SetEncryptionKey(encryption_key.Pass()));
+
+    EXPECT_TRUE(generator_.PushSample(sample));
+    EXPECT_EQ(1u, generator_.NumberOfReadyPesPackets());
+    scoped_ptr<PesPacket> pes_packet = generator_.GetNextPesPacket();
+    ASSERT_TRUE(pes_packet);
+
+    std::vector<uint8_t> expected(expected_output,
+                                  expected_output + expected_output_size);
+    EXPECT_EQ(expected, pes_packet->data());
+  }
+
   PesPacketGenerator generator_;
 };
 
@@ -318,6 +403,323 @@ TEST_F(PesPacketGeneratorTest, TimeStampScaling) {
   EXPECT_TRUE(generator_.Flush());
 }
 
+// The nalu is too small for it to be encrypted. Verify it is not modified.
+TEST_F(PesPacketGeneratorTest, H264SampleEncryptionSmallNalu) {
+  const uint8_t kNaluData[] = {
+      0x00, 0x00, 0x00, 0x01, 0x61, 0xbb, 0xcc, 0xdd,
+  };
+
+  ASSERT_NO_FATAL_FAILURE(H264EncryptionTest(kNaluData, arraysize(kNaluData),
+                                             kNaluData, arraysize(kNaluData)));
+}
+
+// Verify that sample encryption works.
+TEST_F(PesPacketGeneratorTest, H264SampleEncryption) {
+  // Use the following command to encrypt data.
+  // openssl aes-128-cbc -nopad -e -in input -K
+  // "00000000000000000000000000000000" -iv "00000000000000000000000000000000" >
+  // enc
+  const uint8_t kNaluData[] = {
+      0x00, 0x00, 0x00, 0x01,  // Start code.
+      0x61,                    // nalu type 1; this type should get encrypted.
+      // Bogus data but should not be encrypted.
+      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
+      0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+      0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E,
+
+      // Next 16 bytes should be encrypted.
+      0x1F, 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A,
+      0x2B, 0x2C, 0x2D, 0x2E,
+
+      // Next 144 bytes should be in the clear.
+      0x2F, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A,
+      0x3B, 0x3C, 0x3D, 0x3E, 0x3F, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46,
+      0x47, 0x48, 0x49, 0x4A, 0x4B, 0x4C, 0x4D, 0x4E, 0x4F, 0x50, 0x51, 0x52,
+      0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x5B, 0x5C, 0x5D, 0x5E,
+      0x5F, 0x60, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A,
+      0x6B, 0x6C, 0x6D, 0x6E, 0x6F, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
+      0x77, 0x78, 0x79, 0x7A, 0x7B, 0x7C, 0x7D, 0x7E, 0x7F, 0x80, 0x81, 0x82,
+      0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x8B, 0x8C, 0x8D, 0x8E,
+      0x8F, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A,
+      0x9B, 0x9C, 0x9D, 0x9E, 0x9F, 0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6,
+      0xA7, 0xA8, 0xA9, 0xAA, 0xAB, 0xAC, 0xAD, 0xAE, 0xAF, 0xB0, 0xB1, 0xB2,
+      0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xBB, 0xBC, 0xBD, 0xBE,
+
+      // Next 16 bytes should be encrypted.
+      0xBF, 0xC0, 0xC1, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA,
+      0xCB, 0xCC, 0xCD, 0xCE,
+
+      // This last bytes should not be encrypted.
+      0xCF,
+  };
+
+  const uint8_t kEncryptedNaluData[] = {
+      0x00, 0x00, 0x00, 0x01,  // Start code.
+      0x61,                    // nalu type 1; should get encrypted.
+      // Bogus data but should sample encrypted.
+      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
+      0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+      0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E,
+
+      // Encrypted 16 bytes.
+      0x93, 0x3A, 0x2C, 0x38, 0x86, 0x4B, 0x64, 0xE2, 0x62, 0x7E, 0xCC, 0x75,
+      0x71, 0xFB, 0x60, 0x7C,
+
+      // Next 144 bytes should be in the clear.
+      0x2F, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A,
+      0x3B, 0x3C, 0x3D, 0x3E, 0x3F, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46,
+      0x47, 0x48, 0x49, 0x4A, 0x4B, 0x4C, 0x4D, 0x4E, 0x4F, 0x50, 0x51, 0x52,
+      0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x5B, 0x5C, 0x5D, 0x5E,
+      0x5F, 0x60, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A,
+      0x6B, 0x6C, 0x6D, 0x6E, 0x6F, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
+      0x77, 0x78, 0x79, 0x7A, 0x7B, 0x7C, 0x7D, 0x7E, 0x7F, 0x80, 0x81, 0x82,
+      0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x8B, 0x8C, 0x8D, 0x8E,
+      0x8F, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A,
+      0x9B, 0x9C, 0x9D, 0x9E, 0x9F, 0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6,
+      0xA7, 0xA8, 0xA9, 0xAA, 0xAB, 0xAC, 0xAD, 0xAE, 0xAF, 0xB0, 0xB1, 0xB2,
+      0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xBB, 0xBC, 0xBD, 0xBE,
+
+      // Encrypted 16 bytes.
+      0xB7, 0x1C, 0x64, 0xAE, 0x90, 0xA4, 0x35, 0x88, 0x4F, 0xD1, 0x30, 0xC2,
+      0x06, 0x2E, 0xF8, 0xA5,
+
+      // This last bytes should not be encrypted.
+      0xCF,
+  };
+  ASSERT_NO_FATAL_FAILURE(H264EncryptionTest(kNaluData, arraysize(kNaluData),
+                                             kEncryptedNaluData,
+                                             arraysize(kEncryptedNaluData)));
+}
+
+// If any block is encrypted, then the whole nal unit must be re-escaped.
+TEST_F(PesPacketGeneratorTest, H264SampleEncryptionVerifyReescape) {
+  const uint8_t kNaluData[] = {
+      0x00, 0x00, 0x00, 0x01,  // Start code.
+      0x61,                    // nalu type 1; this type should get encrypted.
+      // Bogus data but should not be encrypted.
+      // But 0x00 0x00 0x03 should be re-escaped.
+      0x00, 0x00, 0x03, 0x02, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
+      0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+      0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E,
+
+      // Next 16 bytes should be encrypted.
+      // Note that there is 0x00 0x00 0x03 sequence that will be reescaped.
+      0x1F, 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A,
+      0x2B, 0x2C, 0x2D, 0x2E,
+
+      // Next 144 bytes should be in the clear.
+      0x2F, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A,
+      0x3B, 0x3C, 0x3D, 0x3E, 0x3F, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46,
+      0x47, 0x48, 0x49, 0x4A, 0x4B, 0x4C, 0x4D, 0x4E, 0x4F, 0x50, 0x51, 0x52,
+      0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x5B, 0x5C, 0x5D, 0x5E,
+      0x5F, 0x60, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A,
+      0x6B, 0x6C, 0x6D, 0x6E, 0x6F, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
+      0x77, 0x78, 0x79, 0x7A, 0x7B, 0x7C, 0x7D, 0x7E, 0x7F, 0x80, 0x81, 0x82,
+      0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x8B, 0x8C, 0x8D, 0x8E,
+      0x8F, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A,
+      // Still part of clear data, but this line includes 0x00 0x00 0x03
+      // which should be re-escaped.
+      0x9B, 0x9C, 0x9D, 0x00, 0x00, 0x03, 0x01, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6,
+      0xA7, 0xA8, 0xA9, 0xAA, 0xAB, 0xAC, 0xAD, 0xAE, 0xAF, 0xB0, 0xB1, 0xB2,
+      0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xBB, 0xBC, 0xBD, 0xBE,
+
+      // Next 16 bytes should be encrypted.
+      0xBF, 0xC0, 0xC1, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA,
+      0xCB, 0xCC, 0xCD, 0xCE,
+
+      // This last bytes should not be encrypted.
+      0xCF,
+  };
+
+  const uint8_t kEncryptedNaluData[] = {
+      0x00, 0x00, 0x00, 0x01,  // Start code.
+      0x61,                    // nalu type 1; should get encrypted.
+      // Bogus data but should not be encrypted.
+      0x00, 0x00, 0x03, 0x03, 0x02, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A,
+      0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16,
+      0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E,
+
+      // Encrypted 16 bytes.
+      0x93, 0x3A, 0x2C, 0x38, 0x86, 0x4B, 0x64, 0xE2, 0x62, 0x7E, 0xCC, 0x75,
+      0x71, 0xFB, 0x60, 0x7C,
+
+      // Next 144 bytes should be in the clear.
+      0x2F, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A,
+      0x3B, 0x3C, 0x3D, 0x3E, 0x3F, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46,
+      0x47, 0x48, 0x49, 0x4A, 0x4B, 0x4C, 0x4D, 0x4E, 0x4F, 0x50, 0x51, 0x52,
+      0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x5B, 0x5C, 0x5D, 0x5E,
+      0x5F, 0x60, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A,
+      0x6B, 0x6C, 0x6D, 0x6E, 0x6F, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
+      0x77, 0x78, 0x79, 0x7A, 0x7B, 0x7C, 0x7D, 0x7E, 0x7F, 0x80, 0x81, 0x82,
+      0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x8B, 0x8C, 0x8D, 0x8E,
+      0x8F, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A,
+      // Extra 0x03 is added.
+      0x9B, 0x9C, 0x9D, 0x00, 0x00, 0x03, 0x03, 0x01, 0xA2, 0xA3, 0xA4, 0xA5,
+      0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xAB, 0xAC, 0xAD, 0xAE, 0xAF, 0xB0, 0xB1,
+      0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xBB, 0xBC, 0xBD,
+      0xBE,
+
+      // Encrypted 16 bytes.
+      0xB7, 0x1C, 0x64, 0xAE, 0x90, 0xA4, 0x35, 0x88, 0x4F, 0xD1, 0x30, 0xC2,
+      0x06, 0x2E, 0xF8, 0xA5,
+
+      // This last bytes should not be encrypted.
+      0xCF,
+  };
+  ASSERT_NO_FATAL_FAILURE(H264EncryptionTest(kNaluData, arraysize(kNaluData),
+                                             kEncryptedNaluData,
+                                             arraysize(kEncryptedNaluData)));
+}
+
+// Verify that if the last there are only 16 bytes left, then it doesn't get
+// encrypted.
+TEST_F(PesPacketGeneratorTest, H264SampleEncryptionLast16ByteNotEncrypted) {
+  const uint8_t kNaluData[] = {
+      0x00, 0x00, 0x00, 0x01,  // Start code.
+      0x61,                    // nalu type 1; should get encrypted.
+      // Bogus data but should not be encrypted.
+      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
+      0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+      0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E,
+
+      // Next 16 bytes should be encrypted.
+      0x1F, 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A,
+      0x2B, 0x2C, 0x2D, 0x2E,
+
+      // Next 144 bytes should be in the clear.
+      0x2F, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A,
+      0x3B, 0x3C, 0x3D, 0x3E, 0x3F, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46,
+      0x47, 0x48, 0x49, 0x4A, 0x4B, 0x4C, 0x4D, 0x4E, 0x4F, 0x50, 0x51, 0x52,
+      0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x5B, 0x5C, 0x5D, 0x5E,
+      0x5F, 0x60, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A,
+      0x6B, 0x6C, 0x6D, 0x6E, 0x6F, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
+      0x77, 0x78, 0x79, 0x7A, 0x7B, 0x7C, 0x7D, 0x7E, 0x7F, 0x80, 0x81, 0x82,
+      0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x8B, 0x8C, 0x8D, 0x8E,
+      0x8F, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A,
+      0x9B, 0x9C, 0x9D, 0x9E, 0x9F, 0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6,
+      0xA7, 0xA8, 0xA9, 0xAA, 0xAB, 0xAC, 0xAD, 0xAE, 0xAF, 0xB0, 0xB1, 0xB2,
+      0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xBB, 0xBC, 0xBD, 0xBE,
+
+      // These 16 bytes should not be encrypted.
+      0xBF, 0xC0, 0xC1, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA,
+      0xCB, 0xCC, 0xCD, 0xCE,
+  };
+
+  const uint8_t kEncryptedNaluData[] = {
+      0x00, 0x00, 0x00, 0x01,  // Start code.
+      0x61,                    // nalu type 1; should get encrypted.
+      // Bogus data but should not be encrypted.
+      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
+      0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+      0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E,
+
+      // Encrypted 16 bytes.
+      0x93, 0x3A, 0x2C, 0x38, 0x86, 0x4B, 0x64, 0xE2, 0x62, 0x7E, 0xCC, 0x75,
+      0x71, 0xFB, 0x60, 0x7C,
+
+      // Next 144 bytes should be in the clear.
+      0x2F, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A,
+      0x3B, 0x3C, 0x3D, 0x3E, 0x3F, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46,
+      0x47, 0x48, 0x49, 0x4A, 0x4B, 0x4C, 0x4D, 0x4E, 0x4F, 0x50, 0x51, 0x52,
+      0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x5B, 0x5C, 0x5D, 0x5E,
+      0x5F, 0x60, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A,
+      0x6B, 0x6C, 0x6D, 0x6E, 0x6F, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
+      0x77, 0x78, 0x79, 0x7A, 0x7B, 0x7C, 0x7D, 0x7E, 0x7F, 0x80, 0x81, 0x82,
+      0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x8B, 0x8C, 0x8D, 0x8E,
+      0x8F, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A,
+      0x9B, 0x9C, 0x9D, 0x9E, 0x9F, 0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6,
+      0xA7, 0xA8, 0xA9, 0xAA, 0xAB, 0xAC, 0xAD, 0xAE, 0xAF, 0xB0, 0xB1, 0xB2,
+      0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xBB, 0xBC, 0xBD, 0xBE,
+
+      // These 16 bytes should not be encrypted.
+      0xBF, 0xC0, 0xC1, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA,
+      0xCB, 0xCC, 0xCD, 0xCE,
+  };
+  ASSERT_NO_FATAL_FAILURE(H264EncryptionTest(kNaluData, arraysize(kNaluData),
+                                             kEncryptedNaluData,
+                                             arraysize(kEncryptedNaluData)));
+}
+
+// The sample is too small and it doesn't need to be encrypted.
+TEST_F(PesPacketGeneratorTest, AacSampleEncryptionSmallSample) {
+  const uint8_t kClearData[] = {
+      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A,
+      0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+      0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E,
+  };
+
+  ASSERT_NO_FATAL_FAILURE(AacEncryptionTest(kClearData, arraysize(kClearData),
+                                            kClearData, arraysize(kClearData)));
+}
+
+// Verify that AAC can be encrypted.
+TEST_F(PesPacketGeneratorTest, AacSampleEncryption) {
+  // The data is long enough so that 2 blocks (32 bytes) are encrypted.
+  const uint8_t kClearData[] = {
+      // First 16 bytes are always clear.
+      0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12,
+      0x13, 0x14, 0x15, 0x16,
+
+      // Next 32 bytes (2 blocks) are encrypted.
+      0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F, 0x20, 0x21, 0x22,
+      0x23, 0x24, 0x25, 0x26,
+      0x27, 0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F, 0x30, 0x31, 0x32,
+      0x33, 0x34, 0x35, 0x36,
+
+      // The last 2 bytes are in the clear.
+      0x37, 0x38,
+  };
+
+  const uint8_t kExpectedOutputData[] = {
+      // First 16 bytes are always clear.
+      0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12,
+      0x13, 0x14, 0x15, 0x16,
+
+      // Encrypted bytes.
+      0xE3, 0x42, 0x9B, 0x27, 0x33, 0x67, 0x68, 0x08, 0xA5, 0xB3, 0x3E, 0xB1,
+      0xEE, 0xFC, 0x9E, 0x0A, 0x8E, 0x0C, 0x73, 0xC5, 0x57, 0xEE, 0x58, 0xC7,
+      0x48, 0x74, 0x2A, 0x12, 0x38, 0x4F, 0x4E, 0xAC,
+
+      // The last 2 bytes are in the clear.
+      0x37, 0x38,
+  };
+
+  ASSERT_NO_FATAL_FAILURE(AacEncryptionTest(kClearData, arraysize(kClearData),
+                                            kExpectedOutputData,
+                                            arraysize(kExpectedOutputData)));
+}
+
+// Verify that all the bytes after the leading few bytes are encrypted.
+// Note that this is different from h264 encryption where it doesn't encrypt the
+// last 16.
+TEST_F(PesPacketGeneratorTest, AacSampleEncryptionLastBytesAreEncrypted) {
+  const uint8_t kClearData[] = {
+      // First 16 bytes are always clear.
+      0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12,
+      0x13, 0x14, 0x15, 0x16,
+
+      // Next 32 bytes (2 blocks) are encrypted.
+      0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F, 0x20, 0x21, 0x22,
+      0x23, 0x24, 0x25, 0x26,
+      0x27, 0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F, 0x30, 0x31, 0x32,
+      0x33, 0x34, 0x35, 0x36,
+  };
+
+  const uint8_t kExpectedOutputData[] = {
+      // First 16 bytes are always clear.
+      0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12,
+      0x13, 0x14, 0x15, 0x16,
+
+      // Encrypted bytes.
+      0xE3, 0x42, 0x9B, 0x27, 0x33, 0x67, 0x68, 0x08, 0xA5, 0xB3, 0x3E, 0xB1,
+      0xEE, 0xFC, 0x9E, 0x0A, 0x8E, 0x0C, 0x73, 0xC5, 0x57, 0xEE, 0x58, 0xC7,
+      0x48, 0x74, 0x2A, 0x12, 0x38, 0x4F, 0x4E, 0xAC,
+  };
+  ASSERT_NO_FATAL_FAILURE(AacEncryptionTest(kClearData, arraysize(kClearData),
+                                            kExpectedOutputData,
+                                            arraysize(kExpectedOutputData)));
+}
+
 }  // namespace mp2t
 }  // namespace media
 }  // namespace edash_packager

packager/media/formats/mp4/encrypting_fragmenter.cc

@@ -203,12 +203,16 @@ Status EncryptingFragmenter::CreateEncryptor() {
       break;
     case FOURCC_cens:
       encryptor.reset(new AesPatternCryptor(
-          crypt_byte_block(), skip_byte_block(), AesCryptor::kDontUseConstantIv,
+          crypt_byte_block(), skip_byte_block(),
-          scoped_ptr<AesCryptor>(new AesCtrEncryptor)));
+          AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
+          AesCryptor::kDontUseConstantIv,
+          scoped_ptr<AesCryptor>(new AesCtrEncryptor())));
       break;
     case FOURCC_cbcs:
       encryptor.reset(new AesPatternCryptor(
-          crypt_byte_block(), skip_byte_block(), AesCryptor::kUseConstantIv,
+          crypt_byte_block(), skip_byte_block(),
+          AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
+          AesCryptor::kUseConstantIv,
           scoped_ptr<AesCryptor>(new AesCbcEncryptor(kNoPadding))));
       break;
     default: