// Copyright 2016 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "media/formats/webm/webm_crypto_helpers.h"

#include <array>

#include "base/containers/auto_spanification_helper.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"

using ::testing::ElementsAre;

namespace {

const auto kKeyId = std::to_array<uint8_t>({
    0x01,
    0x02,
    0x03,
    0x04,
    0x05,
    0x06,
    0x07,
    0x08,
});

}  // namespace

namespace media {

TEST(WebMCryptoHelpersTest, EmptyData) {
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_FALSE(WebMCreateDecryptConfig(
      nullptr, 0, kKeyId.data(), base::SpanificationSizeofForStdArray(kKeyId),
      &decrypt_config, &data_offset));
}

TEST(WebMCryptoHelpersTest, ClearData) {
  const uint8_t kData[] = {0x00, 0x0d, 0x0a, 0x0d, 0x0a};
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_TRUE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
  EXPECT_EQ(1u, data_offset);
  EXPECT_FALSE(decrypt_config);
}

TEST(WebMCryptoHelpersTest, EncryptedButNotEnoughBytes) {
  const uint8_t kData[] = {0x01, 0x0d, 0x0a, 0x0d, 0x0a};
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_FALSE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedNotPartitioned) {
  const uint8_t kData[] = {
      // Encrypted
      0x01,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
      // Data
      0x01, 0x02,
  };
  // Extracted from kData and zero extended to 16 bytes.
  const auto kExpectedIv = std::to_array<uint8_t>({
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
  });
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_TRUE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
  EXPECT_TRUE(decrypt_config);
  EXPECT_EQ(
      std::string(kKeyId.data(),
                  base::span(kKeyId)
                      .subspan(base::SpanificationSizeofForStdArray(kKeyId))
                      .data()),
      decrypt_config->key_id());
  EXPECT_EQ(std::string(
                kExpectedIv.data(),
                base::span(kExpectedIv)
                    .subspan(base::SpanificationSizeofForStdArray(kExpectedIv))
                    .data()),
            decrypt_config->iv());
  EXPECT_TRUE(decrypt_config->subsamples().empty());
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedMissingNumPartitionField) {
  const uint8_t kData[] = {
      // Encrypted and Partitioned
      0x03,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
  };
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_FALSE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedNotEnoughBytesForOffsets) {
  const uint8_t kData[] = {
      // Encrypted and Partitioned
      0x03,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
      // Num partitions = 2
      0x02,
      // Partition 0 @ offset 3
      0x00, 0x00, 0x00, 0x03,
  };
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_FALSE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedNotEnoughBytesForData) {
  const uint8_t kData[] = {
      // Encrypted and Partitioned
      0x03,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
      // Num partitions = 2
      0x02,
      // Partition 0 @ offset 3, partition 2 @ offset 5
      0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x05,
      // Should have more than 5 bytes of data
      0x00, 0x01, 0x02, 0x03,
  };
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_FALSE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedNotEnoughBytesForData2) {
  const uint8_t kData[] = {
      // Encrypted and Partitioned
      0x03,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
      // Num partitions = 2
      0x02,
      // Partition 0 @ offset 3, partition 1 @ offset 5
      0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x05,
      // Should have more than 5 bytes of data
      0x00, 0x01, 0x02, 0x03, 0x04,
  };
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_FALSE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedDecreasingOffsets) {
  const uint8_t kData[] = {
      // Encrypted and Partitioned
      0x03,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
      // Num partitions = 2
      0x02,
      // Partition 0 @ offset 3, partition 1 @ offset 2
      0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x02,
      // Should have more than 5 bytes of data
      0x00, 0x01, 0x02, 0x03, 0x04,
  };
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_FALSE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedEvenNumberOfPartitions) {
  const uint8_t kData[] = {
      // Encrypted and Partitioned
      0x03,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
      // Num partitions = 2
      0x02,
      // Partition 0 @ offset 3, partition 1 @ offset 5
      0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x05,
      // Should have more than 5 bytes of data
      0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
  };
  // Extracted from kData and zero extended to 16 bytes.
  const auto kExpectedIv = std::to_array<uint8_t>({
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
  });
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_TRUE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
  EXPECT_TRUE(decrypt_config);
  EXPECT_EQ(
      std::string(kKeyId.data(),
                  base::span(kKeyId)
                      .subspan(base::SpanificationSizeofForStdArray(kKeyId))
                      .data()),
      decrypt_config->key_id());
  EXPECT_EQ(std::string(
                kExpectedIv.data(),
                base::span(kExpectedIv)
                    .subspan(base::SpanificationSizeofForStdArray(kExpectedIv))
                    .data()),
            decrypt_config->iv());
  EXPECT_THAT(decrypt_config->subsamples(),
              ElementsAre(SubsampleEntry(3, 2), SubsampleEntry(1, 0)));
  EXPECT_EQ(18u, data_offset);
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedOddNumberOfPartitions) {
  const uint8_t kData[] = {
      // Encrypted and Partitioned
      0x03,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
      // Num partitions = 1
      0x01,
      // Partition 0 @ offset 3,
      0x00, 0x00, 0x00, 0x03,
      // Should have more than 3 bytes of data
      0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
  };
  // Extracted from kData and zero extended to 16 bytes.
  const auto kExpectedIv = std::to_array<uint8_t>({
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
  });
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_TRUE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
  EXPECT_TRUE(decrypt_config);
  EXPECT_EQ(
      std::string(kKeyId.data(),
                  base::span(kKeyId)
                      .subspan(base::SpanificationSizeofForStdArray(kKeyId))
                      .data()),
      decrypt_config->key_id());
  EXPECT_EQ(std::string(
                kExpectedIv.data(),
                base::span(kExpectedIv)
                    .subspan(base::SpanificationSizeofForStdArray(kExpectedIv))
                    .data()),
            decrypt_config->iv());
  EXPECT_THAT(decrypt_config->subsamples(), ElementsAre(SubsampleEntry(3, 3)));
  EXPECT_EQ(14u, data_offset);
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedZeroNumberOfPartitions) {
  const uint8_t kData[] = {
      // Encrypted and Partitioned
      0x03,
      // IV
      0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a, 0x0d, 0x0a,
      // Num partitions = 0
      0x00,
      // Some random data.
      0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
  };
  // Extracted from kData and zero extended to 16 bytes.
  const auto kExpectedIv = std::to_array<uint8_t>({
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x0d,
      0x0a,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
      0x00,
  });
  std::unique_ptr<DecryptConfig> decrypt_config;
  size_t data_offset;
  ASSERT_TRUE(
      WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId.data(),
                              base::SpanificationSizeofForStdArray(kKeyId),
                              &decrypt_config, &data_offset));
  EXPECT_TRUE(decrypt_config);
  EXPECT_EQ(
      std::string(kKeyId.data(),
                  base::span(kKeyId)
                      .subspan(base::SpanificationSizeofForStdArray(kKeyId))
                      .data()),
      decrypt_config->key_id());
  EXPECT_EQ(std::string(
                kExpectedIv.data(),
                base::span(kExpectedIv)
                    .subspan(base::SpanificationSizeofForStdArray(kExpectedIv))
                    .data()),
            decrypt_config->iv());
  EXPECT_THAT(decrypt_config->subsamples(), ElementsAre(SubsampleEntry(6, 0)));
  EXPECT_EQ(10u, data_offset);
}

}  // namespace media