// Copyright 2021 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
//
//  Reference implementation utilities of OEMCrypto APIs
//
#include <gtest/gtest.h>

#include "cmac.h"

namespace wvoec {
namespace util {
namespace {
// Test vectors are from NIST Special Publication 800-38B for
// CMAC-AES-128 and CMAC-AES-256.  The same data chunks are used
// for both tests.
const uint8_t kDataChunk1[] = {0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96,
                               0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a};
const uint8_t kDataChunk2[] = {0xae, 0x2d, 0x8a, 0x57, 0x1e, 0x03, 0xac, 0x9c,
                               0x9e, 0xb7, 0x6f, 0xac, 0x45, 0xaf, 0x8e, 0x51,
                               0x30, 0xc8, 0x1c, 0x46, 0xa3, 0x5c, 0xe4, 0x11};
const uint8_t kDataChunk3[] = {0xe5, 0xfb, 0xc1, 0x19, 0x1a, 0x0a, 0x52, 0xef,
                               0xf6, 0x9f, 0x24, 0x45, 0xdf, 0x4f, 0x9b, 0x17,
                               0xad, 0x2b, 0x41, 0x7b, 0xe6, 0x6c, 0x37, 0x10};

const uint8_t kKey128[16] = {0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
                             0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c};
const uint8_t kKey256[32] = {0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe,
                             0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d, 0x77, 0x81,
                             0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7,
                             0x2d, 0x98, 0x10, 0xa3, 0x09, 0x14, 0xdf, 0xf4};
}  // namespace

TEST(OEMCryptoCmacTest, BadInit) {
  // If the parameters to the Create() function are invalid, the
  // Cmac instance is not created.
  std::unique_ptr<Cmac> cmac = Cmac::Create(nullptr, sizeof(kKey128));
  EXPECT_FALSE(cmac);
  cmac = Cmac::Create(nullptr, sizeof(kKey256));
  EXPECT_FALSE(cmac);
  cmac = Cmac::Create(kKey128, 0);
  EXPECT_FALSE(cmac);
  cmac = Cmac::Create(kKey256, (sizeof(kKey128) + sizeof(kKey256)) / 2);
  EXPECT_FALSE(cmac);
}

TEST(OEMCryptoCmacTest, CmacAes128) {
  std::unique_ptr<Cmac> cmac = Cmac::Create(kKey128, sizeof(kKey128));
  ASSERT_TRUE(cmac);

  std::vector<uint8_t> digest;
  ASSERT_TRUE(cmac->Finalize(&digest));
  const std::vector<uint8_t> kExpectedEmptyDigest = {
      0xbb, 0x1d, 0x69, 0x29, 0xe9, 0x59, 0x37, 0x28,
      0x7f, 0xa3, 0x7d, 0x12, 0x9b, 0x75, 0x67, 0x46};
  EXPECT_EQ(kExpectedEmptyDigest, digest);
  cmac->Reset();

  ASSERT_TRUE(cmac->Update(kDataChunk1, sizeof(kDataChunk1)));
  ASSERT_TRUE(cmac->Finalize(&digest));
  const std::vector<uint8_t> kExpectedDigest1 = {
      0x07, 0x0a, 0x16, 0xb4, 0x6b, 0x4d, 0x41, 0x44,
      0xf7, 0x9b, 0xdd, 0x9d, 0xd0, 0x4a, 0x28, 0x7c};
  EXPECT_EQ(kExpectedDigest1, digest);
  cmac->Reset();

  ASSERT_TRUE(cmac->Update(kDataChunk1, sizeof(kDataChunk1)));
  ASSERT_TRUE(cmac->Update(kDataChunk2, sizeof(kDataChunk2)));
  ASSERT_TRUE(cmac->Finalize(&digest));
  const std::vector<uint8_t> kExpectedDigest2 = {
      0xdf, 0xa6, 0x67, 0x47, 0xde, 0x9a, 0xe6, 0x30,
      0x30, 0xca, 0x32, 0x61, 0x14, 0x97, 0xc8, 0x27};
  EXPECT_EQ(kExpectedDigest2, digest);
  cmac->Reset();

  ASSERT_TRUE(cmac->Update(kDataChunk1, sizeof(kDataChunk1)));
  for (size_t i = 0; i < sizeof(kDataChunk2); i++) {
    ASSERT_TRUE(cmac->Update(kDataChunk2[i])) << " i = " << i;
  }
  ASSERT_TRUE(cmac->Update(kDataChunk3, sizeof(kDataChunk3)));
  ASSERT_TRUE(cmac->Finalize(&digest));
  const std::vector<uint8_t> kExpectedDigest3 = {
      0x51, 0xf0, 0xbe, 0xbf, 0x7e, 0x3b, 0x9d, 0x92,
      0xfc, 0x49, 0x74, 0x17, 0x79, 0x36, 0x3c, 0xfe};
  EXPECT_EQ(kExpectedDigest3, digest);
}

TEST(OEMCryptoCmacTest, CmacAes256) {
  std::unique_ptr<Cmac> cmac = Cmac::Create(kKey256, sizeof(kKey256));
  ASSERT_TRUE(cmac);

  std::vector<uint8_t> digest;
  ASSERT_TRUE(cmac->Finalize(&digest));
  const std::vector<uint8_t> kExpectedEmptyDigest = {
      0x02, 0x89, 0x62, 0xf6, 0x1b, 0x7b, 0xf8, 0x9e,
      0xfc, 0x6b, 0x55, 0x1f, 0x46, 0x67, 0xd9, 0x83};
  EXPECT_EQ(kExpectedEmptyDigest, digest);
  cmac->Reset();

  ASSERT_TRUE(cmac->Update(kDataChunk1, sizeof(kDataChunk1)));
  ASSERT_TRUE(cmac->Finalize(&digest));
  const std::vector<uint8_t> kExpectedDigest1 = {
      0x28, 0xa7, 0x02, 0x3f, 0x45, 0x2e, 0x8f, 0x82,
      0xbd, 0x4b, 0xf2, 0x8d, 0x8c, 0x37, 0xc3, 0x5c};
  EXPECT_EQ(kExpectedDigest1, digest);
  cmac->Reset();

  ASSERT_TRUE(cmac->Update(kDataChunk1, sizeof(kDataChunk1)));
  ASSERT_TRUE(cmac->Update(kDataChunk2, sizeof(kDataChunk2)));
  ASSERT_TRUE(cmac->Finalize(&digest));
  const std::vector<uint8_t> kExpectedDigest2 = {
      0xaa, 0xf3, 0xd8, 0xf1, 0xde, 0x56, 0x40, 0xc2,
      0x32, 0xf5, 0xb1, 0x69, 0xb9, 0xc9, 0x11, 0xe6};
  EXPECT_EQ(kExpectedDigest2, digest);
  cmac->Reset();

  ASSERT_TRUE(cmac->Update(kDataChunk1, sizeof(kDataChunk1)));
  for (size_t i = 0; i < sizeof(kDataChunk2); i++) {
    ASSERT_TRUE(cmac->Update(kDataChunk2[i])) << " i = " << i;
  }
  ASSERT_TRUE(cmac->Update(kDataChunk3, sizeof(kDataChunk3)));
  ASSERT_TRUE(cmac->Finalize(&digest));
  const std::vector<uint8_t> kExpectedDigest3 = {
      0xe1, 0x99, 0x21, 0x90, 0x54, 0x9f, 0x6e, 0xd5,
      0x69, 0x6a, 0x2c, 0x05, 0x6c, 0x31, 0x54, 0x10};
  EXPECT_EQ(kExpectedDigest3, digest);
}
}  // namespace util
}  // namespace wvoec