// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include <pthread.h>
#include <sched.h>
#include <sys/socket.h>

#include <algorithm>
#include <atomic>
#include <csignal>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <memory>
#include <string>
#include <thread>
#include <vector>

#include <glog/logging.h>
#include <gtest/gtest.h>

#include "kudu/gutil/macros.h"
#include "kudu/security/tls_context.h"
#include "kudu/security/tls_handshake.h"
#include "kudu/util/countdown_latch.h"
#include "kudu/util/monotime.h"
#include "kudu/util/net/sockaddr.h"
#include "kudu/util/net/socket.h"
#include "kudu/util/random.h"
#include "kudu/util/random_util.h"
#include "kudu/util/scoped_cleanup.h"
#include "kudu/util/slice.h"
#include "kudu/util/status.h"
#include "kudu/util/test_macros.h"
#include "kudu/util/test_util.h"

using std::string;
using std::thread;
using std::unique_ptr;
using std::vector;

namespace kudu {
namespace security {

const MonoDelta kTimeout = MonoDelta::FromSeconds(30);

// Size is big enough to not fit into output socket buffer of default size
// (controlled by setsockopt() with SO_SNDBUF).
constexpr size_t kEchoChunkSize = 32 * 1024 * 1024;

class TlsSocketTest : public KuduTest {
 public:
  void SetUp() override {
    KuduTest::SetUp();
    ASSERT_OK(client_tls_.Init());
  }

 protected:
  void ConnectClient(const Sockaddr& addr, unique_ptr<Socket>* sock);
  TlsContext client_tls_;
};

Status DoNegotiationSide(Socket* sock, TlsHandshake* tls, const char* side) {
  tls->set_verification_mode(TlsVerificationMode::VERIFY_NONE);

  bool done = false;
  string received;
  while (!done) {
    string to_send;
    Status s = tls->Continue(received, &to_send);
    if (s.ok()) {
      done = true;
    } else if (!s.IsIncomplete()) {
      RETURN_NOT_OK_PREPEND(s, "unexpected tls error");
    }
    if (!to_send.empty()) {
      size_t nwritten;
      auto deadline = MonoTime::Now() + MonoDelta::FromSeconds(10);
      RETURN_NOT_OK_PREPEND(sock->BlockingWrite(
          reinterpret_cast<const uint8_t*>(to_send.data()),
          to_send.size(), &nwritten, deadline),
                            "error sending");
    }

    if (!done) {
      uint8_t buf[1024];
      int32_t n = 0;
      RETURN_NOT_OK_PREPEND(sock->Recv(buf, arraysize(buf), &n),
                            "error receiving");
      received = string(reinterpret_cast<char*>(&buf[0]), n);
    }
  }
  LOG(INFO) << side << ": negotiation complete";
  return Status::OK();
}

void TlsSocketTest::ConnectClient(const Sockaddr& addr, unique_ptr<Socket>* sock) {
  unique_ptr<Socket> client_sock(new Socket());
  ASSERT_OK(client_sock->Init(addr.family(), 0));
  ASSERT_OK(client_sock->Connect(addr));

  TlsHandshake client(TlsHandshakeType::CLIENT);
  ASSERT_OK(client_tls_.InitiateHandshake(&client));
  ASSERT_OK(DoNegotiationSide(client_sock.get(), &client, "client"));
  ASSERT_OK(client.Finish(&client_sock));
  *sock = std::move(client_sock);
}

class EchoServer {
 public:
  EchoServer()
      : pthread_sync_(1) {
  }
  ~EchoServer() {
    Stop();
    Join();
  }

