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

#include "components/signin/public/webdata/token_service_table.h"

#include <memory>
#include <string>

#include "base/containers/to_vector.h"
#include "base/files/scoped_temp_dir.h"
#include "base/memory/scoped_refptr.h"
#include "base/run_loop.h"
#include "base/strings/string_number_conversions.h"
#include "base/test/bind.h"
#include "base/test/metrics/histogram_tester.h"
#include "base/test/test_future.h"
#include "base/time/time.h"
#include "build/build_config.h"
#include "components/os_crypt/async/browser/os_crypt_async.h"
#include "components/os_crypt/async/browser/test_utils.h"
#include "components/os_crypt/async/common/encryptor.h"
#include "components/webdata/common/web_database.h"
#include "sql/statement.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"

using ::base::Time;
using ::testing::ElementsAre;
using ::testing::IsEmpty;
using ::testing::Key;
using ::testing::Optional;
using ::testing::SizeIs;
using ::testing::UnorderedElementsAre;
using ::testing::UnorderedElementsAreArray;
using TokenWithBindingInfo = TokenServiceTable::TokenWithBindingInfo;

class TokenServiceTableTest : public testing::Test {
 public:
  TokenServiceTableTest()
      : encryptor_(os_crypt_async::GetTestEncryptorForTesting()) {}

  TokenServiceTableTest(const TokenServiceTableTest&) = delete;
  TokenServiceTableTest& operator=(const TokenServiceTableTest&) = delete;

  ~TokenServiceTableTest() override = default;

 protected:
  void SetUp() override {
    ASSERT_TRUE(temp_dir_.CreateUniqueTempDir());
    file_ = temp_dir_.GetPath().AppendASCII("TestWebDatabase");

    table_ = std::make_unique<TokenServiceTable>();
    db_ = std::make_unique<WebDatabase>();
    db_->AddTable(table_.get());
    ASSERT_EQ(sql::INIT_OK, db_->Init(file_, encryptor_));
  }

  base::FilePath file_;
  base::ScopedTempDir temp_dir_;
  scoped_refptr<const os_crypt_async::Encryptor> encryptor_;
  std::unique_ptr<TokenServiceTable> table_;
  std::unique_ptr<WebDatabase> db_;
};

