From d212a3bd8258bb2523280c7b081e4eb9a2194132 Mon Sep 17 00:00:00 2001
From: Hanjie Wu <wu@x-nagi.com>
Date: Sat, 16 Dec 2023 04:50:31 -0800
Subject: add retry different protocols

---
 src/DefaultBtInteractive.cc          |  3 ++
 src/InitiatorMSEHandshakeCommand.cc  | 74 ++++++++++++++++++++----------------
 src/InitiatorMSEHandshakeCommand.h   |  3 +-
 src/Peer.cc                          |  4 +-
 src/Peer.h                           | 25 +++++++++---
 src/PeerAbstractCommand.cc           |  5 +++
 src/PeerAbstractCommand.h            |  1 +
 src/PeerInitiateConnectionCommand.cc | 22 ++++++++++-
 src/PeerInitiateConnectionCommand.h  |  1 +
 src/PeerInteractionCommand.cc        | 17 +++++++++
 src/PeerInteractionCommand.h         |  1 +
 src/UTPexExtensionMessage.cc         | 24 +++++++++---
 src/uTPCore.cc                       |  8 ++--
 13 files changed, 135 insertions(+), 53 deletions(-)

diff --git a/src/DefaultBtInteractive.cc b/src/DefaultBtInteractive.cc
index da68aa95..519e53bd 100644
--- a/src/DefaultBtInteractive.cc
+++ b/src/DefaultBtInteractive.cc
@@ -171,6 +171,9 @@ DefaultBtInteractive::receiveAndSendHandshake()
 
 void DefaultBtInteractive::doPostHandshakeProcessing()
 {
+  if (peer_->isuTP() == true) {
+    peer_->uTPSupport(true);
+  }
   // Set time 0 to haveTimer to cache http/ftp download piece completion
   keepAliveTimer_ = global::wallclock();
   floodingTimer_ = global::wallclock();
diff --git a/src/InitiatorMSEHandshakeCommand.cc b/src/InitiatorMSEHandshakeCommand.cc
index ebcaa11b..9b539a04 100644
--- a/src/InitiatorMSEHandshakeCommand.cc
+++ b/src/InitiatorMSEHandshakeCommand.cc
@@ -40,6 +40,7 @@
 #include "message.h"
 #include "prefs.h"
 #include "SocketCore.h"
+#include "uTPCore.h"
 #include "Logger.h"
 #include "LogFactory.h"
 #include "Peer.h"
@@ -197,7 +198,46 @@ bool InitiatorMSEHandshakeCommand::executeInternal()
   return false;
 }
 
-void InitiatorMSEHandshakeCommand::tryNewPeer()
+bool InitiatorMSEHandshakeCommand::retryDifferentProtocol()
+{
+  if (sequence_ != INITIATOR_SEND_KEY) {
+    // Try legacy handshake when connection can be established
+    if (!(getOption()->getAsBool(PREF_BT_FORCE_ENCRYPTION) ||
+          getOption()->getAsBool(PREF_BT_REQUIRE_CRYPTO))) {
+      // Legacy handshake allowed by configuration
+      A2_LOG_INFO(fmt("CUID#%" PRId64
+                      " - Retry using legacy BitTorrent handshake.",
+                      getCuid()));
+      auto command = make_unique<PeerInitiateConnectionCommand>(
+          getCuid(), requestGroup_, getPeer(), getDownloadEngine(), btRuntime_,
+          false);
+      command->setPeerStorage(peerStorage_);
+      command->setPieceStorage(pieceStorage_);
+      getDownloadEngine()->addCommand(std::move(command));
+      return true;
+    }
+    else {
+      A2_LOG_INFO(fmt("CUID#%" PRId64 " - Establishing connection using legacy"
+                      " BitTorrent handshake is disabled by preference.",
+                      getCuid()));
+    }
+  }
+  // Try different protocol, start from using MSE handshake
+  int tried = getPeer()->triedProtocols();
+  tried |= (getPeer()->isuTP() ? PEER_UTP : PEER_TCP);
+  if (tried == PEER_ALL) {
+    return false;
+  }
+  getPeer()->triedProtocols(tried);
+  auto command = make_unique<PeerInitiateConnectionCommand>(
+      getCuid(), requestGroup_, getPeer(), getDownloadEngine(), btRuntime_);
+  command->setPeerStorage(peerStorage_);
+  command->setPieceStorage(pieceStorage_);
+  getDownloadEngine()->addCommand(std::move(command));
+  return true;
+}
+
+bool InitiatorMSEHandshakeCommand::prepareForNextPeer(time_t wait)
 {
   if (peerStorage_->isPeerAvailable() && btRuntime_->lessThanEqMinPeers()) {
     cuid_t ncuid = getDownloadEngine()->newCUID();
@@ -211,37 +251,7 @@ void InitiatorMSEHandshakeCommand::tryNewPeer()
       getDownloadEngine()->addCommand(std::move(command));
     }
   }
