//===------ SimpleRemoteEPCUtils.cpp - Utils for Simple Remote EPC --------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Message definitions and other utilities for SimpleRemoteEPC and
// SimpleRemoteEPCServer.
//
//===----------------------------------------------------------------------===//

#include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
#include "llvm/Config/llvm-config.h" // for LLVM_ENABLE_THREADS
#include "llvm/Support/Endian.h"

#if !defined(_MSC_VER) && !defined(__MINGW32__)
#include <unistd.h>
#else
#include <io.h>
#endif
#ifndef _WIN32
#include <sys/socket.h>
#endif

namespace {

struct FDMsgHeader {
  static constexpr unsigned MsgSizeOffset = 0;
  static constexpr unsigned OpCOffset = MsgSizeOffset + sizeof(uint64_t);
  static constexpr unsigned SeqNoOffset = OpCOffset + sizeof(uint64_t);
  static constexpr unsigned TagAddrOffset = SeqNoOffset + sizeof(uint64_t);
  static constexpr unsigned Size = TagAddrOffset + sizeof(uint64_t);
};

} // namespace

namespace llvm {
namespace orc {
namespace SimpleRemoteEPCDefaultBootstrapSymbolNames {

const char *ExecutorSessionObjectName =
    "__llvm_orc_SimpleRemoteEPC_dispatch_ctx";
const char *DispatchFnName = "__llvm_orc_SimpleRemoteEPC_dispatch_fn";

} // end namespace SimpleRemoteEPCDefaultBootstrapSymbolNames

shared::WrapperFunctionBuffer encodeHangupPayload(Error Err) {
  using SPSSerialize = shared::SPSArgList<shared::SPSError>;
  auto SE = shared::detail::toSPSSerializable(std::move(Err));
  auto Payload =
      shared::WrapperFunctionBuffer::allocate(SPSSerialize::size(SE));
  shared::SPSOutputBuffer OB(Payload.data(), Payload.size());
  bool Success = SPSSerialize::serialize(OB, SE);
  (void)Success;
  assert(Success && "Hangup payload serialization should not fail");
  return Payload;
}

Error decodeHangupPayload(shared::WrapperFunctionBuffer Payload) {
  assert(!Payload.getOutOfBandError() &&
         "Hangup payload should not be an out-of-band error buffer");

  shared::detail::SPSSerializableError Info;
  shared::SPSInputBuffer IB(Payload.data(), Payload.size());
  if (!shared::SPSArgList<shared::SPSError>::deserialize(IB, Info))
    return make_error<StringError>("Could not deserialize hangup info",
                                   inconvertibleErrorCode());
  return shared::detail::fromSPSSerializable(std::move(Info));
}

std::pair<ExecutorAddr, shared::WrapperFunctionBuffer>
encodeResultMessage(shared::WrapperFunctionBuffer ResultBytes) {
  auto Tag = [](SimpleRemoteEPCResultKind K) {
    return ExecutorAddr(static_cast<uint64_t>(K));
  };

  const char *ErrMsg = ResultBytes.getOutOfBandError();
  if (!ErrMsg)
    return {Tag(SimpleRemoteEPCResultKind::Value), std::move(ResultBytes)};

  using SPSSerialize = shared::SPSArgList<shared::SPSString>;
  StringRef M(ErrMsg);
  auto Payload = shared::WrapperFunctionBuffer::allocate(SPSSerialize::size(M));
  shared::SPSOutputBuffer OB(Payload.data(), Payload.size());
  bool Success = SPSSerialize::serialize(OB, M);
  (void)Success;
  assert(Success && "Out-of-band error serialization should not fail");
  return {Tag(SimpleRemoteEPCResultKind::OutOfBandError), std::move(Payload)};
}

Expected<shared::WrapperFunctionBuffer>
decodeResultMessage(ExecutorAddr TagAddr,
                    shared::WrapperFunctionBuffer Payload) {
  using UT = std::underlying_type_t<SimpleRemoteEPCResultKind>;
  UT KindVal = TagAddr.getValue();
  if (KindVal > static_cast<UT>(SimpleRemoteEPCResultKind::LastResultKind))
    return make_error<StringError>("Unexpected result kind " + Twine(KindVal) +
                                       " in result message",
                                   inconvertibleErrorCode());

  switch (static_cast<SimpleRemoteEPCResultKind>(KindVal)) {
  case SimpleRemoteEPCResultKind::Value:
    return std::move(Payload);
  case SimpleRemoteEPCResultKind::OutOfBandError: {
    // A malformed payload is reported as the out-of-band error itself: the
    // call waiting on this result must be unblocked either way, and an error
    // about the error is more use to the caller than a dead session.
    std::string Msg;
    shared::SPSInputBuffer IB(Payload.data(), Payload.size());
    if (!shared::SPSArgList<shared::SPSString>::deserialize(IB, Msg))
      return shared::WrapperFunctionBuffer::createOutOfBandError(
          "Could not deserialize out-of-band error message");
    return shared::WrapperFunctionBuffer::createOutOfBandError(Msg);
  }
  }
  llvm_unreachable("Invalid result kind");
}

SimpleRemoteEPCTransportClient::~SimpleRemoteEPCTransportClient() = default;
SimpleRemoteEPCTransport::~SimpleRemoteEPCTransport() = default;

Expected<std::unique_ptr<FDSimpleRemoteEPCTransport>>
FDSimpleRemoteEPCTransport::Create(SimpleRemoteEPCTransportClient &C, int InFD,
                                   int OutFD) {
#if LLVM_ENABLE_THREADS
  if (InFD == -1)
    return make_error<StringError>("Invalid input file descriptor " +
                                       Twine(InFD),
                                   inconvertibleErrorCode());
  if (OutFD == -1)
    return make_error<StringError>("Invalid output file descriptor " +
                                       Twine(OutFD),
                                   inconvertibleErrorCode());
  std::unique_ptr<FDSimpleRemoteEPCTransport> FDT(
      new FDSimpleRemoteEPCTransport(C, InFD, OutFD));
  return std::move(FDT);
#else
  return make_error<StringError>("FD-based SimpleRemoteEPC transport requires "
                                 "thread support, but llvm was built with "
                                 "LLVM_ENABLE_THREADS=Off",
                                 inconvertibleErrorCode());
#endif
}

FDSimpleRemoteEPCTransport::~FDSimpleRemoteEPCTransport() {
#if LLVM_ENABLE_THREADS
  ListenerThread.join();
#endif
}

Error FDSimpleRemoteEPCTransport::start() {
#if LLVM_ENABLE_THREADS
  ListenerThread = std::thread([this]() { listenLoop(); });
  return Error::success();
#endif
  llvm_unreachable("Should not be called with LLVM_ENABLE_THREADS=Off");
}

Error FDSimpleRemoteEPCTransport::sendMessage(SimpleRemoteEPCOpcode OpC,
                                              uint64_t SeqNo,
                                              ExecutorAddr TagAddr,
                                              ArrayRef<char> ArgBytes) {
  char HeaderBuffer[FDMsgHeader::Size];

  *((support::ulittle64_t *)(HeaderBuffer + FDMsgHeader::MsgSizeOffset)) =
      FDMsgHeader::Size + ArgBytes.size();
  *((support::ulittle64_t *)(HeaderBuffer + FDMsgHeader::OpCOffset)) =
      static_cast<uint64_t>(OpC);
  *((support::ulittle64_t *)(HeaderBuffer + FDMsgHeader::SeqNoOffset)) = SeqNo;
  *((support::ulittle64_t *)(HeaderBuffer + FDMsgHeader::TagAddrOffset)) =
      TagAddr.getValue();

  std::lock_guard<std::mutex> Lock(M);
  if (Disconnected)
    return make_error<StringError>("FD-transport disconnected",
                                   inconvertibleErrorCode());
  if (int ErrNo = writeBytes(HeaderBuffer, FDMsgHeader::Size))
    return errorCodeToError(std::error_code(ErrNo, std::generic_category()));
  if (int ErrNo = writeBytes(ArgBytes.data(), ArgBytes.size()))
    return errorCodeToError(std::error_code(ErrNo, std::generic_category()));
  return Error::success();
}

void FDSimpleRemoteEPCTransport::disconnect() {
  if (Disconnected)
    return; // Return if already disconnected.

  Disconnected = true;
  bool CloseOutFD = InFD != OutFD;

#ifndef _WIN32
  // We need to shutdown the socket to wake up (and terminate) any ongoing
  // blocking read on this FD. If the FD is not a socket, shutdown will just
  // complain through errno (instead of crashing).
  // FIXME: what about Windows?
  ::shutdown(InFD, CloseOutFD ? SHUT_RD : SHUT_RDWR);
#endif
  // Close InFD.
  while (close(InFD) == -1) {
    if (errno == EBADF)
      break;
  }

  // Close OutFD.
  if (CloseOutFD) {
#ifndef _WIN32
    // FIXME: what about Windows?
    ::shutdown(OutFD, SHUT_WR);
#endif
    while (close(OutFD) == -1) {
      if (errno == EBADF)
        break;
    }
  }
}

static Error makeUnexpectedEOFError() {
  return make_error<StringError>("Unexpected end-of-file",
                                 inconvertibleErrorCode());
}

Error FDSimpleRemoteEPCTransport::readBytes(char *Dst, size_t Size,
                                            bool *IsEOF) {
  assert((Size == 0 || Dst) && "Attempt to read into null.");
  ssize_t Completed = 0;
  while (Completed < static_cast<ssize_t>(Size)) {
    ssize_t Read = ::read(InFD, Dst + Completed, Size - Completed);
    if (Read <= 0) {
      auto ErrNo = errno;
      if (Read == 0) {
        if (Completed == 0 && IsEOF) {
          *IsEOF = true;
          return Error::success();
        } else
          return makeUnexpectedEOFError();
      } else if (ErrNo == EAGAIN || ErrNo == EINTR)
        continue;
      else {
        std::lock_guard<std::mutex> Lock(M);
        if (Disconnected && IsEOF) { // disconnect called,  pretend this is EOF.
          *IsEOF = true;
          return Error::success();
        }
        return errorCodeToError(
            std::error_code(ErrNo, std::generic_category()));
      }
    }
    Completed += Read;
  }
  return Error::success();
}

int FDSimpleRemoteEPCTransport::writeBytes(const char *Src, size_t Size) {
  assert((Size == 0 || Src) && "Attempt to append from null.");
  ssize_t Completed = 0;
  while (Completed < static_cast<ssize_t>(Size)) {
    ssize_t Written = ::write(OutFD, Src + Completed, Size - Completed);
    if (Written < 0) {
      auto ErrNo = errno;
      if (ErrNo == EAGAIN || ErrNo == EINTR)
        continue;
      else
        return ErrNo;
    }
    Completed += Written;
  }
  return 0;
}

void FDSimpleRemoteEPCTransport::listenLoop() {
  Error Err = Error::success();
  do {

    char HeaderBuffer[FDMsgHeader::Size];
    // Read the header buffer.
    {
      bool IsEOF = false;
      if (auto Err2 = readBytes(HeaderBuffer, FDMsgHeader::Size, &IsEOF)) {
        Err = joinErrors(std::move(Err), std::move(Err2));
        break;
      }
      if (IsEOF)
        break;
    }

    // Decode header buffer.
    uint64_t MsgSize;
    SimpleRemoteEPCOpcode OpC;
    uint64_t SeqNo;
    ExecutorAddr TagAddr;

    MsgSize =
        *((support::ulittle64_t *)(HeaderBuffer + FDMsgHeader::MsgSizeOffset));
    OpC = static_cast<SimpleRemoteEPCOpcode>(static_cast<uint64_t>(
        *((support::ulittle64_t *)(HeaderBuffer + FDMsgHeader::OpCOffset))));
    SeqNo =
        *((support::ulittle64_t *)(HeaderBuffer + FDMsgHeader::SeqNoOffset));
    TagAddr.setValue(
        *((support::ulittle64_t *)(HeaderBuffer + FDMsgHeader::TagAddrOffset)));

    if (MsgSize < FDMsgHeader::Size) {
      Err = joinErrors(std::move(Err),
                       make_error<StringError>("Message size too small",
                                               inconvertibleErrorCode()));
      break;
    }

    // Read the argument bytes.
    auto ArgBytes =
        shared::WrapperFunctionBuffer::allocate(MsgSize - FDMsgHeader::Size);
    if (auto Err2 = readBytes(ArgBytes.data(), ArgBytes.size())) {
      Err = joinErrors(std::move(Err), std::move(Err2));
      break;
    }

    if (auto Action =
            C.handleMessage(OpC, SeqNo, TagAddr, std::move(ArgBytes))) {
      if (*Action == SimpleRemoteEPCTransportClient::EndSession)
        break;
    } else {
      Err = joinErrors(std::move(Err), Action.takeError());
      break;
    }
  } while (true);

  // Attempt to close FDs, set Disconnected to true so that subsequent
  // sendMessage calls fail.
  disconnect();

  // Call up to the client to handle the disconnection.
  C.handleDisconnect(std::move(Err));
}

} // end namespace orc
} // end namespace llvm