allocate proper sized buffer for aes encryption

shaka-project · Jul 12, 2023 · 9c60090 · 9c60090
1 parent 7598735
commit 9c60090
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 9 deletions.
diff --git a/packager/media/base/aes_cryptor.h b/packager/media/base/aes_cryptor.h
@@ -41,6 +41,10 @@ class AesCryptor {
   virtual bool InitializeWithIv(const std::vector<uint8_t>& key,
                                 const std::vector<uint8_t>& iv) = 0;
 
+  virtual size_t RequiredCiphertextSize(size_t plaintext_size) {
+    return plaintext_size;
+  }
+
   /// @name Various forms of crypt (Encrypt/Decrypt) calls.
   /// It is an Encrypt function for encryptor and a Decrypt function for
   /// decryptor. The text and crypt_text pointers can be the same address for

diff --git a/packager/media/base/aes_encryptor.cc b/packager/media/base/aes_encryptor.cc
@@ -125,15 +125,20 @@ bool AesCbcEncryptor::InitializeWithIv(const std::vector<uint8_t>& key,
   return SetIv(iv);
 }
 
+size_t AesCbcEncryptor::RequiredCiphertextSize(size_t plaintext_size) {
+  // mbedtls requires a buffer large enough for one extra block.
+  return plaintext_size + NumPaddingBytes(plaintext_size) + AES_BLOCK_SIZE;
+}
+
 bool AesCbcEncryptor::CryptInternal(const uint8_t* plaintext,
                                     size_t plaintext_size,
                                     uint8_t* ciphertext,
                                     size_t* ciphertext_size) {
   const size_t residual_block_size = plaintext_size % AES_BLOCK_SIZE;
   const size_t num_padding_bytes = NumPaddingBytes(plaintext_size);
-  // mbedtls requires a buffer large enough for one extra block.
   const size_t required_ciphertext_size =
-      plaintext_size + num_padding_bytes + AES_BLOCK_SIZE;
+      RequiredCiphertextSize(plaintext_size);
+
   if (*ciphertext_size < required_ciphertext_size) {
     LOG(ERROR) << "Expecting output size of at least "
                << required_ciphertext_size << " bytes.";

diff --git a/packager/media/base/aes_encryptor.h b/packager/media/base/aes_encryptor.h
@@ -84,6 +84,8 @@ class AesCbcEncryptor : public AesCryptor {
   bool InitializeWithIv(const std::vector<uint8_t>& key,
                         const std::vector<uint8_t>& iv) override;
 
+  size_t RequiredCiphertextSize(size_t plaintext_size) override;
+
  private:
   bool CryptInternal(const uint8_t* plaintext,
                      size_t plaintext_size,

diff --git a/packager/media/crypto/encryption_handler.cc b/packager/media/crypto/encryption_handler.cc
@@ -299,8 +299,11 @@ Status EncryptionHandler::ProcessMediaSample(
     return DispatchMediaSample(kStreamIndex, std::move(clear_sample));
   }
 
-  std::shared_ptr<uint8_t> cipher_sample_data(
-      new uint8_t[clear_sample->data_size()], std::default_delete<uint8_t[]>());
+  size_t ciphertext_size =
+      encryptor_->RequiredCiphertextSize(clear_sample->data_size());
+
+  std::shared_ptr<uint8_t> cipher_sample_data(new uint8_t[ciphertext_size],
+                                              std::default_delete<uint8_t[]>());
 
   const uint8_t* source = clear_sample->data();
   uint8_t* dest = cipher_sample_data.get();
@@ -314,15 +317,15 @@ Status EncryptionHandler::ProcessMediaSample(
         total_size += subsample.clear_bytes;
       }
       if (subsample.cipher_bytes > 0) {
-        EncryptBytes(source, subsample.cipher_bytes, dest);
+        EncryptBytes(source, subsample.clear_bytes, dest, ciphertext_size);
         source += subsample.cipher_bytes;
         dest += subsample.cipher_bytes;
         total_size += subsample.cipher_bytes;
       }
     }
     DCHECK_EQ(total_size, clear_sample->data_size());
   } else {
-    EncryptBytes(source, clear_sample->data_size(), dest);
+    EncryptBytes(source, clear_sample->data_size(), dest, ciphertext_size);
   }
 
   std::shared_ptr<MediaSample> cipher_sample(clear_sample->Clone());
@@ -386,11 +389,12 @@ bool EncryptionHandler::CreateEncryptor(const EncryptionKey& encryption_key) {
 
 void EncryptionHandler::EncryptBytes(const uint8_t* source,
                                      size_t source_size,
-                                     uint8_t* dest) {
+                                     uint8_t* dest,
+                                     size_t dest_size) {
   DCHECK(source);
   DCHECK(dest);
   DCHECK(encryptor_);
-  CHECK(encryptor_->Crypt(source, source_size, dest));
+  CHECK(encryptor_->Crypt(source, source_size, dest, &dest_size));
 }
 
 void EncryptionHandler::InjectSubsampleGeneratorForTesting(

diff --git a/packager/media/crypto/encryption_handler.h b/packager/media/crypto/encryption_handler.h
@@ -53,7 +53,10 @@ class EncryptionHandler : public MediaHandler {
                                  uint8_t* dest);
   // Encrypt an array with size |source_size|. |dest| should have at
   // least |source_size| bytes.
-  void EncryptBytes(const uint8_t* source, size_t source_size, uint8_t* dest);
+  void EncryptBytes(const uint8_t* source,
+                    size_t source_size,
+                    uint8_t* dest,
+                    size_t dest_size);
 
   // An E-AC3 frame comprises of one or more syncframes. This function extracts
   // the syncframe sizes from the source bytes.