-}
-
-bool InitiatorMSEHandshakeCommand::prepareForNextPeer(time_t wait)
-{
-  if (sequence_ == INITIATOR_SEND_KEY) {
-    // We don't try legacy handshake when connection did not
-    // established.
-    tryNewPeer();
-    return true;
-  }
-  else if (getOption()->getAsBool(PREF_BT_FORCE_ENCRYPTION) ||
-           getOption()->getAsBool(PREF_BT_REQUIRE_CRYPTO)) {
-    A2_LOG_INFO(fmt("CUID#%" PRId64 " - Establishing connection using legacy"
-                    " BitTorrent handshake is disabled by preference.",
-                    getCuid()));
-    tryNewPeer();
-    return true;
-  }
-  else {
-    // try legacy BitTorrent handshake
-    A2_LOG_INFO(fmt("CUID#%" PRId64
-                    " - Retry using legacy BitTorrent handshake.",
-                    getCuid()));
-    auto command = make_unique<PeerInitiateConnectionCommand>(
-        getCuid(), requestGroup_, getPeer(), getDownloadEngine(), btRuntime_,
-        false);
-    command->setPeerStorage(peerStorage_);
-    command->setPieceStorage(pieceStorage_);
-    getDownloadEngine()->addCommand(std::move(command));
-    return true;
-  }
+  return true;
 }
 
 void InitiatorMSEHandshakeCommand::onAbort()
diff --git a/src/InitiatorMSEHandshakeCommand.h b/src/InitiatorMSEHandshakeCommand.h
index 91f1a065..1ec94fb6 100644
--- a/src/InitiatorMSEHandshakeCommand.h
+++ b/src/InitiatorMSEHandshakeCommand.h
@@ -72,10 +72,9 @@ private:
 
   const std::shared_ptr<Option>& getOption() const;
 
-  void tryNewPeer();
-
 protected:
   virtual bool executeInternal() CXX11_OVERRIDE;
+  virtual bool retryDifferentProtocol() CXX11_OVERRIDE;
   virtual bool prepareForNextPeer(time_t wait) CXX11_OVERRIDE;
   virtual void onAbort() CXX11_OVERRIDE;
   virtual bool exitBeforeExecute() CXX11_OVERRIDE;
diff --git a/src/Peer.cc b/src/Peer.cc
index cdc609fb..c6efae3c 100644
--- a/src/Peer.cc
+++ b/src/Peer.cc
@@ -49,10 +49,12 @@ Peer::Peer(std::string ipaddr, uint16_t port, bool incoming)
     : ipaddr_(std::move(ipaddr)),
       port_(port),
       origPort_(port),
-      supportUtp_(false),
       cuid_(0),
       firstContactTime_(global::wallclock()),
       dropStartTime_(Timer::zero()),
+      triedProtocols_(0),
+      isuTP_(false),
+      uTPSupport_(false),
       seeder_(false),
       incoming_(incoming),
       localPeer_(false),
