// Copyright 2021 The gVisor Authors.
//
// Licensed 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 <err.h>
#include <pthread.h>
#include <stdarg.h>
#include <stdio.h>
#include <sys/epoll.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

#include <array>
#include <cerrno>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <string>
#include <vector>

#include "absl/cleanup/cleanup.h"
#include "absl/strings/str_replace.h"
#include "absl/strings/string_view.h"
#include "pkg/sentry/seccheck/points/common.pb.h"
#include "pkg/sentry/seccheck/points/container.pb.h"
#include "pkg/sentry/seccheck/points/sentry.pb.h"
#include "pkg/sentry/seccheck/points/syscall.pb.h"
#include "google/protobuf/text_format.h"

namespace {

typedef std::function<void(absl::string_view buf)> Callback;

constexpr size_t kMaxEventSize = 300 * 1024;

bool quiet = false;

#pragma pack(push, 1)
struct Header {
  uint16_t header_size;
  uint16_t message_type;
  uint32_t dropped_count;
};
#pragma pack(pop)

void Log(const char* fmt, ...) {
  if (!quiet) {
    va_list ap;
    va_start(ap, fmt);
    vprintf(fmt, ap);
    va_end(ap);
  }
}

template <class T>
std::string ShortFmt(T msg) {
  std::string short_text_msg;
  google::protobuf::TextFormat::PrintToString(msg, &short_text_msg);
  return absl::StrReplaceAll(short_text_msg,
                             {{"\r\n", " "}, {"\n", " "}, {"\r", " "}});
}

template <class T>
void UnpackSyscall(absl::string_view buf) {
  T evt;
  if (!evt.ParseFromString(buf)) {
    err(1, "ParseFromString(): %.*s", static_cast<int>(buf.size()), buf.data());
  }
  absl::string_view name = evt.GetDescriptor()->name();
  Log("%s %.*s %s\n", evt.has_exit() ? "X" : "E", static_cast<int>(name.size()),
      name.data(), ShortFmt(evt).c_str());
}

template <class T>
void Unpack(absl::string_view buf) {
  T evt;
  if (!evt.ParseFromString(buf)) {
    err(1, "ParseFromString(): %.*s", static_cast<int>(buf.size()), buf.data());
  }
  absl::string_view name = evt.GetDescriptor()->name();
  Log("%.*s => %s\n", static_cast<int>(name.size()), name.data(),
      ShortFmt(evt).c_str());
}

// List of dispatchers indexed based on MessageType enum values.
// LINT.IfChange
const std::vector<Callback> dispatchers = [] {
  std::vector<Callback> result(::gvisor::common::MessageType_MAX + 1, nullptr);
  result[::gvisor::common::MESSAGE_CONTAINER_START] =
      Unpack<::gvisor::container::Start>;
  result[::gvisor::common::MESSAGE_SENTRY_CLONE] =
      Unpack<::gvisor::sentry::CloneInfo>;
  result[::gvisor::common::MESSAGE_SENTRY_EXEC] =
      Unpack<::gvisor::sentry::ExecveInfo>;
  result[::gvisor::common::MESSAGE_SENTRY_EXIT_NOTIFY_PARENT] =
      Unpack<::gvisor::sentry::ExitNotifyParentInfo>;
  result[::gvisor::common::MESSAGE_SENTRY_TASK_EXIT] =
      Unpack<::gvisor::sentry::TaskExit>;
  result[::gvisor::common::MESSAGE_SYSCALL_RAW] =
      UnpackSyscall<::gvisor::syscall::Syscall>;
  result[::gvisor::common::MESSAGE_SYSCALL_OPEN] =
      UnpackSyscall<::gvisor::syscall::Open>;
  result[::gvisor::common::MESSAGE_SYSCALL_CLOSE] =
      UnpackSyscall<::gvisor::syscall::Close>;
  result[::gvisor::common::MESSAGE_SYSCALL_READ] =
      UnpackSyscall<::gvisor::syscall::Read>;
  result[::gvisor::common::MESSAGE_SYSCALL_CONNECT] =
      UnpackSyscall<::gvisor::syscall::Connect>;
  result[::gvisor::common::MESSAGE_SYSCALL_EXECVE] =
      UnpackSyscall<::gvisor::syscall::Execve>;
  result[::gvisor::common::MESSAGE_SYSCALL_SOCKET] =
      UnpackSyscall<::gvisor::syscall::Socket>;
  result[::gvisor::common::MESSAGE_SYSCALL_CHDIR] =
      UnpackSyscall<::gvisor::syscall::Chdir>;
  result[::gvisor::common::MESSAGE_SYSCALL_SETID] =
      UnpackSyscall<::gvisor::syscall::Setid>;
  result[::gvisor::common::MESSAGE_SYSCALL_SETRESID] =
      UnpackSyscall<::gvisor::syscall::Setresid>;
  result[::gvisor::common::MESSAGE_SYSCALL_DUP] =
      UnpackSyscall<::gvisor::syscall::Dup>;
  result[::gvisor::common::MESSAGE_SYSCALL_PRLIMIT64] =
      UnpackSyscall<::gvisor::syscall::Prlimit>;
  result[::gvisor::common::MESSAGE_SYSCALL_PIPE] =
      UnpackSyscall<::gvisor::syscall::Pipe>;
  result[::gvisor::common::MESSAGE_SYSCALL_FCNTL] =
      UnpackSyscall<::gvisor::syscall::Fcntl>;
  result[::gvisor::common::MESSAGE_SYSCALL_SIGNALFD] =
      UnpackSyscall<::gvisor::syscall::Signalfd>;
  result[::gvisor::common::MESSAGE_SYSCALL_EVENTFD] =
      UnpackSyscall<::gvisor::syscall::Eventfd>;
  result[::gvisor::common::MESSAGE_SYSCALL_CHROOT] =
      UnpackSyscall<::gvisor::syscall::Chroot>;
  result[::gvisor::common::MESSAGE_SYSCALL_CLONE] =
      UnpackSyscall<::gvisor::syscall::Clone>;
  result[::gvisor::common::MESSAGE_SYSCALL_BIND] =
      UnpackSyscall<::gvisor::syscall::Bind>;
  result[::gvisor::common::MESSAGE_SYSCALL_ACCEPT] =
      UnpackSyscall<::gvisor::syscall::Accept>;
  result[::gvisor::common::MESSAGE_SYSCALL_TIMERFD_CREATE] =
      UnpackSyscall<::gvisor::syscall::TimerfdCreate>;
  result[::gvisor::common::MESSAGE_SYSCALL_TIMERFD_SETTIME] =
      UnpackSyscall<::gvisor::syscall::TimerfdSetTime>;
  result[::gvisor::common::MESSAGE_SYSCALL_TIMERFD_GETTIME] =
      UnpackSyscall<::gvisor::syscall::TimerfdGetTime>;
  result[::gvisor::common::MESSAGE_SYSCALL_FORK] =
      UnpackSyscall<::gvisor::syscall::Fork>;
  result[::gvisor::common::MESSAGE_SYSCALL_INOTIFY_INIT] =
      UnpackSyscall<::gvisor::syscall::InotifyInit>;
  result[::gvisor::common::MESSAGE_SYSCALL_INOTIFY_ADD_WATCH] =
      UnpackSyscall<::gvisor::syscall::InotifyAddWatch>;
  result[::gvisor::common::MESSAGE_SYSCALL_INOTIFY_RM_WATCH] =
      UnpackSyscall<::gvisor::syscall::InotifyRmWatch>;
  result[::gvisor::common::MESSAGE_SYSCALL_SOCKETPAIR] =
      UnpackSyscall<::gvisor::syscall::SocketPair>;
  result[::gvisor::common::MESSAGE_SYSCALL_WRITE] =
      UnpackSyscall<::gvisor::syscall::Write>;
  result[::gvisor::common::MESSAGE_SENTRY_MMAP] =
      Unpack<::gvisor::sentry::MmapInfo>;
  result[::gvisor::common::MESSAGE_SYSCALL_MMAP] =
      UnpackSyscall<::gvisor::syscall::Mmap>;
  result[::gvisor::common::MESSAGE_SYSCALL_LISTEN] =
      UnpackSyscall<::gvisor::syscall::Listen>;
  result[::gvisor::common::MESSAGE_SYSCALL_PTRACE] =
      UnpackSyscall<::gvisor::syscall::Ptrace>;
  return result;
}();
// LINT.ThenChange(../../pkg/sentry/seccheck/points/common.proto)

void DispatchMessage(absl::string_view buf) {
  const Header* hdr = reinterpret_cast<const Header*>(&buf[0]);

  // Payload size can be zero when proto object contains only defaults values.
  size_t payload_size = buf.size() - hdr->header_size;
  if (payload_size < 0) {
    printf("Header size (%u) is larger than message %lu\n", hdr->header_size,
           buf.size());
    return;
  }

  auto proto = buf.substr(hdr->header_size);
  if (proto.size() < payload_size) {
    printf("Message was truncated, size: %lu, expected: %zu\n", proto.size(),
           payload_size);
    return;
  }

  if (hdr->message_type == 0 || hdr->message_type >= dispatchers.size()) {
    printf("Invalid message type: %u\n", hdr->message_type);
    return;
  }
  Callback cb = dispatchers[hdr->message_type];
  if (cb) {
    cb(proto);
  } else {
    printf("No dispatcher configured for message type: %u\n",
           hdr->message_type);
  }
}

bool ReadAndUnpack(int client) {
  std::array<char, kMaxEventSize> buf;
  int bytes = read(client, buf.data(), buf.size());
  if (bytes < 0) {
    err(1, "read");
  }
  if (bytes == 0) {
    return false;
  }
  DispatchMessage(absl::string_view(buf.data(), bytes));
  return true;
}

void* PollLoop(void* ptr) {
  const int poll_fd = static_cast<int>(reinterpret_cast<intptr_t>(ptr));
  for (;;) {
    epoll_event evts[64];
    int nfds = epoll_wait(poll_fd, evts, 64, -1);
    if (nfds < 0) {
      if (errno == EINTR) {
        continue;
      }
      err(1, "epoll_wait");
    }

    for (int i = 0; i < nfds; ++i) {
      if (evts[i].events & EPOLLIN) {
        int client = evts[i].data.fd;
        ReadAndUnpack(client);
      }
      if ((evts[i].events & (EPOLLRDHUP | EPOLLHUP)) != 0) {
        int client = evts[i].data.fd;
        // Drain any remaining messages before closing the socket.
        while (ReadAndUnpack(client)) {
        }
        close(client);
        printf("Connection closed\n");
      }
      if (evts[i].events & EPOLLERR) {
        printf("error\n");
      }
    }
  }
}

void StartPollThread(int poll_fd) {
  pthread_t thread;
  if (pthread_create(&thread, nullptr, PollLoop,
                     reinterpret_cast<void*>(poll_fd)) != 0) {
    err(1, "pthread_create");
  }
  pthread_detach(thread);
}

// Handshake performs version exchange with client. See common.proto for details
// about the protocol.
bool Handshake(int client_fd) {
  std::vector<char> buf(10240);
  int bytes = read(client_fd, buf.data(), buf.size());
  if (bytes < 0) {
    printf("Error receiving Handshake message: %d\n", errno);
    return false;
  } else if (bytes == (int)buf.size()) {
    // Protect against the Handshake becoming larger than the buffer allocated
    // for it.
    printf("Handshake message too big\n");
    return false;
  }
  ::gvisor::common::Handshake in = {};
  if (!in.ParseFromString(absl::string_view(buf.data(), bytes))) {
    printf("Error parsing Handshake message\n");
    return false;
  }

  constexpr uint32_t minSupportedVersion = 1;
  if (in.version() < minSupportedVersion) {
    printf("Client has unsupported version %u\n", in.version());
    return false;
  }

  ::gvisor::common::Handshake out;
  out.set_version(1);
  if (!out.SerializeToFileDescriptor(client_fd)) {
    printf("Error sending Handshake message: %d\n", errno);
    return false;
  }
  return true;
}
}  // namespace

