// 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 "remoting/host/security_key/security_key_auth_handler_posix.h"

#include <stddef.h>
#include <sys/socket.h>

#include <cstdint>
#include <memory>
#include <string>

#include "base/callback_list.h"
#include "base/files/file_path.h"
#include "base/files/file_util.h"
#include "base/files/scoped_temp_dir.h"
#include "base/functional/bind.h"
#include "base/functional/callback_helpers.h"
#include "base/memory/ptr_util.h"
#include "base/run_loop.h"
#include "base/strings/string_view_util.h"
#include "base/task/sequenced_task_runner.h"
#include "base/task/thread_pool.h"
#include "base/test/run_until.h"
#include "base/test/task_environment.h"
#include "base/time/time.h"
#include "net/base/io_buffer.h"
#include "net/base/net_errors.h"
#include "net/base/sockaddr_storage.h"
#include "net/base/sockaddr_util_posix.h"
#include "net/base/test_completion_callback.h"
#include "net/socket/socket_posix.h"
#include "net/socket/unix_domain_client_socket_posix.h"
#include "net/traffic_annotation/network_traffic_annotation_test_helper.h"
#include "remoting/host/security_key/security_key_socket.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace remoting {

namespace {

const char kSocketFilename[] = "socket_for_testing";

// Test security key request data.
const uint8_t kRequestData[] = {
    0x00, 0x00, 0x00, 0x9a, 0x65, 0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00,
    0x00, 0x90, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x20, 0x60, 0x90,
    0x24, 0x71, 0xf8, 0xf2, 0xe5, 0xdf, 0x7f, 0x81, 0xc7, 0x49, 0xc4, 0xa3,
    0x58, 0x5c, 0xf6, 0xcc, 0x40, 0x14, 0x28, 0x0c, 0xa0, 0xfa, 0x03, 0x18,
    0x38, 0xd8, 0x7d, 0x77, 0x2b, 0x3a, 0x00, 0x00, 0x00, 0x20, 0x64, 0x46,
    0x47, 0x2f, 0xdf, 0x6e, 0xed, 0x7b, 0xf3, 0xc3, 0x37, 0x20, 0xf2, 0x36,
    0x67, 0x6c, 0x36, 0xe1, 0xb4, 0x5e, 0xbe, 0x04, 0x85, 0xdb, 0x89, 0xa3,
    0xcd, 0xfd, 0xd2, 0x4b, 0xd6, 0x9f, 0x00, 0x00, 0x00, 0x40, 0x38, 0x35,
    0x05, 0x75, 0x1d, 0x13, 0x6e, 0xb3, 0x6b, 0x1d, 0x29, 0xae, 0xd3, 0x43,
    0xe6, 0x84, 0x8f, 0xa3, 0x9d, 0x65, 0x4e, 0x2f, 0x57, 0xe3, 0xf6, 0xe6,
    0x20, 0x3c, 0x00, 0xc6, 0xe1, 0x73, 0x34, 0xe2, 0x23, 0x99, 0xc4, 0xfa,
    0x91, 0xc2, 0xd5, 0x97, 0xc1, 0x8b, 0xd0, 0x3c, 0x13, 0xba, 0xf0, 0xd7,
    0x5e, 0xa3, 0xbc, 0x02, 0x5b, 0xec, 0xe4, 0x4b, 0xae, 0x0e, 0xf2, 0xbd,
    0xc8, 0xaa};

const uint8_t kResponseData[] = {0x00, 0x00, 0x00, 0x01, 0x42};

const uint8_t kSshErrorData[] = {0x00, 0x00, 0x00, 0x01, 0x05};

}  // namespace