  void Start() {
    ASSERT_OK(server_tls_.Init());
    ASSERT_OK(server_tls_.GenerateSelfSignedCertAndKey());
    ASSERT_OK(listen_addr_.ParseString("127.0.0.1", 0));
    ASSERT_OK(listener_.Init(listen_addr_.family(), 0));
    ASSERT_OK(listener_.BindAndListen(listen_addr_, /*listen_queue_size=*/10));
    ASSERT_OK(listener_.GetSocketAddress(&listen_addr_));

    thread_ = thread([&] {
        pthread_ = pthread_self();
        pthread_sync_.CountDown();
        unique_ptr<Socket> sock(new Socket());
        Sockaddr remote;
        CHECK_OK(listener_.Accept(sock.get(), &remote, /*flags=*/0));

        TlsHandshake server(TlsHandshakeType::SERVER);
        CHECK_OK(server_tls_.InitiateHandshake(&server));
        CHECK_OK(DoNegotiationSide(sock.get(), &server, "server"));
        CHECK_OK(server.Finish(&sock));

        CHECK_OK(sock->SetRecvTimeout(kTimeout));
        unique_ptr<uint8_t[]> buf(new uint8_t[kEchoChunkSize]);
        // An "echo" loop for kEchoChunkSize byte buffers.
        while (!is_stopped_) {
          size_t n;
          Status s = sock->BlockingRecv(buf.get(), kEchoChunkSize, &n, MonoTime::Now() + kTimeout);
          if (!s.ok()) {
            CHECK(is_stopped_) << "unexpected error reading: " << s.ToString();
          }

          LOG(INFO) << "server echoing " << n << " bytes";
          size_t written;
          s = sock->BlockingWrite(buf.get(), n, &written, MonoTime::Now() + kTimeout);
          if (!s.ok()) {
            CHECK(is_stopped_) << "unexpected error writing: " << s.ToString();
          }
          if (slow_read_) {
            SleepFor(MonoDelta::FromMilliseconds(10));
          }
        }

        CHECK_OK(listener_.Close());
      });
  }

  void EnableSlowRead() {
    slow_read_ = true;
  }

  const Sockaddr& listen_addr() const {
    return listen_addr_;
  }

  bool is_stopped() const {
    return is_stopped_;
  }

  void Stop() {
    is_stopped_ = true;
  }
  void Join() {
    thread_.join();
  }

  const pthread_t& pthread() {
    pthread_sync_.Wait();
    return pthread_;
  }

 private:
  TlsContext server_tls_;
  Socket listener_;
  Sockaddr listen_addr_;
  thread thread_;
  pthread_t pthread_;
  CountDownLatch pthread_sync_;
  std::atomic<bool> is_stopped_ { false };

