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

#include <latch>
#include <memory>
#include <thread>

#include "agent/src/agent_win.h"
#include "agent/src/event_win.h"
#include "browser/src/client_win.h"
#include "gtest/gtest.h"

namespace content_analysis {
namespace sdk {
namespace testing {

// A handler that counts the number of times the callback methods are invoked.
// Also remembers the last BrowserInfo structure passed to it from any of the
// callbacks.
struct TestHandler : public AgentEventHandler {
  void OnBrowserConnected(const BrowserInfo& info) override {
    last_info_ = info;
    ++connect_count_;
  }
  void OnBrowserDisconnected(const BrowserInfo& info) override {
    last_info_ = info;
    ++disconnect_count_;
  }
  void OnAnalysisRequested(
      std::unique_ptr<ContentAnalysisEvent> event) override {
    ++request_count_;
    ResultCode ret = event->Send();
    ASSERT_EQ(ResultCode::OK, ret);
  }
  void OnResponseAcknowledged(
      const ContentAnalysisAcknowledgement& ack) override {
    ++ack_count_;
  }
  void OnCancelRequests(
      const ContentAnalysisCancelRequests& cancel) override {
    ++cancel_count_;
  }

  int connect_count_ = 0;
  int disconnect_count_ = 0;
  int request_count_ = 0;
  int ack_count_ = 0;
  int cancel_count_ = 0;
  BrowserInfo last_info_;
};

// A test handler that closes its event before sending the response.
struct CloseEventTestHandler : public TestHandler {
  void OnAnalysisRequested(
    std::unique_ptr<ContentAnalysisEvent> event) override {
    ++request_count_;

    // Closing the event before sending should generate an error.
    ResultCode ret = event->Close();
    ASSERT_EQ(ResultCode::OK, ret);

    ret = event->Send();
    ASSERT_NE(ResultCode::OK, ret);
  }
};

// A test handler that attempts to send two responses for a given request.
struct DoubleSendTestHandler : public TestHandler {
  void OnAnalysisRequested(
      std::unique_ptr<ContentAnalysisEvent> event) override {
    ++request_count_;

    ResultCode ret = event->Send();
    ASSERT_EQ(ResultCode::OK, ret);

    // Trying to send again fails.
    ret = event->Send();
    ASSERT_NE(ResultCode::OK, ret);
  }
};

// A test handler that signals a latch after a client connects.
// Can only be used with one client.
struct SignalClientConnectedTestHandler : public TestHandler {
  void OnBrowserConnected(const BrowserInfo& info) override {
    TestHandler::OnBrowserConnected(info);
    wait_for_client.count_down();
  }

  std::latch wait_for_client{ 1 };
};

// A test handler that signals a latch after a request is processed.
// Can only be used with one request.
struct SignalClientRequestedTestHandler : public TestHandler {
  void OnAnalysisRequested(
      std::unique_ptr<ContentAnalysisEvent> event) override {
    TestHandler::OnAnalysisRequested(std::move(event));
    wait_for_request.count_down();
  }