diff --git a/src/Peer.h b/src/Peer.h
index f669ad05..7fd304ae 100644
--- a/src/Peer.h
+++ b/src/Peer.h
@@ -53,6 +53,8 @@ namespace aria2 {
 class PeerSessionResource;
 class BtMessageDispatcher;
 
+enum PeerProtocols { PEER_UTP = 1, PEER_TCP = 2, PEER_ALL = 3 };
+
 class Peer {
 private:
   std::string ipaddr_;
@@ -64,8 +66,6 @@ private:
   // is used to distinguish peer identity.
   uint16_t origPort_;
 
-  bool supportUtp_;
-
   cuid_t cuid_;
 
   unsigned char peerId_[PEER_ID_LENGTH];
@@ -74,10 +74,15 @@ private:
 
   Timer dropStartTime_;
 
-  bool seeder_;
-
   std::unique_ptr<PeerSessionResource> res_;
 
+  int triedProtocols_ = 0;
+
+  bool isuTP_;
+  bool uTPSupport_;
+
+  bool seeder_;
+
   // If true, port is assumed not to be a listening port.
   bool incoming_;
 
@@ -105,9 +110,17 @@ public:
 
   uint16_t getOrigPort() const { return origPort_; }
 
-  bool supportUtp() const { return supportUtp_; }
+  bool isuTP() const { return isuTP_; }
+
+  void isuTP(bool uTP) { isuTP_ = uTP; }
+
+  bool uTPSupport() const { return uTPSupport_; }
+
+  void uTPSupport(bool support) { uTPSupport_ = support; }
+
+  int triedProtocols() const { return triedProtocols_; }
 
-  void supportUtp(bool support) { supportUtp_ = support; }
+  void triedProtocols(int tried) { triedProtocols_ = tried; }
 
   void usedBy(cuid_t cuid);
 
diff --git a/src/PeerAbstractCommand.cc b/src/PeerAbstractCommand.cc
index 31f85441..03df3492 100644
--- a/src/PeerAbstractCommand.cc
+++ b/src/PeerAbstractCommand.cc
@@ -107,6 +107,9 @@ bool PeerAbstractCommand::execute()
     return true;
   }
   catch (RecoverableException& err) {
+    if (retryDifferentProtocol() == true) {
+      return true;
+    }
     A2_LOG_DEBUG_EX(fmt(MSG_TORRENT_DOWNLOAD_ABORTED, getCuid()), err);
     A2_LOG_DEBUG(fmt(MSG_PEER_BANNED, getCuid(), peer_->getIPAddress().c_str(),
                      peer_->getPort()));
@@ -115,6 +118,8 @@ bool PeerAbstractCommand::execute()
   }
 }
 
+bool PeerAbstractCommand::retryDifferentProtocol() { return false; }
+
 // TODO this method removed when PeerBalancerCommand is implemented
 bool PeerAbstractCommand::prepareForNextPeer(time_t wait) { return true; }
 
diff --git a/src/PeerAbstractCommand.h b/src/PeerAbstractCommand.h
index c025a961..acc305a1 100644
--- a/src/PeerAbstractCommand.h
+++ b/src/PeerAbstractCommand.h
@@ -76,6 +76,7 @@ protected:
     timeout_ = std::move(timeout);
   }
 
+  virtual bool retryDifferentProtocol();
   virtual bool prepareForNextPeer(time_t wait);
   virtual void onAbort(){};
   // This function is called when DownloadFailureException is caught right after
diff --git a/src/PeerInitiateConnectionCommand.cc b/src/PeerInitiateConnectionCommand.cc
index fbc91486..632ba861 100644
--- a/src/PeerInitiateConnectionCommand.cc
+++ b/src/PeerInitiateConnectionCommand.cc
@@ -40,6 +40,7 @@
 #include "message.h"
 #include "prefs.h"
 #include "SocketCore.h"
+#include "uTPCore.h"
 #include "Logger.h"
 #include "LogFactory.h"
 #include "Peer.h"
@@ -50,7 +51,6 @@
 #include "RequestGroup.h"
 #include "util.h"
 #include "fmt.h"