class SecurityKeyAuthHandlerPosixTest : public testing::Test {
 public:
  SecurityKeyAuthHandlerPosixTest()
      : run_loop_(new base::RunLoop()),
        expected_request_data_(
            base::as_string_view(base::span(kRequestData).subspan<4>())),
        client_response_data_(
            base::as_string_view(base::span(kResponseData).subspan<4>())) {
    EXPECT_TRUE(temp_dir_.CreateUniqueTempDir());
    socket_path_ = temp_dir_.GetPath().Append(kSocketFilename);

    send_message_callback_ = base::BindRepeating(
        &SecurityKeyAuthHandlerPosixTest::SendMessageToClient,
        base::Unretained(this));

    auto file_task_runner = base::ThreadPool::CreateSequencedTaskRunner(
        {base::MayBlock(), base::TaskPriority::BEST_EFFORT});

    auth_handler_ = remoting::SecurityKeyAuthHandlerPosix::CreateForTesting(
        socket_path_, file_task_runner);
    auth_handler_->SetSendMessageCallback(send_message_callback_, this);

    EXPECT_NE(auth_handler_.get(), nullptr);
  }

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

  void CreateSocketAndWait() {
    ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 0u);
    auth_handler_->CreateSecurityKeyConnection();

    // Wait until the socket file is created on disk, which indicates the
    // background file task has completed and the handler is listening.
    // This is 100% deterministic and avoids any manual RunLoops or banned
    // RunUntilIdle patterns.
    ASSERT_TRUE(
        base::test::RunUntil([&]() { return base::PathExists(socket_path_); }));

    ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 0u);
  }

  void SendMessageToClient(int connection_id, const std::string& data) {
    last_connection_id_received_ = connection_id;
    last_message_received_ = data;
    run_loop_->Quit();
  }

  void WaitForSendMessageToClient() {
    run_loop_->Run();
    run_loop_ = std::make_unique<base::RunLoop>();
  }

  void CheckHostDataMessage(int id) {
    ASSERT_EQ(last_connection_id_received_, id);
    ASSERT_EQ(last_message_received_.length(), expected_request_data_.length());
    ASSERT_EQ(last_message_received_, expected_request_data_);
  }

  void WriteRequestData(net::UnixDomainClientSocket* client_socket) {
    int request_len = sizeof(kRequestData);
    auto request_buffer = base::MakeRefCounted<net::DrainableIOBuffer>(
        base::MakeRefCounted<net::WrappedIOBuffer>(kRequestData), request_len);
    net::TestCompletionCallback write_callback;
    int bytes_written = 0;
    while (bytes_written < request_len) {
      int write_result = client_socket->Write(
          request_buffer.get(), request_buffer->BytesRemaining(),
          write_callback.callback(), TRAFFIC_ANNOTATION_FOR_TESTS);
      write_result = write_callback.GetResult(write_result);
      ASSERT_GT(write_result, 0);
      bytes_written += write_result;
      ASSERT_LE(bytes_written, request_len);
      request_buffer->DidConsume(write_result);
    }
    ASSERT_EQ(bytes_written, request_len);
  }

  void WaitForResponseData(net::UnixDomainClientSocket* client_socket) {
    WaitForData(client_socket, sizeof(kResponseData));
  }

  void WaitForErrorData(net::UnixDomainClientSocket* client_socket) {
    WaitForData(client_socket, sizeof(kSshErrorData));
  }

  void WaitForData(net::UnixDomainClientSocket* socket, int request_len) {
    auto buffer = base::MakeRefCounted<net::IOBufferWithSize>(request_len);
    auto read_buffer = base::MakeRefCounted<net::DrainableIOBuffer>(
        std::move(buffer), request_len);
    net::TestCompletionCallback read_callback;
    int bytes_read = 0;
    while (bytes_read < request_len) {
      int read_result =
          socket->Read(read_buffer.get(), read_buffer->BytesRemaining(),
                       read_callback.callback());
      read_result = read_callback.GetResult(read_result);
      ASSERT_GT(read_result, 0);
      bytes_read += read_result;
      ASSERT_LE(bytes_read, request_len);
      read_buffer->DidConsume(bytes_read);
    }
    ASSERT_EQ(bytes_read, request_len);
  }

  void WaitForSocketClose(net::UnixDomainClientSocket* socket) {
    auto buffer = base::MakeRefCounted<net::IOBufferWithSize>(1);
    net::TestCompletionCallback read_callback;
    int rv = socket->Read(buffer.get(), 1, read_callback.callback());
    rv = read_callback.GetResult(rv);
    ASSERT_TRUE(rv == 0 || rv == net::ERR_CONNECTION_CLOSED);
  }

 protected:
  base::test::TaskEnvironment task_environment_{
      base::test::TaskEnvironment::MainThreadType::IO};
  std::unique_ptr<base::RunLoop> run_loop_;

  // Object under test.
  std::unique_ptr<SecurityKeyAuthHandler> auth_handler_;

  SecurityKeyAuthHandler::SendMessageCallback send_message_callback_;

  int last_connection_id_received_ = -1;
  std::string last_message_received_;

  const std::string expected_request_data_;

  const std::string client_response_data_;

  base::ScopedTempDir temp_dir_;
  base::FilePath socket_path_;
};