int main(int argc, char** argv) {
  for (int c = 0; (c = getopt(argc, argv, "q")) != -1;) {
    switch (c) {
      case 'q':
        quiet = true;
        break;
      default:
        exit(1);
    }
  }

  if (!quiet) {
    setbuf(stdout, NULL);
    setbuf(stderr, NULL);
  }
  std::string path("/tmp/gvisor_events.sock");
  if (optind < argc) {
    path = argv[optind];
  }
  if (path.empty()) {
    err(1, "empty file name");
  }
  printf("Socket address %s\n", path.c_str());
  unlink(path.c_str());

  int sock = socket(AF_UNIX, SOCK_SEQPACKET, 0);
  if (sock < 0) {
    err(1, "socket");
  }
  auto sock_closer = absl::MakeCleanup([sock] { close(sock); });

  struct sockaddr_un addr;
  addr.sun_family = AF_UNIX;
  strncpy(addr.sun_path, path.c_str(), path.size() + 1);
  if (bind(sock, reinterpret_cast<struct sockaddr*>(&addr), sizeof(addr))) {
    err(1, "bind");
  }
  if (listen(sock, 5) < 0) {
    err(1, "listen");
  }

  int epoll_fd = epoll_create(1);
  if (epoll_fd < 0) {
    err(1, "epoll_create");
  }
  auto epoll_closer = absl::MakeCleanup([epoll_fd] { close(epoll_fd); });
  StartPollThread(epoll_fd);

  for (;;) {
    int client = accept(sock, nullptr, nullptr);
    if (client < 0) {
      if (errno == EINTR) {
        continue;
      }
      err(1, "accept");
    }
    printf("Connection accepted\n");

    if (!Handshake(client)) {
      close(client);
      continue;
    }

    struct epoll_event evt;
    evt.data.fd = client;
    evt.events = EPOLLIN;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client, &evt) < 0) {
      err(1, "epoll_ctl(ADD)");
    }
  }
}
