From 4b3c02267a7c52f95997b24c69ca4df0f130daa2 Mon Sep 17 00:00:00 2001
From: Fred Gylys-Colwell <fredgc@google.com>
Date: Wed, 1 Jul 2015 13:21:21 -0700
Subject: [PATCH] Add Short Key ID Test

Merge from widevine repo of http://go/wvgerrit/14970

Even if devices cannot handle key ids with different lengths in the
same license, they should still handle keys with a shorter key id.

This is a partial fix for:
bug: 21935358

Change-Id: Ibc84f0b5d7d9bc5d24a2081f0581a2b256e51f44
---
 .../oemcrypto/test/oemcrypto_test.cpp         | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp b/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
index 6049fdd8..d909b157 100644
--- a/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
+++ b/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
@@ -4725,6 +4725,39 @@ TEST_F(GenericCryptoTest, KeyDurationVerify) {
   ASSERT_EQ(OEMCrypto_ERROR_KEY_EXPIRED, sts);
 }
 
+TEST_F(GenericCryptoTest, UniformShortKeyId) {
+  OEMCryptoResult sts;
+  Session s;
+  s.open();
+  s.GenerateTestSessionKeys();
+  MakeFourKeys(&s);
+  s.SetKeyId(0, "123456789012"); // 12 bytes.
+  s.SetKeyId(1, "aaaaaaaaaaaa");
+  s.SetKeyId(2, "bbbbbbbbbbbb");
+  s.SetKeyId(3, "cccccccccccc");
+  s.EncryptAndSign();
+  s.LoadTestKeys();
+  unsigned int key_index = 1;
+  vector<uint8_t> encrypted;
+
+  // To make sure OEMCrypto is not expecting the key_id to be zero padded, we
+  // will create a buffer that is padded with 'Z'.
+  vector<uint8_t> key_id_buffer(s.license().keys[key_index].key_id_length + 5,
+                                'Z');  // Fill a bigger buffer with letter 'Z'.
+  memcpy(key_id_buffer.data(), s.license().keys[key_index].key_id,
+         s.license().keys[key_index].key_id_length);
+  EncryptBuffer(&s, key_index, clear_buffer_, &encrypted);
+  sts = OEMCrypto_SelectKey(s.session_id(), key_id_buffer.data(),
+                            s.license().keys[key_index].key_id_length);
+  ASSERT_EQ(OEMCrypto_SUCCESS, sts);
+  vector<uint8_t> resultant(encrypted.size());
+  sts = OEMCrypto_Generic_Decrypt(s.session_id(), &encrypted[0], encrypted.size(),
+                                  iv_, OEMCrypto_AES_CBC_128_NO_PADDING,
+                                  &resultant[0]);
+  ASSERT_EQ(OEMCrypto_SUCCESS, sts);
+  ASSERT_EQ(clear_buffer_, resultant);
+}
+
 TEST_F(GenericCryptoTest, DISABLED_ShortKeyId) {
   OEMCryptoResult sts;
   Session s;