TEST_F(SecurityKeyAuthHandlerPosixTest, HandleSingleRequest) {
  CreateSocketAndWait();

  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;

  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // Write the request and verify the response.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();
  CheckHostDataMessage(1);

  // Verify the connection is now valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));

  // Verify that completing a request/response cycle didn't close the socket.
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 1u);

  auth_handler_->SendClientResponse(1, client_response_data_);
  WaitForResponseData(&client_socket);

  // Verify that completing a request/response cycle didn't close the socket.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 1u);
}

TEST_F(SecurityKeyAuthHandlerPosixTest, HandleSingleRequestWithEof) {
  CreateSocketAndWait();

  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;

  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // Write the request.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();

  // Verify the connection is valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 1u);

  net::SocketPosix raw_socket;
  net::SockaddrStorage address;
  ASSERT_TRUE(net::FillUnixAddress(socket_path_.value(), false, &address));
  raw_socket.AdoptConnectedSocket(client_socket.ReleaseConnectedSocket(),
                                  address);

  // Close the write end of the socket.
  ASSERT_EQ(shutdown(raw_socket.socket_fd(), SHUT_WR), 0);

  // Verify that socket has not been closed yet.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 1u);

  // Wait for the response to be received.
  CheckHostDataMessage(1);

  // Verify that socket has not been closed yet.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 1u);

  auth_handler_->SendClientResponse(1, client_response_data_);

  // Verify the connection has been closed and is no longer valid.
  ASSERT_FALSE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 0u);
}

TEST_F(SecurityKeyAuthHandlerPosixTest, HandleTwoRequests) {
  CreateSocketAndWait();

  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;

  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // Write the request and verify the response.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();
  CheckHostDataMessage(1);

  // Verify the connection is now valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));

  // Send a 'client' response to the socket and verify the data is received.
  auth_handler_->SendClientResponse(1, client_response_data_);
  WaitForResponseData(&client_socket);

  // Verify the connection is still valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));

  // Repeat the request/response cycle.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();
  CheckHostDataMessage(1);

  // Verify the connection is still valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));

  auth_handler_->SendClientResponse(1, client_response_data_);
  WaitForResponseData(&client_socket);

  // Verify the connection is still valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));

  // Verify that completing two request/response cycles didn't close the
  // socket.
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 1u);
}

TEST_F(SecurityKeyAuthHandlerPosixTest, HandleTwoIndependentRequests) {
  CreateSocketAndWait();

  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;

  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // Write the request and verify the response.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();
  CheckHostDataMessage(1);

  // Verify the first connection is now valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));

  // Send a 'client' response to the socket and verify the data is received.
  auth_handler_->SendClientResponse(1, client_response_data_);
  WaitForResponseData(&client_socket);

  // Verify the connection is still valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));

  // Disconnect and establish a new connection.
  client_socket.Disconnect();

  net::TestCompletionCallback connect_callback2;
  rv = client_socket.Connect(connect_callback2.callback());
  ASSERT_EQ(connect_callback2.GetResult(rv), net::OK);

  // Repeat the request/response cycle.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();
  CheckHostDataMessage(2);

  // Verify the connection is now valid.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(2));

  // Send a 'client' response to the socket and verify the data is received.
  auth_handler_->SendClientResponse(2, client_response_data_);
  WaitForResponseData(&client_socket);

  // Verify the second connection is valid and the first is not.
  ASSERT_TRUE(auth_handler_->IsValidConnectionId(2));
  ASSERT_FALSE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 1u);
}

