1 files changed, 323 insertions, 0 deletions
diff --git a/src/uTPCore.cc b/src/uTPCore.cc
new file mode 100644
index 00000000..964b0883
--- /dev/null
+++ b/src/uTPCore.cc
@@ -0,0 +1,323 @@
+#include "common.h"
+
+#include <chrono>
+#include <array>
+
+#include "a2netcompat.h"
+#include "a2io.h"
+#include "a2time.h"
+#include "message.h"
+#include "DlRetryEx.h"
+#include "DlAbortEx.h"
+#include "fmt.h"
+#include "util.h"
+#include "TimeA2.h"
+#include "a2functional.h"
+#include "Logger.h"
+#include "LogFactory.h"
+#include "uTPCore.h"
+#include "Option.h"
+#include "DownloadEngine.h"
+#include "SocketCore.h"
+#include "RequestGroupMan.h"
+
+namespace aria2 {
+
+utp_context* uTPCore::ctx_;
+utp_context* uTPCore::ctx6_;
+std::shared_ptr<SocketCore> uTPCore::sock_;
+std::shared_ptr<SocketCore> uTPCore::sock6_;
+Timer uTPCore::start_;
+
+uint64 uTPCore::getMilliseconds(utp_callback_arguments* arg)
+{
+  auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
+      start_.difference());
+  return (uint64)ms.count();
+}
+
+uint64 uTPCore::getMicroseconds(utp_callback_arguments* arg)
+{
+  auto us = std::chrono::duration_cast<std::chrono::microseconds>(
+      start_.difference());
+  return (uint64)us.count();
+}
+
+uint64 uTPCore::getReadBufferSize(utp_callback_arguments* arg)
+{
+  uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
+  if (utp == nullptr) {
+    return 0;
+  }
+  return utp->end_ - utp->pos_;
+}
+
+uint64 uTPCore::writeLog(utp_callback_arguments* arg)
+{
+  A2_LOG_DEBUG(fmt("uTP: %s", arg->buf));
+  return 0;
+}
+
+uint64 uTPCore::sendTo(utp_callback_arguments* arg)
+{
+  if (arg->address->sa_family == AF_INET) {
+    sock_->writeData(arg->buf, arg->len, arg->address, arg->address_len,
+                     arg->flags);
+  }
+  else {
+    sock6_->writeData(arg->buf, arg->len, arg->address, arg->address_len,
+                      arg->flags);
+  }
+  return 0;
+}
+
+uint64 uTPCore::stateChangeCallback(utp_callback_arguments* arg)
+{
+  uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
+  if (arg->state == UTP_STATE_CONNECT || arg->state == UTP_STATE_WRITABLE) {
+    utp->isWritable_ = true;
+    if (utp->command_ != nullptr && utp->writeCheck_ == true) {
+      utp->command_->setStatusActive();
+    }
+  }
+  return 0;
+}
+
+uint64 uTPCore::readCallback(utp_callback_arguments* arg)
+{
+  uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
+  size_t bufLen = utp->buf_.size() - utp->end_;
+  if (arg->len > bufLen) {
+    utp->buf_.resize(utp->end_ + arg->len);
+  }
+  std::memcpy(utp->buf_.data() + utp->end_, arg->buf, arg->len);
+  utp->end_ += arg->len;
+  if (utp->end_ == utp->buf_.size()) {
+    utp->needDrainRead_ = true;
+  }
+  else {
+    utp_read_drained(utp->s_);
+  }
+  if (utp->command_ != nullptr && utp->readCheck_ == true) {
+    utp->command_->setStatusActive();
+  }
+  return 0;
+}
+
+uint64 uTPCore::errorCallback(utp_callback_arguments* arg)
+{
+  uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
+  utp->errNum_ = arg->error_code;
+  if (utp->command_ != nullptr) {
+    utp->command_->setStatusActive();
+  }
+  return 0;
+}
+
+namespace {
+void setupContextAndSock(utp_context*& ctx, std::shared_ptr<SocketCore>& sock,
+                         int family, DownloadEngine* e)
+{
+  sock = std::make_shared<SocketCore>(SOCK_DGRAM);
+  sock->create(family);
+  sock->setNonBlockingMode();
+  ctx = utp_init(2);
+  utp_set_callback(ctx, UTP_GET_MILLISECONDS, uTPCore::getMilliseconds);
+  utp_set_callback(ctx, UTP_GET_MICROSECONDS, uTPCore::getMicroseconds);
+  utp_set_callback(ctx, UTP_GET_READ_BUFFER_SIZE, uTPCore::getReadBufferSize);
+  if (A2_LOG_DEBUG_ENABLED) {
+    utp_context_set_option(ctx, UTP_LOG_DEBUG, 1);
+    utp_context_set_option(ctx, UTP_LOG_NORMAL, 1);
+    utp_context_set_option(ctx, UTP_LOG_MTU, 1);
+  }
+  utp_set_callback(ctx, UTP_LOG, uTPCore::writeLog);
+  utp_set_callback(ctx, UTP_SENDTO, uTPCore::sendTo);
+  utp_set_callback(ctx, UTP_ON_STATE_CHANGE, uTPCore::stateChangeCallback);
+  utp_set_callback(ctx, UTP_ON_READ, uTPCore::readCallback);
+
+  e->addRoutineCommand(
+      std::make_unique<uTPListenCommand>(e->newCUID(), e, ctx, sock));
+}
+} // namespace
+
+void uTPCore::initContext(uint16_t port, DownloadEngine* e)
+{
+  start_ = global::wallclock();
+  setupContextAndSock(ctx_, sock_, AF_INET, e);
+  setupContextAndSock(ctx6_, sock6_, AF_INET6, e);
+}
+
+void uTPCore::destroyContext()
+{
+  utp_destroy(ctx_);
+  ctx_ = nullptr;
+}
+
+uTPCore::uTPCore()
+    : s_(nullptr),
+      isWritable_(false),
+      needDrainRead_(false),
+      pos_(0),
+      end_(0),
+      errNum_(0),
+      readCheck_(false),
+      writeCheck_(false),
+      command_(nullptr),
+      SocketCore(SOCK_DGRAM)
+{
+  buf_.resize(16_k);
+}
+
+uTPCore::~uTPCore() { closeConnection(); }
+
+void uTPCore::establishConnection(const std::string& host, uint16_t port,
+                                  bool tcpNodelay)
+{
+  addrinfo* res;
+  int rv;
+  rv = callGetaddrinfo(&res, host.c_str(), util::uitos(port).c_str(),
+                       protocolFamily_, SOCK_DGRAM, 0, 0);
+  if (rv != 0) {
+    throw DL_ABORT_EX(fmt(EX_RESOLVE_HOSTNAME, host.c_str(), gai_strerror(rv)));
+  }
+  std::unique_ptr<addrinfo, decltype(&freeaddrinfo)> resDeleter(res,
+                                                                freeaddrinfo);
+  struct addrinfo* rp;
+  for (rp = res; rp; rp = rp->ai_next) {
+    s_ = utp_create_socket(rp->ai_family == AF_INET ? ctx_ : ctx6_);
+    utp_set_userdata(s_, (void*)this);
+    if (utp_connect(s_, rp->ai_addr, rp->ai_addrlen) != 0) {
+      utp_close(s_);
+      s_ = nullptr;
+      continue;
+    }
+    return;
+  }
+  throw DL_ABORT_EX(
+      fmt("Failed to create uTP connection to %s.", host.c_str()));
+}
+
+bool uTPCore::isWritable(time_t timeout) { return isWritable_; }
+
+bool uTPCore::isReadable(time_t timeout) { return pos_ < end_; }
+
+ssize_t uTPCore::writeData(const void* data, size_t len)
+{
+  a2iovec iov = {(void*)data, len};
+  return writeData(&iov, 1);
+}
+
+ssize_t uTPCore::writeVector(a2iovec* iov, size_t iovcnt)
+{
+  if (errNum_ != 0) {
+    throw DL_RETRY_EX(
+        fmt("Failed to send data, cause: %s.", utp_error_code_names[errNum_]));
+  }
+  ssize_t ret = 0;
+  wantRead_ = false;
+  wantWrite_ = false;
+  size_t len =
+      std::accumulate(iov, iov + iovcnt, (size_t)0,
+                      [](size_t s, a2iovec& elem) { return s + elem.iov_len; });
+  ret = utp_writev(s_, (utp_iovec*)iov, iovcnt);
+  if (ret == -1) {
+    throw DL_RETRY_EX("Failed to send data over uTP.");
+  }
+  if ((size_t)ret < len) {
+    wantWrite_ = true;
+    isWritable_ = false;
+  }
+  return ret;
+}
+
+void uTPCore::readData(void* data, size_t& len)
+{
+  if (errNum_ != 0) {
+    throw DL_RETRY_EX(
+        fmt("Failed to read data, cause: %s.", utp_error_code_names[errNum_]));
+  }
+  wantRead_ = false;
+  wantWrite_ = false;
+  size_t dataLen = end_ - pos_;
+  if (dataLen == 0) {
+    len = 0;
+    wantRead_ = true;
+    return;
+  }
+  if (len > dataLen) {
+    len = dataLen;
+  }
+  std::memcpy(data, buf_.data() + pos_, len);
+  pos_ += len;
+  if (pos_ > buf_.size() / 2) {
+    dataLen = end_ - pos_;
+    std::memmove(buf_.data(), buf_.data() + pos_, dataLen);
+    pos_ = 0;
+    end_ = dataLen;
+    if (needDrainRead_ == true) {
+      utp_read_drained(s_);
+      needDrainRead_ = false;
+    }
+  }
+}
+
+void uTPCore::closeConnection()
+{
+  if (s_ != nullptr) {
+    utp_set_userdata(s_, nullptr);
+    utp_close(s_);
+  }
+  s_ = nullptr;
+  command_ = nullptr;
+}
+
+uTPListenCommand::uTPListenCommand(cuid_t cuid, DownloadEngine* e,
+                                   utp_context* ctx,
+                                   std::shared_ptr<SocketCore> sock)
+    : e_(e), ctx_(ctx), sock_(sock), Command(cuid)
+{
+  e_->addSocketForReadCheck(sock_, this);
+}
+
+uTPListenCommand::~uTPListenCommand()
+{
+  if (sock_ != nullptr) {
+    e_->deleteSocketForReadCheck(sock_, this);
+  }
+};
+
+bool uTPListenCommand::execute()
+{
+  if (e_->isHaltRequested() || e_->getRequestGroupMan()->downloadFinished()) {
+    return true;
+  }
+  std::array<char, 2_k> data;
+  try {
+    while (1) {
+      sockaddr_union sockaddr;
+      socklen_t socklen = sizeof(sockaddr);
+      ssize_t length =
+          sock_->readDataFrom(data.data(), data.size(), sockaddr, socklen);
+      if (length == 0) {
+        utp_issue_deferred_acks(ctx_);
+        break;
+      }
+      utp_process_udp(ctx_, (byte*)data.data(), (size_t)length, &sockaddr.sa,
+                      socklen);
+    }
+  }
+  catch (RecoverableException& e) {
+    A2_LOG_INFO_EX("Exception thrown while receiving UDP message.", e);
+  }
+  utp_check_timeouts(ctx_);
+  e_->addRoutineCommand(std::unique_ptr<Command>(this));
+  return false;
+}
+
+void uTPCore::SetRWCheck(Command* command, int is_add, int is_read)
+{
+  command_ = command;
+  (is_read ? readCheck_ : writeCheck_) = is_add;
+}
+
+} // namespace aria2