-#include "uTPCore.h"
 
 namespace aria2 {
 
@@ -77,11 +77,13 @@ bool PeerInitiateConnectionCommand::executeInternal()
 {
   A2_LOG_INFO(fmt(MSG_CONNECTING_TO_SERVER, getCuid(),
                   getPeer()->getIPAddress().c_str(), getPeer()->getPort()));
-  if (getPeer()->supportUtp() || 1) {
+  if ((getPeer()->triedProtocols() & PEER_UTP) == 0) {
     setSocket(std::make_shared<uTPCore>());
+    getPeer()->isuTP(true);
   }
   else {
     setSocket(std::make_shared<SocketCore>());
+    getPeer()->isuTP(false);
   }
   getSocket()->establishConnection(getPeer()->getIPAddress(),
                                    getPeer()->getPort(), false);
@@ -103,6 +105,22 @@ bool PeerInitiateConnectionCommand::executeInternal()
   return true;
 }
 
+bool PeerInitiateConnectionCommand::retryDifferentProtocol()
+{
+  int tried = getPeer()->triedProtocols();
+  tried |= (getPeer()->isuTP() ? PEER_UTP : PEER_TCP);
+  if (tried == PEER_ALL) {
+    return false;
+  }
+  getPeer()->triedProtocols(tried);
+  auto command = make_unique<PeerInitiateConnectionCommand>(
+      getCuid(), requestGroup_, getPeer(), getDownloadEngine(), btRuntime_);
+  command->setPeerStorage(peerStorage_);
+  command->setPieceStorage(pieceStorage_);
+  getDownloadEngine()->addCommand(std::move(command));
+  return true;
+}
+
 // TODO this method removed when PeerBalancerCommand is implemented
 bool PeerInitiateConnectionCommand::prepareForNextPeer(time_t wait)
 {
diff --git a/src/PeerInitiateConnectionCommand.h b/src/PeerInitiateConnectionCommand.h
index e1d56879..6a827743 100644
--- a/src/PeerInitiateConnectionCommand.h
+++ b/src/PeerInitiateConnectionCommand.h
@@ -58,6 +58,7 @@ private:
 
 protected:
   virtual bool executeInternal() CXX11_OVERRIDE;
+  virtual bool retryDifferentProtocol() CXX11_OVERRIDE;
   virtual bool prepareForNextPeer(time_t wait) CXX11_OVERRIDE;
   virtual void onAbort() CXX11_OVERRIDE;
   virtual bool exitBeforeExecute() CXX11_OVERRIDE;
diff --git a/src/PeerInteractionCommand.cc b/src/PeerInteractionCommand.cc
index b10fb3ce..e451a8e2 100644
--- a/src/PeerInteractionCommand.cc
+++ b/src/PeerInteractionCommand.cc
@@ -43,6 +43,7 @@
 #include "message.h"
 #include "prefs.h"
 #include "SocketCore.h"
+#include "uTPCore.h"
 #include "Option.h"
 #include "DownloadContext.h"
 #include "Peer.h"
@@ -376,6 +377,22 @@ bool PeerInteractionCommand::executeInternal()
   return false;
 }
 
+bool PeerInteractionCommand::retryDifferentProtocol()
+{
+  int tried = getPeer()->triedProtocols();
+  tried |= (getPeer()->isuTP() ? PEER_UTP : PEER_TCP);
+  if (tried == PEER_ALL) {
+    return false;
+  }
+  getPeer()->triedProtocols(tried);
+  auto command = make_unique<PeerInitiateConnectionCommand>(
+      getCuid(), requestGroup_, getPeer(), getDownloadEngine(), btRuntime_);
+  command->setPeerStorage(peerStorage_);
+  command->setPieceStorage(pieceStorage_);
+  getDownloadEngine()->addCommand(std::move(command));
+  return true;
+}
+
 // TODO this method removed when PeerBalancerCommand is implemented
 bool PeerInteractionCommand::prepareForNextPeer(time_t wait)
 {
diff --git a/src/PeerInteractionCommand.h b/src/PeerInteractionCommand.h
index 78a84cca..3c0be90d 100644
--- a/src/PeerInteractionCommand.h
+++ b/src/PeerInteractionCommand.h
@@ -72,6 +72,7 @@ private:
 
 protected:
   virtual bool executeInternal() CXX11_OVERRIDE;
+  virtual bool retryDifferentProtocol() CXX11_OVERRIDE;
   virtual bool prepareForNextPeer(time_t wait) CXX11_OVERRIDE;
   virtual void onAbort() CXX11_OVERRIDE;
   virtual void onFailure(const Exception& err) CXX11_OVERRIDE;
diff --git a/src/UTPexExtensionMessage.cc b/src/UTPexExtensionMessage.cc
index 2929ea7e..ab6a4258 100644
--- a/src/UTPexExtensionMessage.cc
+++ b/src/UTPexExtensionMessage.cc
@@ -91,6 +91,11 @@ std::string UTPexExtensionMessage::getPayload()
   return bencode2::encode(&dict);
 }
 
+namespace {
+const char PEX_FLAG_SEEDER = 0x20;
+const char PEX_FLAG_UTP = 0x40;
+} // namespace
+
 std::pair<std::pair<std::string, std::string>,
           std::pair<std::string, std::string>>
 UTPexExtensionMessage::createCompactPeerListAndFlag(
@@ -100,19 +105,26 @@ UTPexExtensionMessage::createCompactPeerListAndFlag(
   std::string flagstring;
   std::string addrstring6;
   std::string flagstring6;
-  for (auto itr = std::begin(peers), eoi = std::end(peers); itr != eoi; ++itr) {
+  std::for_each(std::begin(peers), std::end(peers), [&](auto& peer) {
     unsigned char compact[COMPACT_LEN_IPV6];
-    int compactlen = bittorrent::packcompact(compact, (*itr)->getIPAddress(),
-                                             (*itr)->getPort());
+    int compactlen =
+        bittorrent::packcompact(compact, peer->getIPAddress(), peer->getPort());
+    char flags = 0x00;
+    if (peer->isSeeder()) {
+      flags |= PEX_FLAG_SEEDER;
+    }
+    if (peer->uTPSupport()) {
+      flags |= PEX_FLAG_UTP;
+    }
     if (compactlen == COMPACT_LEN_IPV4) {
       addrstring.append(&compact[0], &compact[compactlen]);
-      flagstring += (*itr)->isSeeder() ? 0x02u : 0x00u;
+      flagstring += flags;
     }
     else if (compactlen == COMPACT_LEN_IPV6) {
       addrstring6.append(&compact[0], &compact[compactlen]);
-      flagstring6 += (*itr)->isSeeder() ? 0x02u : 0x00u;
+      flagstring6 += flags;
     }
-  }
+  });
   return std::make_pair(
       std::make_pair(std::move(addrstring), std::move(flagstring)),
       std::make_pair(std::move(addrstring6), std::move(flagstring6)));
diff --git a/src/uTPCore.cc b/src/uTPCore.cc
index 72cb32dc..2e3672e0 100644
--- a/src/uTPCore.cc
+++ b/src/uTPCore.cc
@@ -117,10 +117,10 @@ uint64 uTPCore::errorCallback(utp_callback_arguments* arg)
 
 namespace {
 void setupContextAndSock(utp_context*& ctx, std::shared_ptr<SocketCore>& sock,
-                         int family, DownloadEngine* e)
+                         int family, uint16_t port, DownloadEngine* e)
 {
   sock = std::make_shared<SocketCore>(SOCK_DGRAM);
-  sock->create(family);
+  sock->bind("", port, family);
   sock->setNonBlockingMode();
   ctx = utp_init(2);
   utp_set_callback(ctx, UTP_GET_MILLISECONDS, uTPCore::getMilliseconds);
@@ -144,8 +144,8 @@ void setupContextAndSock(utp_context*& ctx, std::shared_ptr<SocketCore>& sock,
 void uTPCore::initContext(uint16_t port, DownloadEngine* e)
 {
   start_ = global::wallclock();
-  setupContextAndSock(ctx_, sock_, AF_INET, e);
-  setupContextAndSock(ctx6_, sock6_, AF_INET6, e);
+  setupContextAndSock(ctx_, sock_, AF_INET, port, e);
+  setupContextAndSock(ctx6_, sock6_, AF_INET6, port, e);
 }
 
 void uTPCore::destroyContext()
-- 
cgit