TEST_F(SecurityKeyAuthHandlerPosixTest, HandleReadTimeout) {
  CreateSocketAndWait();

  auth_handler_->SetRequestTimeoutForTest(base::TimeDelta());

  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;
  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // SSH Error should be received when the connection times out.
  WaitForErrorData(&client_socket);

  // Connection should no longer be valid.
  ASSERT_FALSE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 0u);
}

TEST_F(SecurityKeyAuthHandlerPosixTest, HandleClientErrorMessage) {
  CreateSocketAndWait();

  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;
  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // Write the request and verify the response.  This ensures the socket has
  // been created and is working before sending the error to tear it down.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();
  CheckHostDataMessage(1);

  // Send a 'client' response to the socket and verify the data is received.
  auth_handler_->SendClientResponse(1, client_response_data_);
  WaitForResponseData(&client_socket);

  ASSERT_TRUE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 1u);

  auth_handler_->SendErrorAndCloseConnection(1);

  // Connection should be removed immediately.
  ASSERT_FALSE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 0u);

  // SSH Error should be received.
  WaitForErrorData(&client_socket);
}

TEST_F(SecurityKeyAuthHandlerPosixTest,
       OnSecurityKeyRequest_SafeWhenCallbackNull) {
  CreateSocketAndWait();

  // 1. Establish the connection by connecting a client socket.
  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;
  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // 2. Clear the callback.
  auth_handler_->ClearSendMessageCallback(this);

  // 3. Write request data (which will trigger OnIncomingData).
  WriteRequestData(&client_socket);

  // 4. Verify that the connection was closed immediately (EOF).
  WaitForSocketClose(&client_socket);

  // 5. Verify the connection was cleaned up.
  ASSERT_FALSE(auth_handler_->IsValidConnectionId(1));
  ASSERT_EQ(auth_handler_->GetActiveConnectionCountForTest(), 0u);
}

TEST_F(SecurityKeyAuthHandlerPosixTest, CreateConnectionIsIdempotent) {
  CreateSocketAndWait();

  // Calling it again should not recreate the socket or crash.
  auth_handler_->CreateSecurityKeyConnection();

  // Verify we can still connect to the socket.
  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;
  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // Verify the connection is tracked.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();
  CheckHostDataMessage(1);
}

TEST_F(SecurityKeyAuthHandlerPosixTest,
       ClearSendMessageCallbackIgnoresMismatchedClient) {
  // Attempt to clear the callback with a mismatched client ID.
  const void* mismatched_client_id = reinterpret_cast<const void*>(0xdeadbeef);
  auth_handler_->ClearSendMessageCallback(mismatched_client_id);

  CreateSocketAndWait();

  net::UnixDomainClientSocket client_socket(socket_path_.value(), false);
  net::TestCompletionCallback connect_callback;
  int rv = client_socket.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  // Verify that the callback is STILL active and receives the message.
  WriteRequestData(&client_socket);
  WaitForSendMessageToClient();
  CheckHostDataMessage(1);

  // Now clear it with the CORRECT client ID.
  auth_handler_->ClearSendMessageCallback(this);

  // Verify that subsequent requests are dropped and connection is closed.
  net::UnixDomainClientSocket client_socket2(socket_path_.value(), false);
  rv = client_socket2.Connect(connect_callback.callback());
  ASSERT_EQ(connect_callback.GetResult(rv), net::OK);

  WriteRequestData(&client_socket2);
  WaitForSocketClose(&client_socket2);
}

}  // namespace remoting