/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this file,
 * You can obtain one at http://mozilla.org/MPL/2.0/. */

#include <array>
#include <cstdint>
#include <numeric>
#include <set>
#include <vector>

#include "NSSCipherStrategy.h"
#include "gtest/gtest.h"
#include "mozilla/Span.h"

namespace mozilla::dom::quota::test {

namespace {

constexpr size_t kPayloadLength = 64;
constexpr size_t kNonceLength = 12;

using IvArray = std::array<uint8_t, NSSCipherStrategy::BlockPrefixLength>;
using Payload = std::array<uint8_t, kPayloadLength>;

NSSCipherStrategy::KeyType MakeKey(uint8_t aInitialValue = 1) {
  NSSCipherStrategy::KeyType key;
  std::iota(key.begin(), key.end(), aInitialValue);
  return key;
}

Payload MakePlaintext(uint8_t aInitialValue = 1) {
  Payload plaintext;
  std::iota(plaintext.begin(), plaintext.end(), aInitialValue);
  return plaintext;
}

std::vector<uint8_t> NonceOf(const IvArray& aIv) {
  return std::vector<uint8_t>(aIv.begin(), aIv.begin() + kNonceLength);
}

}  // namespace

TEST(NSSCipherStrategyTest, NonceDiffersAcrossCallsInOneContext)
{
  const auto key = MakeKey();
  const auto plaintext = MakePlaintext();

  NSSCipherStrategy encrypt;
  ASSERT_EQ(NS_OK, encrypt.Init(CipherMode::Encrypt,
                                NSSCipherStrategy::SerializeKey(key)));

  std::set<std::vector<uint8_t>> nonces;
  for (size_t i = 0; i < 32; ++i) {
    IvArray iv{};
    Payload ciphertext{};
    ASSERT_EQ(NS_OK,
              encrypt.Cipher(Span{iv}, Span{plaintext}, Span{ciphertext}));
    EXPECT_TRUE(nonces.insert(NonceOf(iv)).second)
        << "nonce repeated within a single context on call " << i;
  }
}

// Regression test for bug 2055694. ObfuscatingVFS creates a separate encrypt
// context per SQLite file handle and per reopen, all sharing one key. When the
// nonce came from a per-context counter starting at zero, those contexts
// produced identical (key, nonce) pairs, so XOR of two ciphertexts disclosed
// XOR of their plaintexts.
TEST(NSSCipherStrategyTest, NonceDiffersAcrossContextsWithTheSameKey)
{
  const auto key = MakeKey();
  const auto plaintext = MakePlaintext();

  const auto encryptOnce = [&key, &plaintext](IvArray& aIv,
                                              Payload& aCiphertext) {
    NSSCipherStrategy encrypt;
    ASSERT_EQ(NS_OK, encrypt.Init(CipherMode::Encrypt,
                                  NSSCipherStrategy::SerializeKey(key)));
    ASSERT_EQ(NS_OK,
              encrypt.Cipher(Span{aIv}, Span{plaintext}, Span{aCiphertext}));
  };

  IvArray firstIv{};
  Payload firstCiphertext{};
  encryptOnce(firstIv, firstCiphertext);

  IvArray secondIv{};
  Payload secondCiphertext{};
  encryptOnce(secondIv, secondCiphertext);

  EXPECT_NE(NonceOf(firstIv), NonceOf(secondIv));
  EXPECT_NE(firstCiphertext, secondCiphertext);
}

// Everything ahead of the tag is generated by Cipher, so a caller that leaves
// the block prefix uninitialized cannot leak its contents to disk.
TEST(NSSCipherStrategyTest, OverwritesTheWholePrefixRegardlessOfCallerInput)
{
  const auto key = MakeKey();
  const auto plaintext = MakePlaintext();

  NSSCipherStrategy encrypt;
  ASSERT_EQ(NS_OK, encrypt.Init(CipherMode::Encrypt,
                                NSSCipherStrategy::SerializeKey(key)));

  constexpr size_t kTagLength = 16;
  constexpr uint8_t kSentinel = 0xAA;

  std::set<std::vector<uint8_t>> fillers;
  for (size_t i = 0; i < 32; ++i) {
    IvArray iv;
    iv.fill(kSentinel);
    Payload ciphertext{};
    ASSERT_EQ(NS_OK,
              encrypt.Cipher(Span{iv}, Span{plaintext}, Span{ciphertext}));

    fillers.emplace(iv.begin() + kNonceLength, iv.end() - kTagLength);
  }

  // Had the filler bytes been left to the caller, every sample would still
  // hold the sentinel and this set would have exactly one entry.
  EXPECT_GT(fillers.size(), 1u);
}

TEST(NSSCipherStrategyTest, RoundTripsUsingTheStoredNonce)
{
  const auto key = MakeKey();
  const auto plaintext = MakePlaintext();

  IvArray iv{};
  Payload ciphertext{};
  {
    NSSCipherStrategy encrypt;
    ASSERT_EQ(NS_OK, encrypt.Init(CipherMode::Encrypt,
                                  NSSCipherStrategy::SerializeKey(key)));
    ASSERT_EQ(NS_OK,
              encrypt.Cipher(Span{iv}, Span{plaintext}, Span{ciphertext}));
  }
  EXPECT_NE(plaintext, ciphertext);

  Payload decrypted{};
  {
    NSSCipherStrategy decrypt;
    ASSERT_EQ(NS_OK, decrypt.Init(CipherMode::Decrypt,
                                  NSSCipherStrategy::SerializeKey(key)));
    ASSERT_EQ(NS_OK,
              decrypt.Cipher(Span{iv}, Span{ciphertext}, Span{decrypted}));
  }
  EXPECT_EQ(plaintext, decrypted);
}

// Decryption must depend only on the nonce stored alongside the ciphertext,
// never on a context-local counter. This is what keeps blocks written by older
// builds -- whose nonces were sequential -- readable.
TEST(NSSCipherStrategyTest, DecryptionIsIndependentOfBlockOrder)
{
  const auto key = MakeKey();
  const auto firstPlaintext = MakePlaintext(1);
  const auto secondPlaintext = MakePlaintext(101);

  IvArray firstIv{};
  Payload firstCiphertext{};
  IvArray secondIv{};
  Payload secondCiphertext{};
  {
    NSSCipherStrategy encrypt;
    ASSERT_EQ(NS_OK, encrypt.Init(CipherMode::Encrypt,
                                  NSSCipherStrategy::SerializeKey(key)));
    ASSERT_EQ(NS_OK, encrypt.Cipher(Span{firstIv}, Span{firstPlaintext},
                                    Span{firstCiphertext}));
    ASSERT_EQ(NS_OK, encrypt.Cipher(Span{secondIv}, Span{secondPlaintext},
                                    Span{secondCiphertext}));
  }

  NSSCipherStrategy decrypt;
  ASSERT_EQ(NS_OK, decrypt.Init(CipherMode::Decrypt,
                                NSSCipherStrategy::SerializeKey(key)));

  Payload secondDecrypted{};
  ASSERT_EQ(NS_OK, decrypt.Cipher(Span{secondIv}, Span{secondCiphertext},
                                  Span{secondDecrypted}));
  EXPECT_EQ(secondPlaintext, secondDecrypted);

  Payload firstDecrypted{};
  ASSERT_EQ(NS_OK, decrypt.Cipher(Span{firstIv}, Span{firstCiphertext},
                                  Span{firstDecrypted}));
  EXPECT_EQ(firstPlaintext, firstDecrypted);
}

}  // namespace mozilla::dom::quota::test