From e288d10026305bcc1782946b3991bb55d67b5355 Mon Sep 17 00:00:00 2001
From: Fred Gylys-Colwell <fredgc@google.com>
Date: Tue, 18 Sep 2018 16:48:13 -0700
Subject: [PATCH] Unit test with small buffer

Merge from Widevine repo of http://go/wvgerrit/57640

Add a unit test for OEMCrypto to handle a small buffer.

Test: unit tests
bug: 78233951
Change-Id: I5efe088705e2d248ab9ea45d8576daf69ad8bcdb
---
 .../oemcrypto/test/oemcrypto_test.cpp         | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp b/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
index 0f77d08a..2322e565 100644
--- a/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
+++ b/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
@@ -1862,9 +1862,9 @@ class OEMCryptoSessionTestsPartialBlockTests
     : public OEMCryptoSessionTestsDecryptTests {};
 
 TEST_P(OEMCryptoSessionTestsDecryptTests, SingleLargeSubsample) {
-  // This subsample size should be larger a few encrypt/skip patterns.  Most
+  // This subsample size is larger than a few encrypt/skip patterns.  Most
   // test cases use a pattern length of 160, so we'll run through at least two
-  // full patterns.
+  // full patterns if we have more than 320 -- round up to 400.
   subsample_size_.push_back(SampleSize(0, 400));
   FindTotalSize();
   vector<uint8_t> unencryptedData(total_size_);
@@ -2043,6 +2043,21 @@ TEST_P(OEMCryptoSessionTestsDecryptTests, DecryptLargeBuffer) {
   TestDecryptCENC(key, encryptionIv, encryptedData, unencryptedData);
 }
 
+TEST_P(OEMCryptoSessionTestsDecryptTests, DecryptSmallBuffer) {
+  // There are probably no frames this small, but we should handle them anyway.
+  subsample_size_.push_back(SampleSize(5, 5));
+  FindTotalSize();
+  vector<uint8_t> unencryptedData(total_size_);
+  vector<uint8_t> encryptedData(total_size_);
+  vector<uint8_t> encryptionIv(AES_BLOCK_SIZE);
+  vector<uint8_t> key(AES_BLOCK_SIZE);
+  EXPECT_EQ(1, GetRandBytes(&encryptionIv[0], AES_BLOCK_SIZE));
+  EXPECT_EQ(1, GetRandBytes(&key[0], AES_BLOCK_SIZE));
+  for (size_t i = 0; i < total_size_; i++) unencryptedData[i] = i % 256;
+  EncryptData(key, encryptionIv, unencryptedData, &encryptedData);
+  TestDecryptCENC(key, encryptionIv, encryptedData, unencryptedData);
+}
+
 TEST_P(OEMCryptoSessionTestsDecryptTests, DecryptUnencrypted) {
   subsample_size_.push_back(SampleSize(256, 0));
   FindTotalSize();