/*
 *  Copyright (c) 2018-present, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under the BSD-style license found in the
 *  LICENSE file in the root directory of this source tree.
 */
#include <fizz/backend/openssl/OpenSSL.h>
#include <fizz/client/PskSerializationUtils.h>
#include <fizz/client/test/Utilities.h>
#include <fizz/crypto/test/TestUtil.h>
#include <fizz/protocol/test/Mocks.h>
#include <folly/portability/GMock.h>
#include <folly/portability/GTest.h>

using namespace testing;

namespace fizz {
namespace client {
namespace test {

static std::shared_ptr<const fizz::Cert> getCert(
    folly::StringPiece encodedPem) {
  auto certs = folly::ssl::OpenSSLCertUtils::readCertsFromBuffer(encodedPem);
  CHECK_EQ(certs.size(), 1);
  return openssl::CertUtils::makePeerCert(std::move(certs[0]));
}
class PskSerializationTest : public Test {
 public:
  void SetUp() override {
    ticketTime_ = std::chrono::system_clock::now();
  }

 protected:
  CachedPsk getCachedPsk(std::string pskName) {
    auto psk = getTestPsk(pskName, ticketTime_);
    psk.clientCert = getCert(fizz::test::kClientAuthClientCert);
    psk.serverCert = getCert(fizz::test::kRSACertificate);
    return psk;
  }

