1 files changed, 21 insertions, 5 deletions
diff --git a/src/PeerAbstractCommand.cc b/src/PeerAbstractCommand.cc
index 6ab8aa77..31f85441 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"
@@ -179,6 +180,26 @@ void PeerAbstractCommand::setWriteCheckSocket(
   }
 }
 
+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 +207,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));