  std::latch wait_for_request{ 1 };
};

std::unique_ptr<AgentWin> CreateAgent(
    Agent::Config config,
    TestHandler** handler_ptr,
    ResultCode expected_rc=ResultCode::OK) {
  ResultCode rc;
  auto handler = std::make_unique<TestHandler>();
  *handler_ptr = handler.get();
  auto agent = std::make_unique<AgentWin>(
      std::move(config), std::move(handler), &rc);
  EXPECT_EQ(expected_rc, rc);
  return agent;
}

std::unique_ptr<ClientWin> CreateClient(
    Client::Config config) {
  int rc;
  auto client = std::make_unique<ClientWin>(std::move(config), &rc);
  return rc == 0 ? std::move(client) : nullptr;
}

ContentAnalysisRequest BuildRequest(std::string content=std::string()) {
  ContentAnalysisRequest request;
  request.set_request_token("req-token");
  *request.add_tags() = "dlp";
  request.set_text_content(content);  // Moved.
  return request;
}

TEST(AgentTest, Create) {
  const Agent::Config config{"test", false};
  TestHandler* handler_ptr;
  auto agent = CreateAgent(config, &handler_ptr);
  ASSERT_TRUE(agent);
  ASSERT_TRUE(handler_ptr);

  ASSERT_EQ(config.name, agent->GetConfig().name);
  ASSERT_EQ(config.user_specific, agent->GetConfig().user_specific);
}

TEST(AgentTest, Create_InvalidPipename) {
  // TODO(rogerta): The win32 docs say that a backslash is an invalid
  // character for a pipename, but it seemed to work correctly in testing.
  // Using an empty name was the easiest way to generate an invalid pipe
  // name.
  const Agent::Config config{"", false};
  TestHandler* handler_ptr;
  auto agent = CreateAgent(config, &handler_ptr,
      ResultCode::ERR_INVALID_CHANNEL_NAME);
  ASSERT_TRUE(agent);

  ASSERT_EQ(ResultCode::ERR_AGENT_NOT_INITIALIZED,
            agent->HandleOneEventForTesting());
}

// Can't create two agents with the same name.
TEST(AgentTest, Create_SecondFails) {
  const Agent::Config config{ "test", false };
  TestHandler* handler_ptr1;
  auto agent1 = CreateAgent(config, &handler_ptr1);
  ASSERT_TRUE(agent1);

  TestHandler* handler_ptr2;
  auto agent2 = CreateAgent(config, &handler_ptr2,
      ResultCode::ERR_AGENT_ALREADY_EXISTS);
  ASSERT_TRUE(agent2);

  ASSERT_EQ(ResultCode::ERR_AGENT_NOT_INITIALIZED,
            agent2->HandleOneEventForTesting());
}

TEST(AgentTest, Stop) {
  TestHandler* handler_ptr;
  auto agent = CreateAgent({ "test", false }, &handler_ptr);
  ASSERT_TRUE(agent);

  // Create a separate thread to stop the agent.
  std::thread thread([&agent]() {
    agent->Stop();
  });

  agent->HandleEvents();
  thread.join();
}

TEST(AgentTest, ConnectAndStop) {
  ResultCode rc;
  auto handler = std::make_unique<SignalClientConnectedTestHandler>();
  auto* handler_ptr = handler.get();
  auto agent = std::make_unique<AgentWin>(
    Agent::Config{"test", false}, std::move(handler), &rc);
  ASSERT_TRUE(agent);
  ASSERT_EQ(ResultCode::OK, rc);

  // Client thread waits until latch reaches zero.
  std::latch stop_client{ 1 };

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([&stop_client]() {
    auto client = CreateClient({ "test", false });
    ASSERT_TRUE(client);
    stop_client.wait();
  });

  // A thread that stops the agent after one client connects.
  std::thread stop_agent([&handler_ptr, &agent]() {
    handler_ptr->wait_for_client.wait();
    agent->Stop();
  });

  agent->HandleEvents();

  stop_client.count_down();
  client_thread.join();
  stop_agent.join();
}

TEST(AgentTest, Connect_UserSpecific) {
  ResultCode rc;
  auto handler = std::make_unique<SignalClientConnectedTestHandler>();
  auto* handler_ptr = handler.get();
  auto agent = std::make_unique<AgentWin>(
    Agent::Config{ "test", true }, std::move(handler), &rc);
  ASSERT_TRUE(agent);
  ASSERT_EQ(ResultCode::OK, rc);

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([]() {
    // If the user_specific does not match the agent, the client should not
    // connect.
    auto client = CreateClient({ "test", false });
    ASSERT_FALSE(client);

    auto client2 = CreateClient({ "test", true });
    ASSERT_TRUE(client2);
  });

  // A thread that stops the agent after one client connects.
  std::thread stop_agent([&handler_ptr, &agent]() {
    handler_ptr->wait_for_client.wait();
    agent->Stop();
  });

  agent->HandleEvents();

  client_thread.join();
  stop_agent.join();
}

TEST(AgentTest, ConnectRequestAndStop) {
  ResultCode rc;
  auto handler = std::make_unique<SignalClientRequestedTestHandler>();
  auto* handler_ptr = handler.get();
  auto agent = std::make_unique<AgentWin>(
      Agent::Config{"test", false}, std::move(handler), &rc);
  ASSERT_TRUE(agent);
  ASSERT_EQ(ResultCode::OK, rc);

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([]() {
    auto client = CreateClient({ "test", false });
    ASSERT_TRUE(client);

    ContentAnalysisRequest request = BuildRequest("test");
    ContentAnalysisResponse response;
    client->Send(request, &response);
  });

  // A thread that stops the agent after one client connects.
  std::thread stop_agent([&handler_ptr, &agent]() {
    handler_ptr->wait_for_request.wait();
    agent->Stop();
  });

  agent->HandleEvents();

  client_thread.join();
  stop_agent.join();
}

TEST(AgentTest, ConnectAndClose) {
  const Agent::Config aconfig{ "test", false };
  const Client::Config cconfig{ "test", false };

  // Create an agent and client that connects to it.
  TestHandler* handler_ptr;
  auto agent = CreateAgent(aconfig, &handler_ptr);
  ASSERT_TRUE(agent);
  auto client = CreateClient(cconfig);
  ASSERT_TRUE(client);
  ASSERT_EQ(cconfig.name, client->GetConfig().name);
  ASSERT_EQ(cconfig.user_specific, client->GetConfig().user_specific);

  agent->HandleOneEventForTesting();
  ASSERT_EQ(1, handler_ptr->connect_count_);
  ASSERT_EQ(0, handler_ptr->disconnect_count_);
  ASSERT_EQ(0, handler_ptr->cancel_count_);
  ASSERT_EQ(GetCurrentProcessId(), handler_ptr->last_info_.pid);

  // Close the client and make sure a disconnect is received.
  client.reset();
  agent->HandleOneEventForTesting();
  ASSERT_EQ(1, handler_ptr->connect_count_);
  ASSERT_EQ(1, handler_ptr->disconnect_count_);
  ASSERT_EQ(0, handler_ptr->cancel_count_);
  ASSERT_EQ(GetCurrentProcessId(), handler_ptr->last_info_.pid);
}

TEST(AgentTest, Request) {
  TestHandler* handler_ptr;
  auto agent = CreateAgent({"test", false}, &handler_ptr);
  ASSERT_TRUE(agent);

  bool done = false;

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([&done]() {
    auto client = CreateClient({"test", false});
    ASSERT_TRUE(client);

    // Send a request and wait for a response.
    ContentAnalysisRequest request = BuildRequest();
    ContentAnalysisResponse response;
    int ret = client->Send(request, &response);
    ASSERT_EQ(0, ret);
    ASSERT_EQ(request.request_token(), response.request_token());

    done = true;
  });

  while (!done) {
    agent->HandleOneEventForTesting();
  }
  ASSERT_EQ(1, handler_ptr->request_count_);
  ASSERT_EQ(0, handler_ptr->ack_count_);
  ASSERT_EQ(0, handler_ptr->cancel_count_);

  client_thread.join();
}

TEST(AgentTest, Request_Large) {
  TestHandler* handler_ptr;
  auto agent = CreateAgent({"test", false}, &handler_ptr);
  ASSERT_TRUE(agent);

  bool done = false;

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([&done]() {
    auto client = CreateClient({"test", false});
    ASSERT_TRUE(client);

    // Send a request and wait for a response.  Create a large string, which
    // means larger than the initial mesasge buffer size specified when
    // creating the pipes (4096 bytes).
    ContentAnalysisRequest request = BuildRequest(std::string(5000, 'a'));
    ContentAnalysisResponse response;
    int ret = client->Send(request, &response);
    ASSERT_EQ(0, ret);
    ASSERT_EQ(request.request_token(), response.request_token());

    done = true;
  });

  while (!done) {
    agent->HandleOneEventForTesting();
  }
  ASSERT_EQ(1, handler_ptr->request_count_);
  ASSERT_EQ(0, handler_ptr->ack_count_);
  ASSERT_EQ(0, handler_ptr->cancel_count_);

  client_thread.join();
}

TEST(AgentTest, Request_DoubleSend) {
  ResultCode rc;
  auto handler = std::make_unique<DoubleSendTestHandler>();
  DoubleSendTestHandler* handler_ptr = handler.get();
  auto agent = std::make_unique<AgentWin>(
    Agent::Config{"test", false}, std::move(handler), &rc);
  ASSERT_TRUE(agent);
  ASSERT_EQ(ResultCode::OK, rc);

  bool done = false;

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([&done]() {
    auto client = CreateClient({ "test", false });
    ASSERT_TRUE(client);

    // Send a request and wait for a response.
    ContentAnalysisRequest request = BuildRequest();
    ContentAnalysisResponse response;
    int ret = client->Send(request, &response);
    ASSERT_EQ(0, ret);
    ASSERT_EQ(request.request_token(), response.request_token());

    done = true;
    });

  while (!done) {
    agent->HandleOneEventForTesting();
  }
  ASSERT_EQ(1, handler_ptr->request_count_);
  ASSERT_EQ(0, handler_ptr->ack_count_);
  ASSERT_EQ(0, handler_ptr->cancel_count_);

  client_thread.join();
}

TEST(AgentTest, Request_CloseEvent) {
  ResultCode rc;
  auto handler = std::make_unique<CloseEventTestHandler>();
  CloseEventTestHandler* handler_ptr = handler.get();
  auto agent = std::make_unique<AgentWin>(
      Agent::Config{"test", false}, std::move(handler), &rc);
  ASSERT_TRUE(agent);
  ASSERT_EQ(ResultCode::OK, rc);

  bool done = false;

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([&done]() {
    auto client = CreateClient({"test", false});
    ASSERT_TRUE(client);

    // Send a request and wait for a response.
    ContentAnalysisRequest request = BuildRequest();
    ContentAnalysisResponse response;
    int ret = client->Send(request, &response);
    ASSERT_EQ(0, ret);
    ASSERT_EQ(request.request_token(), response.request_token());

    done = true;
  });

  while (!done) {
    agent->HandleOneEventForTesting();
  }
  ASSERT_EQ(1, handler_ptr->request_count_);

  client_thread.join();
}

TEST(AgentTest, Ack) {
  TestHandler* handler_ptr;
  auto agent = CreateAgent({ "test", false }, &handler_ptr);
  ASSERT_TRUE(agent);

  bool done = false;

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([&done]() {
    auto client = CreateClient({"test", false});
    ASSERT_TRUE(client);

    // Send a request and wait for a response.
    ContentAnalysisRequest request = BuildRequest();
    ContentAnalysisResponse response;
    int ret = client->Send(request, &response);
    ASSERT_EQ(0, ret);

    ContentAnalysisAcknowledgement ack;
    ack.set_request_token(request.request_token());
    ret = client->Acknowledge(ack);
    ASSERT_EQ(0, ret);

    done = true;
  });

  while (!done) {
    agent->HandleOneEventForTesting();
  }
  ASSERT_EQ(1, handler_ptr->request_count_);
  ASSERT_EQ(1, handler_ptr->ack_count_);
  ASSERT_EQ(0, handler_ptr->cancel_count_);

  client_thread.join();
}

TEST(AgentTest, Cancel) {
  TestHandler* handler_ptr;
  auto agent = CreateAgent({ "test", false }, &handler_ptr);
  ASSERT_TRUE(agent);

  // Create a thread to handle the client.  Since the client is sync, it can't
  // run in the same thread as the agent.
  std::thread client_thread([]() {
    auto client = CreateClient({"test", false});
    ASSERT_TRUE(client);

    ContentAnalysisCancelRequests cancel;
    cancel.set_user_action_id("1234567890");
    int ret = client->CancelRequests(cancel);
    ASSERT_EQ(0, ret);
  });

  while (handler_ptr->cancel_count_ == 0) {
    agent->HandleOneEventForTesting();
  }
  ASSERT_EQ(0, handler_ptr->request_count_);
  ASSERT_EQ(0, handler_ptr->ack_count_);

  client_thread.join();
}

}  // namespace testing
}  // namespace sdk
}  // namespace content_analysis