1 files changed, 34 insertions, 13 deletions

diff --git a/src/PeerAbstractCommand.cc b/src/PeerAbstractCommand.cc
index 6ab8aa77..65995c98 100644
--- a/src/PeerAbstractCommand.cc
+++ b/src/PeerAbstractCommand.cc
@@ -38,6 +38,7 @@
 #include "Option.h"
 #include "DlAbortEx.h"
 #include "SocketCore.h"
+#include "uTPCore.h"
 #include "Logger.h"
 #include "LogFactory.h"
 #include "message.h"
@@ -106,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()));
@@ -114,13 +118,15 @@ bool PeerAbstractCommand::execute()
   }
 }
 
+bool PeerAbstractCommand::retryDifferentProtocol() { return false; }
+
 // TODO this method removed when PeerBalancerCommand is implemented
 bool PeerAbstractCommand::prepareForNextPeer(time_t wait) { return true; }
 
 void PeerAbstractCommand::disableReadCheckSocket()
 {
   if (checkSocketIsReadable_) {
-    e_->deleteSocketForReadCheck(readCheckTarget_, this);
+    setRWCheck(e_, readCheckTarget_, false, true);
     checkSocketIsReadable_ = false;
     readCheckTarget_.reset();
   }
@@ -135,13 +141,13 @@ void PeerAbstractCommand::setReadCheckSocket(
   else {
     if (checkSocketIsReadable_) {
       if (*readCheckTarget_ != *socket) {
-        e_->deleteSocketForReadCheck(readCheckTarget_, this);
-        e_->addSocketForReadCheck(socket, this);
+        setRWCheck(e_, readCheckTarget_, false, true);
+        setRWCheck(e_, socket, true, true);
         readCheckTarget_ = socket;
       }
     }
     else {
-      e_->addSocketForReadCheck(socket, this);
+      setRWCheck(e_, socket, true, true);
       checkSocketIsReadable_ = true;
       readCheckTarget_ = socket;
     }
@@ -151,7 +157,7 @@ void PeerAbstractCommand::setReadCheckSocket(
 void PeerAbstractCommand::disableWriteCheckSocket()
 {
   if (checkSocketIsWritable_) {
-    e_->deleteSocketForWriteCheck(writeCheckTarget_, this);
+    setRWCheck(e_, writeCheckTarget_, false, false);
     checkSocketIsWritable_ = false;
     writeCheckTarget_.reset();
   }
@@ -166,19 +172,39 @@ void PeerAbstractCommand::setWriteCheckSocket(
   else {
     if (checkSocketIsWritable_) {
       if (*writeCheckTarget_ != *socket) {
-        e_->deleteSocketForWriteCheck(writeCheckTarget_, this);
-        e_->addSocketForWriteCheck(socket, this);
+        setRWCheck(e_, writeCheckTarget_, false, false);
+        setRWCheck(e_, socket, true, false);
         writeCheckTarget_ = socket;
       }
     }
     else {
-      e_->addSocketForWriteCheck(socket, this);
+      setRWCheck(e_, socket, true, false);
       checkSocketIsWritable_ = true;
       writeCheckTarget_ = socket;
     }
   }
 }
 
+void PeerAbstractCommand::setRWCheck(DownloadEngine* e,
+                                     const std::shared_ptr<SocketCore>& socket,
+                                     int is_add, int is_read)
+{
+  auto utp = std::dynamic_pointer_cast<uTPCore>(socket);
+  if (utp == nullptr) {
+    if (is_add) {
+      is_read ? e_->addSocketForReadCheck(socket, this)
+              : e_->addSocketForWriteCheck(socket, this);
+    }
+    else {
+      is_read ? e_->deleteSocketForReadCheck(socket, this)
+              : e_->deleteSocketForWriteCheck(socket, this);
+    }
+  }
+  else {
+    utp->setRWCheck(this, is_add, is_read);
+  }
+}
+
 void PeerAbstractCommand::setNoCheck(bool check) { noCheck_ = check; }
 
 void PeerAbstractCommand::updateKeepAlive()
@@ -186,11 +212,6 @@ void PeerAbstractCommand::updateKeepAlive()
   checkPoint_ = global::wallclock();
 }
 
-void PeerAbstractCommand::createSocket()
-{
-  socket_ = std::make_shared<SocketCore>();
-}
-
 void PeerAbstractCommand::addCommandSelf()
 {
   e_->addCommand(std::unique_ptr<Command>(this));