TEST_F(TokenServiceTableTest, TokenServiceGetAllRemoveAll) {
  std::map<std::string, TokenWithBindingInfo> out_map;
  std::string service;
  std::string service2;
  service = "testservice";
  service2 = "othertestservice";
  bool should_reencrypt = false;

  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_TRUE(out_map.empty());

  // Check that get all tokens works
  EXPECT_TRUE(table_->SetTokenForService(service, "pepperoni", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->SetTokenForService(service2, "steak", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo("pepperoni"), out_map.find(service)->second);
  EXPECT_EQ(TokenWithBindingInfo("steak"), out_map.find(service2)->second);
  out_map.clear();

  // Purge
  EXPECT_TRUE(table_->RemoveAllTokens());
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_TRUE(out_map.empty());

  // Check that you can still add it back in
  EXPECT_TRUE(table_->SetTokenForService(service, "cheese", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo("cheese"), out_map.find(service)->second);
}

TEST_F(TokenServiceTableTest, TokenServiceGetAllWrappedBindingKeys) {
  EXPECT_THAT(table_->GetAllWrappedBindingKeys(), Optional(IsEmpty()));

  EXPECT_TRUE(table_->SetTokenForService("service1", "token1", {1, 2, 3},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->SetTokenForService("service2", "token2", {4, 5, 6},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->SetTokenForService("service3", "token3", {7, 8, 9},
                                         /*mtls_token_binding=*/false));
  EXPECT_THAT(
      table_->GetAllWrappedBindingKeys(),
      Optional(UnorderedElementsAre(ElementsAre(1, 2, 3), ElementsAre(4, 5, 6),
                                    ElementsAre(7, 8, 9))));

  EXPECT_TRUE(table_->RemoveTokenForService("service1"));
  EXPECT_THAT(table_->GetAllWrappedBindingKeys(),
              Optional(UnorderedElementsAre(ElementsAre(4, 5, 6),
                                            ElementsAre(7, 8, 9))));
}

TEST_F(TokenServiceTableTest, TokenServiceGetSet) {
  std::map<std::string, TokenWithBindingInfo> out_map;
  std::string service;
  service = "testservice";
  bool should_reencrypt = false;

  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_TRUE(out_map.empty());

  EXPECT_TRUE(table_->SetTokenForService(service, "pepperoni", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo("pepperoni"), out_map.find(service)->second);
  out_map.clear();

  // try blanking it - won't remove it from the db though!
  EXPECT_TRUE(table_->SetTokenForService(service, std::string(), {},
                                         /*mtls_token_binding=*/false));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo(""), out_map.find(service)->second);
  out_map.clear();

  // try mutating it
  EXPECT_TRUE(table_->SetTokenForService(service, "ham", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo("ham"), out_map.find(service)->second);
}

TEST_F(TokenServiceTableTest, TokenServiceRemove) {
  std::map<std::string, TokenWithBindingInfo> out_map;
  std::string service;
  std::string service2;
  service = "testservice";
  service2 = "othertestservice";
  bool should_reencrypt = false;

  EXPECT_TRUE(table_->SetTokenForService(service, "pepperoni", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->SetTokenForService(service2, "steak", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->RemoveTokenForService(service));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(0u, out_map.count(service));
  EXPECT_EQ(TokenWithBindingInfo("steak"), out_map.find(service2)->second);
}

TEST_F(TokenServiceTableTest, TokenServiceRemoveOther) {
  EXPECT_TRUE(table_->SetTokenForService("a", "1", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->SetTokenForService("b", "2", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->SetTokenForService("c", "3", {},
                                         /*mtls_token_binding=*/false));

  base::HistogramTester histogram_tester;
  EXPECT_TRUE(table_->RemoveOtherTokens({"a", "c", "zzz"}));

  std::map<std::string, TokenWithBindingInfo> out_map;
  bool should_reencrypt = false;
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_THAT(out_map, UnorderedElementsAre(
                           std::make_pair("a", TokenWithBindingInfo("1")),
                           std::make_pair("c", TokenWithBindingInfo("3"))));
  histogram_tester.ExpectUniqueSample(
      "Signin.TokenTable.RemoveOtherTokensCount",
      /*sample=*/1, /*expected_bucket_count=*/1);
}

TEST_F(TokenServiceTableTest, TokenServiceRemoveOtherKeepNone) {
  EXPECT_TRUE(table_->SetTokenForService("a", "1", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->SetTokenForService("b", "2", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_TRUE(table_->SetTokenForService("c", "3", {},
                                         /*mtls_token_binding=*/false));

  base::HistogramTester histogram_tester;
  EXPECT_TRUE(table_->RemoveOtherTokens({}));

  std::map<std::string, TokenWithBindingInfo> out_map;
  bool should_reencrypt = false;
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_THAT(out_map, IsEmpty());
  histogram_tester.ExpectUniqueSample(
      "Signin.TokenTable.RemoveOtherTokensCount",
      /*sample=*/3, /*expected_bucket_count=*/1);
}

class TokenServiceTableRemoveOtherStressTest
    : public TokenServiceTableTest,
      public testing::WithParamInterface<size_t> {};

// Tests variable `services_to_keep` vector sizes in `RemoveOtherTokens()`.
TEST_P(TokenServiceTableRemoveOtherStressTest, TokenServiceRemoveOtherStress) {
  const size_t test_size = GetParam();

  std::vector<std::string> services_to_keep;
  for (size_t i = 0; i < test_size; ++i) {
    services_to_keep.push_back("keep_" + base::NumberToString(i));
    EXPECT_TRUE(table_->SetTokenForService(services_to_keep[i], "keep_token",
                                           {},
                                           /*mtls_token_binding=*/false));
    EXPECT_TRUE(table_->SetTokenForService("remove_" + base::NumberToString(i),
                                           "remove_token", {},
                                           /*mtls_token_binding=*/false));
  }

  base::HistogramTester histogram_tester;
  EXPECT_TRUE(table_->RemoveOtherTokens(services_to_keep));

  std::map<std::string, TokenWithBindingInfo> out_map;
  bool should_reencrypt = false;
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));

  std::vector<std::pair<std::string, TokenWithBindingInfo>> expected_pairs =
      base::ToVector(services_to_keep, [](const std::string& service) {
        return std::make_pair(service, TokenWithBindingInfo("keep_token"));
      });
  EXPECT_THAT(out_map, UnorderedElementsAreArray(expected_pairs));
  histogram_tester.ExpectUniqueSample(
      "Signin.TokenTable.RemoveOtherTokensCount",
      /*sample=*/test_size,
      /*expected_bucket_count=*/1);
}

INSTANTIATE_TEST_SUITE_P(,
                         TokenServiceTableRemoveOtherStressTest,
                         testing::Values(0u, 1u, 10u, 42u, 100u));

TEST_F(TokenServiceTableTest, GetSetWithBidningKey) {
  bool should_reencrypt = false;
  std::map<std::string, TokenWithBindingInfo> out_map;
  const std::string kService = "testservice";
  const std::vector<uint8_t> kBindingKey = {1, 4, 2};

  EXPECT_TRUE(table_->SetTokenForService(kService, "pepperoni", kBindingKey,
                                         /*mtls_token_binding=*/false));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo("pepperoni", kBindingKey),
            out_map.find(kService)->second);
  out_map.clear();

  // Override with a new token with a new binding key.
  const std::vector<uint8_t> kNewBindingKey = {4, 8, 15, 23};
  EXPECT_TRUE(table_->SetTokenForService(kService, "ham", kNewBindingKey,
                                         /*mtls_token_binding=*/false));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo("ham", kNewBindingKey),
            out_map.find(kService)->second);
  out_map.clear();

  // Override with a new token without a binding key.
  EXPECT_TRUE(table_->SetTokenForService(kService, "steak", {},
                                         /*mtls_token_binding=*/false));
  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo("steak"), out_map.find(kService)->second);
  out_map.clear();
}

TEST_F(TokenServiceTableTest, GetNullMtlsTokenBinding) {
  bool should_reencrypt = false;
  std::map<std::string, TokenWithBindingInfo> out_map;
  const std::string kService = "testservice";

  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_TRUE(out_map.find(kService) == out_map.end());

  // Test reading null values
  std::string encrypted_token;
  ASSERT_TRUE(encryptor_->EncryptString("pepperoni", &encrypted_token));

  // Manually insert an entry without setting `mtls_token_binding` to verify
  // that having a null value is read back as `false`.
  sql::Statement s(db_->GetSQLConnection()->GetUniqueStatement(
      "INSERT OR REPLACE INTO token_service "
      "(service, encrypted_token, binding_key) VALUES (?, ?, ?)"));
  s.BindString(0, kService);
  s.BindBlob(1, encrypted_token);
  s.BindBlob(2, std::vector<uint8_t>{});
  ASSERT_TRUE(s.Run());

  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_EQ(TokenWithBindingInfo("pepperoni", /*wrapped_binding_key=*/{},
                                 /*mtls_token_binding=*/false),
            out_map.find(kService)->second);
}

TEST_F(TokenServiceTableTest, SetMtlsTokenBinding) {
  bool should_reencrypt = false;
  std::map<std::string, TokenWithBindingInfo> out_map;
  const std::string kService = "testservice";

  EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
            table_->GetAllTokens(&out_map, should_reencrypt));
  EXPECT_TRUE(out_map.find(kService) == out_map.end());

  {  // Test setting to true.
    EXPECT_TRUE(table_->SetTokenForService(kService, "pepperoni",
                                           /*wrapped_binding_key=*/{},
                                           /*mtls_token_binding=*/true));
    EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
              table_->GetAllTokens(&out_map, should_reencrypt));
    EXPECT_EQ(TokenWithBindingInfo("pepperoni", /*wrapped_binding_key=*/{},
                                   /*mtls_token_binding=*/true),
              out_map.find(kService)->second);
    out_map.clear();
  }

  {  // Test setting to false
    EXPECT_TRUE(table_->SetTokenForService(kService, "pepperoni",
                                           /*wrapped_binding_key=*/{},
                                           /*mtls_token_binding=*/false));
    EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
              table_->GetAllTokens(&out_map, should_reencrypt));
    EXPECT_EQ(TokenWithBindingInfo("pepperoni", /*wrapped_binding_key=*/{},
                                   /*mtls_token_binding=*/false),
              out_map.find(kService)->second);
  }
}

TEST_F(TokenServiceTableTest, TokenMetrics) {
  const std::string service = "testservice";
  {
    base::HistogramTester histograms;
    EXPECT_TRUE(table_->SetTokenForService(service, "pepperoni", {},
                                           /*mtls_token_binding=*/false));
    histograms.ExpectUniqueSample("Signin.TokenTable.SetTokenResult",
                                  /*kSuccess*/ 0, 1u);
  }
  {
    base::HistogramTester histograms;
    std::map<std::string, TokenWithBindingInfo> out_map;
    bool should_reencrypt;
    EXPECT_EQ(TokenServiceTable::Result::TOKEN_DB_RESULT_SUCCESS,
              table_->GetAllTokens(&out_map, should_reencrypt));
    histograms.ExpectUniqueSample("Signin.TokenTable.ReadTokenFromDBResult",
                                  /*READ_ONE_TOKEN_SUCCESS*/ 0, 1u);
  }

  {
    base::HistogramTester histograms;
    EXPECT_THAT(table_->GetAllWrappedBindingKeys(),
                Optional(UnorderedElementsAre(IsEmpty())));
    histograms.ExpectUniqueSample(
        "Signin.TokenTable.GetAllWrappedBindingKeysResult",
        /*kSuccess*/ 0, 1u);
  }
}