add uTP accept callback

2 files changed, 44 insertions, 15 deletions

diff --git a/src/uTPCore.cc b/src/uTPCore.cc
index 2e3672e0..a279ae56 100644
--- a/src/uTPCore.cc
+++ b/src/uTPCore.cc
@@ -13,14 +13,15 @@
 #include "fmt.h"
 #include "util.h"
 #include "TimeA2.h"
-#include "a2functional.h"
+#include "Peer.h"
 #include "Logger.h"
 #include "LogFactory.h"
-#include "uTPCore.h"
 #include "Option.h"
 #include "DownloadEngine.h"
 #include "SocketCore.h"
+#include "uTPCore.h"
 #include "RequestGroupMan.h"
+#include "ReceiverMSEHandshakeCommand.h"
 
 namespace aria2 {
 
@@ -53,13 +54,13 @@ uint64 uTPCore::getReadBufferSize(utp_callback_arguments* arg)
   return utp->end_ - utp->pos_;
 }
 
-uint64 uTPCore::writeLog(utp_callback_arguments* arg)
+uint64 uTPCore::writeLogCallback(utp_callback_arguments* arg)
 {
   A2_LOG_DEBUG(fmt("uTP: %s", arg->buf));
   return 0;
 }
 
-uint64 uTPCore::sendTo(utp_callback_arguments* arg)
+uint64 uTPCore::sendtoCallback(utp_callback_arguments* arg)
 {
   if (arg->address->sa_family == AF_INET) {
     sock_->writeData(arg->buf, arg->len, arg->address, arg->address_len,
@@ -72,6 +73,24 @@ uint64 uTPCore::sendTo(utp_callback_arguments* arg)
   return 0;
 }
 
+uint64 uTPCore::acceptCallback(utp_callback_arguments* arg)
+{
+  std::shared_ptr<uTPCore> utp = std::make_shared<uTPCore>(arg->socket);
+  utp_set_userdata(arg->socket, utp.get());
+  utp->isWritable_ = true;
+  Endpoint endpoint = util::getNumericNameInfo(arg->address, arg->address_len);
+  DownloadEngine* e = (DownloadEngine*)utp_context_get_userdata(arg->context);
+  std::shared_ptr<Peer> peer =
+      std::make_shared<Peer>(endpoint.addr, endpoint.port, true);
+  cuid_t cuid = e->newCUID();
+  e->addCommand(make_unique<ReceiverMSEHandshakeCommand>(cuid, peer, e, utp));
+  A2_LOG_DEBUG(fmt("Accepted the connection from %s:%u.",
+                   peer->getIPAddress().c_str(), peer->getPort()));
+  A2_LOG_DEBUG(
+      fmt("Added CUID#%" PRId64 " to receive BitTorrent/MSE handshake.", cuid));
+  return 0;
+}
+
 uint64 uTPCore::stateChangeCallback(utp_callback_arguments* arg)
 {
   uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
@@ -123,6 +142,7 @@ void setupContextAndSock(utp_context*& ctx, std::shared_ptr<SocketCore>& sock,
   sock->bind("", port, family);
   sock->setNonBlockingMode();
   ctx = utp_init(2);
+  utp_context_set_userdata(ctx, (void*)e);
   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);
@@ -131,8 +151,9 @@ void setupContextAndSock(utp_context*& ctx, std::shared_ptr<SocketCore>& sock,
     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_LOG, uTPCore::writeLogCallback);
+  utp_set_callback(ctx, UTP_SENDTO, uTPCore::sendtoCallback);
+  utp_set_callback(ctx, UTP_ON_ACCEPT, uTPCore::acceptCallback);
   utp_set_callback(ctx, UTP_ON_STATE_CHANGE, uTPCore::stateChangeCallback);
   utp_set_callback(ctx, UTP_ON_READ, uTPCore::readCallback);
 
@@ -145,17 +166,25 @@ void uTPCore::initContext(uint16_t port, DownloadEngine* e)
 {
   start_ = global::wallclock();
   setupContextAndSock(ctx_, sock_, AF_INET, port, e);
-  setupContextAndSock(ctx6_, sock6_, AF_INET6, port, e);
+  if (e->getOption()->getAsBool(PREF_DISABLE_IPV6) == false) {
+    setupContextAndSock(ctx6_, sock6_, AF_INET6, port, e);
+  }
 }
 
 void uTPCore::destroyContext()
 {
-  utp_destroy(ctx_);
-  ctx_ = nullptr;
+  if (ctx_ != nullptr) {
+    utp_destroy(ctx_);
+    ctx_ = nullptr;
+  }
+  if (ctx6_ != nullptr) {
+    utp_destroy(ctx6_);
+    ctx6_ = nullptr;
+  }
 }
 
-uTPCore::uTPCore()
-    : s_(nullptr),
+uTPCore::uTPCore(utp_socket* s)
+    : s_(s),
       isWritable_(false),
       needDrainRead_(false),
       pos_(0),
diff --git a/src/uTPCore.h b/src/uTPCore.h
index afca8047..afc13fd6 100644
--- a/src/uTPCore.h
+++ b/src/uTPCore.h
@@ -36,7 +36,7 @@ private:
   static Timer start_;
 
 public:
-  uTPCore();
+  uTPCore(utp_socket* s = nullptr);
   virtual ~uTPCore();
 
   static void initContext(uint16_t port, DownloadEngine* e);
@@ -57,9 +57,9 @@ public:
   static uint64 getMicroseconds(utp_callback_arguments* arg);
   static uint64 getReadBufferSize(utp_callback_arguments* arg);
 
-  static uint64 writeLog(utp_callback_arguments* arg);
-  static uint64 sendTo(utp_callback_arguments* arg);
-
+  static uint64 writeLogCallback(utp_callback_arguments* arg);
+  static uint64 sendtoCallback(utp_callback_arguments* arg);
+  static uint64 acceptCallback(utp_callback_arguments* arg);
   static uint64 stateChangeCallback(utp_callback_arguments* arg);
   static uint64 readCallback(utp_callback_arguments* arg);
   static uint64 errorCallback(utp_callback_arguments* arg);