  bool slow_read_ = false;
};

void handler(int /* signal */) {}

// Test that errors returned when reading from a TLS socket are reasonable and
// contain the address of the remote.
TEST_F(TlsSocketTest, TestRecvFailure) {
  EchoServer server;
  server.Start();
  unique_ptr<Socket> client_sock;
  NO_FATALS(ConnectClient(server.listen_addr(), &client_sock));
  unique_ptr<uint8_t[]> buf(new uint8_t[kEchoChunkSize]);

  // The client writes to the server. This blocks on the server receiving,
  // so once this completes the server is in its "echo loop".
  size_t nwritten_unused;
  ASSERT_OK(client_sock->BlockingWrite(buf.get(),
                                       kEchoChunkSize,
                                       &nwritten_unused,
                                       MonoTime::Now() + kTimeout));

  // Stop the server. The client can still Recv once from the server because
  // the server is trying to echo what was sent by the client.
  server.Stop();

  // The first Recv succeeds. This unblocks the server and causes it to exit.
  size_t nread_unused;
  ASSERT_OK(client_sock->BlockingRecv(buf.get(),
                                      kEchoChunkSize,
                                      &nread_unused,
                                      MonoTime::Now() + kTimeout));

  // The server has closed the connection, so the second Recv will fail.
  Status s = client_sock->BlockingRecv(buf.get(),
                                       kEchoChunkSize,
                                       &nread_unused,
                                       MonoTime::Now() + kTimeout);
  ASSERT_TRUE(s.IsNetworkError()) << s.ToString();
  ASSERT_STR_MATCHES(s.message().ToString(), "BlockingRecv error: failed to read from "
                                             "TLS socket \\(remote: 127.0.0.1:[0-9]+\\): ");
}

// Test for failures to handle EINTR during TLS connection
// negotiation and data send/receive.
TEST_F(TlsSocketTest, TestTlsSocketInterrupted) {
  // Set up a no-op signal handler for SIGUSR2.
  struct sigaction sa, sa_old;
  memset(&sa, 0, sizeof(sa));
  sa.sa_handler = &handler;
  sigaction(SIGUSR2, &sa, &sa_old);
  SCOPED_CLEANUP({ sigaction(SIGUSR2, &sa_old, nullptr); });

  EchoServer server;
  NO_FATALS(server.Start());

  // Start a thread to send signals to the server thread.
  thread killer([&]() {
      while (!server.is_stopped()) {
        PCHECK(pthread_kill(server.pthread(), SIGUSR2) == 0);
        SleepFor(MonoDelta::FromMicroseconds(rand() % 10));
      }
    });
  SCOPED_CLEANUP({ killer.join(); });

  unique_ptr<Socket> client_sock;
  NO_FATALS(ConnectClient(server.listen_addr(), &client_sock));

  unique_ptr<uint8_t[]> buf(new uint8_t[kEchoChunkSize]);
  for (int i = 0; i < 10; i++) {
    SleepFor(MonoDelta::FromMilliseconds(1));
    size_t nwritten;
    ASSERT_OK(client_sock->BlockingWrite(buf.get(), kEchoChunkSize, &nwritten,
        MonoTime::Now() + kTimeout));
    size_t n;
    ASSERT_OK(client_sock->BlockingRecv(buf.get(), kEchoChunkSize, &n,
        MonoTime::Now() + kTimeout));
  }
  server.Stop();
  ASSERT_OK(client_sock->Close());
  LOG(INFO) << "client done";
}

// Return an iovec containing the same data as the buffer 'buf' with the length 'len',
// but split into random-sized chunks. The chunks are sized randomly between 1 and
// 'max_chunk_size' bytes.
vector<struct iovec> ChunkIOVec(Random* rng, uint8_t* buf, int len, int max_chunk_size) {
  vector<struct iovec> ret;
  uint8_t* p = buf;
  int rem = len;
  while (rem > 0) {
    int len = rng->Uniform(max_chunk_size) + 1;
    len = std::min(len, rem);
    ret.push_back({p, static_cast<size_t>(len)});
    p += len;
    rem -= len;
  }
  return ret;
}

// Regression test for KUDU-2218, a bug in which Writev would improperly handle
// partial writes in non-blocking mode.
TEST_F(TlsSocketTest, TestNonBlockingWritev) {
  Random rng(GetRandomSeed32());

  EchoServer server;
  server.EnableSlowRead();
  NO_FATALS(server.Start());

  unique_ptr<Socket> client_sock;
  NO_FATALS(ConnectClient(server.listen_addr(), &client_sock));

  unique_ptr<uint8_t[]> buf(new uint8_t[kEchoChunkSize]);
  unique_ptr<uint8_t[]> rbuf(new uint8_t[kEchoChunkSize]);
  RandomString(buf.get(), kEchoChunkSize, &rng);

  for (int i = 0; i < 10; i++) {
    ASSERT_OK(client_sock->SetNonBlocking(true));

    // Prepare an IOV with the input data split into a bunch of randomly-sized
    // chunks.
    vector<struct iovec> iov = ChunkIOVec(&rng, buf.get(), kEchoChunkSize, 1024 * 1024);

    // Loop calling writev until the iov is exhausted
    int rem = kEchoChunkSize;
    while (rem > 0) {
      CHECK(!iov.empty()) << rem;
      int64_t n;
      Status s = client_sock->Writev(&iov[0], iov.size(), &n);
      if (Socket::IsTemporarySocketError(s.posix_code())) {
        sched_yield();
        continue;
      }
      ASSERT_OK(s);
      ASSERT_LE(n, rem);
      rem -= n;
      ASSERT_GE(n, 0);
      while (n > 0) {
        if (n < iov[0].iov_len) {
          iov[0].iov_len -= n;
          iov[0].iov_base = reinterpret_cast<uint8_t*>(iov[0].iov_base) + n;
          n = 0;
        } else {
          n -= iov[0].iov_len;
          iov.erase(iov.begin());
        }
      }
    }
    LOG(INFO) << "client waiting";

    size_t n;
    ASSERT_OK(client_sock->SetNonBlocking(false));
    ASSERT_OK(client_sock->BlockingRecv(rbuf.get(), kEchoChunkSize, &n,
        MonoTime::Now() + kTimeout));
    LOG(INFO) << "client got response";

    ASSERT_EQ(0, memcmp(buf.get(), rbuf.get(), kEchoChunkSize));
  }

  server.Stop();
  ASSERT_OK(client_sock->Close());
}

} // namespace security
} // namespace kudu