  std::chrono::system_clock::time_point ticketTime_;
};

TEST_F(PskSerializationTest, TestDeserializationCompatibility) {
  // This is the output of the `TestSerialization` serializePsk call, at
  // a specific commit. If this test fails, this implies that your change
  // may cause old PSKs to no longer be recognized.
  constexpr folly::StringPiece kSerializedPSK =
      " ";

  auto psk = fizz::client::deserializePsk(
      fizz::openssl::certificateSerializer(),
      folly::ByteRange(folly::unhexlify(kSerializedPSK)));
  EXPECT_EQ(psk.psk, "SomePsk");
  EXPECT_EQ(psk.secret, "resumptionsecret");
  EXPECT_EQ(psk.type, PskType::Resumption);
  EXPECT_EQ(psk.version, ProtocolVersion::tls_1_3);
  EXPECT_EQ(psk.cipher, CipherSuite::TLS_AES_128_GCM_SHA256);
  EXPECT_EQ(psk.group, NamedGroup::x25519);
  EXPECT_EQ(psk.maxEarlyDataSize, 0);
  EXPECT_EQ(psk.ticketAgeAdd, 0x11111111);
  EXPECT_EQ(psk.alpn, "h2");
  EXPECT_EQ(psk.clientCert->getIdentity(), "Fizz Client");
  EXPECT_EQ(psk.serverCert->getIdentity(), "Fizz");
}

static std::shared_ptr<fizz::Cert> makeMockCertWithIdentity(
    const std::string& ident) {
  auto cert = std::make_shared<fizz::test::MockCert>();
  EXPECT_CALL(*cert, getIdentity()).WillRepeatedly(Return(ident));
  return cert;
}

TEST_F(PskSerializationTest, TestSerializationWithMock) {
  auto psk = getCachedPsk("SomePsk");
  psk.clientCert = makeMockCertWithIdentity("Fizz Client");
  psk.serverCert = makeMockCertWithIdentity("Fizz Server");
  StrictMock<fizz::test::MockCertificateSerialization> mockSerializer;
  EXPECT_CALL(mockSerializer, serialize(_))
      .Times(2)
      .WillRepeatedly(Invoke([](auto&& cert) {
        std::string repr = fmt::format("serialized: {}", cert.getIdentity());
        return folly::IOBuf::copyBuffer(repr);
      }));
  EXPECT_CALL(mockSerializer, deserialize(_))
      .Times(2)
      .WillRepeatedly(Invoke([](folly::ByteRange range) {
        folly::StringPiece sp(
            reinterpret_cast<const char*>(range.data()), range.size());
        EXPECT_NE(sp.find("serialized: "), std::string::npos);
        if (sp.find("Fizz Client") != std::string::npos) {
          return makeMockCertWithIdentity("client cert");
        } else {
          return makeMockCertWithIdentity("server cert");
        }
      }));
  std::string serializedPsk = fizz::client::serializePsk(mockSerializer, psk);
  EXPECT_TRUE(serializedPsk.find("serialized: ") != std::string::npos);
  auto psk2 = fizz::client::deserializePsk(
      mockSerializer, folly::ByteRange(serializedPsk));
  EXPECT_NE(psk2.clientCert, nullptr);
  EXPECT_NE(psk2.serverCert, nullptr);

  EXPECT_EQ(psk2.clientCert->getIdentity(), "client cert");
  EXPECT_EQ(psk2.serverCert->getIdentity(), "server cert");
}

TEST_F(PskSerializationTest, TestSerializationThatFails) {
  auto psk = getCachedPsk("SomePsk");
  psk.clientCert = makeMockCertWithIdentity("Fizz Client");
  psk.serverCert = makeMockCertWithIdentity("Fizz Server");
  StrictMock<fizz::test::MockCertificateSerialization> mockSerializer;
  Sequence s;
  EXPECT_CALL(mockSerializer, serialize(_))
      .InSequence(s)
      .WillOnce(Return(nullptr));
  EXPECT_CALL(mockSerializer, serialize(_))
      .InSequence(s)
      .WillOnce(Throw(std::runtime_error("failed serialization")));

  // When serialization fails, then we expect that deserialize is not called
  EXPECT_CALL(mockSerializer, deserialize(_)).Times(0);
  std::string serializedPsk = fizz::client::serializePsk(mockSerializer, psk);
  auto psk2 = fizz::client::deserializePsk(
      mockSerializer, folly::ByteRange(serializedPsk));
  EXPECT_EQ(psk2.clientCert, nullptr);
  EXPECT_EQ(psk2.serverCert, nullptr);
}

TEST_F(PskSerializationTest, TestDeserializationThatFails) {
  auto psk = getCachedPsk("SomePsk");
  psk.clientCert = makeMockCertWithIdentity("Fizz Client");
  psk.serverCert = makeMockCertWithIdentity("Fizz Server");
  StrictMock<fizz::test::MockCertificateSerialization> mockSerializer;
  EXPECT_CALL(mockSerializer, serialize(_))
      .Times(2)
      .WillRepeatedly(
          Invoke([](auto&&) { return folly::IOBuf::copyBuffer("hello"); }));

  Sequence s;
  EXPECT_CALL(mockSerializer, deserialize(_))
      .InSequence(s)
      .WillOnce(Return(nullptr));
  EXPECT_CALL(mockSerializer, deserialize(_))
      .InSequence(s)
      .WillOnce(Throw(std::runtime_error("failed deserialization")));
  std::string serializedPsk = fizz::client::serializePsk(mockSerializer, psk);
  auto psk2 = fizz::client::deserializePsk(
      mockSerializer, folly::ByteRange(serializedPsk));
  EXPECT_EQ(psk2.clientCert, nullptr);
  EXPECT_EQ(psk2.serverCert, nullptr);
}

TEST_F(PskSerializationTest, TestSerialization) {
  auto psk = getCachedPsk("SomePsk");
  std::string serializedPsk =
      fizz::client::serializePsk(fizz::openssl::certificateSerializer(), psk);
  EXPECT_FALSE(serializedPsk.empty());
  auto psk2 = fizz::client::deserializePsk(
      fizz::openssl::certificateSerializer(), folly::ByteRange(serializedPsk));
  EXPECT_EQ(psk.psk, psk2.psk);
  EXPECT_EQ(psk.secret, psk2.secret);
  EXPECT_EQ(psk.type, psk2.type);
  EXPECT_EQ(psk.version, psk2.version);
  EXPECT_EQ(psk.cipher, psk2.cipher);
  EXPECT_EQ(psk.group, psk2.group);
  EXPECT_EQ(psk.maxEarlyDataSize, psk2.maxEarlyDataSize);
  EXPECT_EQ(psk.ticketAgeAdd, psk2.ticketAgeAdd);
  EXPECT_EQ(psk.alpn, psk2.alpn);
  ASSERT_NE(psk.clientCert, nullptr);
  EXPECT_EQ(psk.clientCert->getIdentity(), "Fizz Client");
  ASSERT_NE(psk.serverCert, nullptr);
  EXPECT_EQ(psk.serverCert->getIdentity(), "Fizz");
  // Time may have slight rounding errors, make sure its negligible
  auto ticketHandshakeTimeDiff =
      std::chrono::duration_cast<std::chrono::milliseconds>(
          psk.ticketHandshakeTime - psk2.ticketHandshakeTime);
  auto ticketExpirationTimeDiff =
      std::chrono::duration_cast<std::chrono::seconds>(
          psk.ticketExpirationTime - psk2.ticketExpirationTime);
  auto ticketIssueTimeDiff =
      std::chrono::duration_cast<std::chrono::milliseconds>(
          psk.ticketIssueTime - psk2.ticketIssueTime);
  EXPECT_EQ(ticketHandshakeTimeDiff.count(), 0);
  EXPECT_EQ(ticketExpirationTimeDiff.count(), 0);
  EXPECT_EQ(ticketIssueTimeDiff.count(), 0);
}

TEST_F(PskSerializationTest, TestTimestampOverflow) {
  // The timestamps for all of the time points are
  // numeric_limits<int64_t>::max()
  constexpr folly::StringPiece kValidPskWithBadTimestamps =
      "0007536f6d6550736b0010726573756d7074696f6e7365637265740304130101001d026832111111117fffffffffffffff7fffffffffffffff0000000000000000000000017fffffffffffffff";
  auto psk = fizz::client::deserializePsk(
      fizz::openssl::certificateSerializer(),
      folly::ByteRange(folly::unhexlify(kValidPskWithBadTimestamps)));

  EXPECT_GT(psk.ticketIssueTime.time_since_epoch().count(), 0);
  EXPECT_GT(psk.ticketExpirationTime.time_since_epoch().count(), 0);
  EXPECT_GT(psk.ticketHandshakeTime.time_since_epoch().count(), 0);
}

TEST_F(PskSerializationTest, TestInvalidSerializationThrows) {
  auto psk = getCachedPsk("SomePsk");
  std::string serializedPsk =
      fizz::client::serializePsk(fizz::openssl::certificateSerializer(), psk);
  EXPECT_FALSE(serializedPsk.empty());
  serializedPsk.insert(10, "HI!");
  ASSERT_ANY_THROW({
    fizz::client::deserializePsk(
        fizz::openssl::certificateSerializer(),
        folly::ByteRange(serializedPsk));
  });
}

} // namespace test
} // namespace client
} // namespace fizz
