Support UDP tracker

It shares UDP listening port with IPv4 DHT. At the moment, in order to
enable UDP tracker support, enable IPv4 DHT.

1 files changed, 64 insertions, 9 deletions

diff --git a/src/DHTInteractionCommand.cc b/src/DHTInteractionCommand.cc
index 4fc3f371..096812a8 100644
--- a/src/DHTInteractionCommand.cc
+++ b/src/DHTInteractionCommand.cc
@@ -46,10 +46,16 @@
 #include "LogFactory.h"
 #include "DHTMessageCallback.h"
 #include "DHTNode.h"
+#include "DHTConnection.h"
+#include "UDPTrackerClient.h"
+#include "UDPTrackerRequest.h"
 #include "fmt.h"
+#include "wallclock.h"
 
 namespace aria2 {
 
+// TODO This name of this command is misleading, because now it also
+// handles UDP trackers as well as DHT.
 DHTInteractionCommand::DHTInteractionCommand(cuid_t cuid, DownloadEngine* e)
   : Command(cuid),
     e_(e)
@@ -77,23 +83,60 @@ void DHTInteractionCommand::disableReadCheckSocket(const SharedHandle<SocketCore
 
 bool DHTInteractionCommand::execute()
 {
-  if(e_->getRequestGroupMan()->downloadFinished() || e_->isHaltRequested()) {
+  // We need to keep this command alive while TrackerWatcherCommand
+  // needs this.
+  if(e_->getRequestGroupMan()->downloadFinished() ||
+     (e_->isHaltRequested() && udpTrackerClient_->getNumWatchers() == 0)) {
+    return true;
+  } else if(e_->isForceHaltRequested()) {
+    udpTrackerClient_->failAll();
     return true;
   }
 
   taskQueue_->executeTask();
 
-  while(1) {
-    SharedHandle<DHTMessage> m = receiver_->receiveMessage();
-    if(!m) {
-      break;
+  std::string remoteAddr;
+  uint16_t remotePort;
+  unsigned char data[64*1024];
+  try {
+    while(1) {
+      ssize_t length = connection_->receiveMessage(data, sizeof(data),
+                                                   remoteAddr, remotePort);
+      if(length <= 0) {
+        break;
+      }
+      if(data[0] == 'd') {
+        // udp tracker response does not start with 'd', so assume
+        // this message belongs to DHT. nothrow.
+        receiver_->receiveMessage(remoteAddr, remotePort, data, length);
+      } else {
+        // this may be udp tracker response. nothrow.
+        udpTrackerClient_->receiveReply(data, length, remoteAddr, remotePort,
+                                        global::wallclock());
+      }
     }
+  } catch(RecoverableException& e) {
+    A2_LOG_INFO_EX("Exception thrown while receiving UDP message.", e);
   }
   receiver_->handleTimeout();
-  try {
-    dispatcher_->sendMessages();
-  } catch(RecoverableException& e) {
-    A2_LOG_ERROR_EX(EX_EXCEPTION_CAUGHT, e);
+  udpTrackerClient_->handleTimeout(global::wallclock());
+  dispatcher_->sendMessages();
+  while(!udpTrackerClient_->getPendingRequests().empty()) {
+    // no throw
+    ssize_t length = udpTrackerClient_->createRequest(data, sizeof(data),
+                                                      remoteAddr, remotePort,
+                                                      global::wallclock());
+    if(length == -1) {
+      break;
+    }
+    try {
+      // throw
+      connection_->sendMessage(data, length, remoteAddr, remotePort);
+      udpTrackerClient_->requestSent(global::wallclock());
+    } catch(RecoverableException& e) {
+      A2_LOG_INFO_EX("Exception thrown while sending UDP tracker request.", e);
+      udpTrackerClient_->requestFail(UDPT_ERR_NETWORK);
+    }
   }
   e_->addCommand(this);
   return false;
@@ -114,4 +157,16 @@ void DHTInteractionCommand::setTaskQueue(const SharedHandle<DHTTaskQueue>& taskQ
   taskQueue_ = taskQueue;
 }
 
+void DHTInteractionCommand::setConnection
+(const SharedHandle<DHTConnection>& connection)
+{
+  connection_ = connection;
+}
+
+void DHTInteractionCommand::setUDPTrackerClient
+(const SharedHandle<UDPTrackerClient>& udpTrackerClient)
+{
+  udpTrackerClient_ = udpTrackerClient;
+}
+
 } // namespace aria2