peer whitelist

src/base/CMakeLists.txt

@@ -106,6 +106,7 @@ add_library(qbt_base STATIC
     bittorrent/peer_blacklist.hpp
     bittorrent/peer_filter_plugin.hpp
     bittorrent/peer_logger.hpp
+    bittorrent/peer_whitelist.hpp
     bittorrent/portforwarderimpl.cpp
     bittorrent/resumedatasavingmanager.cpp
     bittorrent/session.cpp

src/base/bittorrent/peer_blacklist.hpp

@@ -50,8 +50,9 @@ void drop_connection(lt::peer_connection_handle ph)
 template<typename F>
 auto wrap_filter(F filter, const std::string& tag)
 {
-  return [=](const lt::peer_info& info, bool) {
+  return [=](const lt::peer_info& info, bool handshake, bool* stop_filtering) {
     bool matched = filter(info);
+    *stop_filtering = !handshake && !matched;
     if (matched)
       peer_logger_singleton::instance().log_peer(info, tag);
     return matched;

src/base/bittorrent/peer_filter_plugin.hpp

@@ -3,7 +3,7 @@
 #include <libtorrent/extensions.hpp>
 #include <libtorrent/peer_connection_handle.hpp>
 
-using filter_function = std::function<bool(const lt::peer_info&, bool)>;
+using filter_function = std::function<bool(const lt::peer_info&, bool, bool*)>;
 using action_function = std::function<void(lt::peer_connection_handle)>;
 
 class peer_filter_plugin final : public lt::peer_plugin
@@ -15,12 +15,6 @@ class peer_filter_plugin final : public lt::peer_plugin
     , m_action(std::move(action))
   {}
 
-  void add_handshake(lt::entry& e) override
-  {
-    handle_peer(true);
-    peer_plugin::add_handshake(e);
-  }
-
   bool on_handshake(lt::span<char const> d) override
   {
     handle_peer(true);
@@ -33,13 +27,64 @@ class peer_filter_plugin final : public lt::peer_plugin
     return peer_plugin::on_extension_handshake(d);
   }
 
+  bool on_interested() override
+  {
+    handle_peer();
+    return peer_plugin::on_interested();
+  }
+
+  bool on_not_interested() override
+  {
+    handle_peer();
+    return peer_plugin::on_not_interested();
+  }
+
+  bool on_have(lt::piece_index_t p) override
+  {
+    handle_peer();
+    return peer_plugin::on_have(p);
+  }
+
+  bool on_dont_have(lt::piece_index_t p) override
+  {
+    handle_peer();
+    return peer_plugin::on_dont_have(p);
+  }
+
+  bool on_bitfield(lt::bitfield const& bitfield) override
+  {
+    handle_peer();
+    return peer_plugin::on_bitfield(bitfield);
+  }
+
+  bool on_have_all() override
+  {
+    handle_peer();
+    return peer_plugin::on_have_all();
+  }
+
+  bool on_have_none() override
+  {
+    handle_peer();
+    return peer_plugin::on_have_none();
+  }
+
+  bool on_request(lt::peer_request const& r) override
+  {
+    handle_peer();
+    return peer_plugin::on_request(r);
+  }
+
 protected:
   void handle_peer(bool handshake = false)
   {
+    if (m_stop_filtering)
+      return;
+
     lt::peer_info info;
     m_peer_connection.get_peer_info(info);
 
-    if (m_filter(info, handshake))
+    if (m_filter(info, handshake, &m_stop_filtering))
       m_action(m_peer_connection);
   }
 
@@ -48,6 +93,8 @@ class peer_filter_plugin final : public lt::peer_plugin
 
   filter_function m_filter;
   action_function m_action;
+
+  bool m_stop_filtering = false;
 };
 
 

src/base/bittorrent/peer_whitelist.hpp

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <algorithm>
+#include <fstream>
+#include <string>
+
+#include <QDir>
+#include <QRegularExpression>
+#include <QString>
+#include <QVector>
+
+#include <libtorrent/peer_info.hpp>
+
+#include "base/logger.h"
+#include "base/profile.h"
+
+#include "peer_filter_plugin.hpp"
+#include "peer_logger.hpp"
+
+
+namespace {
+
+bool qregex_has_match(const QRegularExpression& re, const QString& str)
+{
+  auto m = re.match(str);
+  return m.hasMatch();
+}
+
+}
+
+class peer_filter
+{
+public:
+  explicit peer_filter(const QString& filter_file)
+  {
+    std::ifstream ifs(filter_file.toStdString());
+    std::string peer_id, client;
+    while (ifs >> peer_id >> client) {
+      QRegularExpression peer_id_re(QString::fromStdString(peer_id));
+      QRegularExpression client_re(QString::fromStdString(client));
+
+      QString msg_tmpl("whitelist: invalid %1 matching expression '%2' detected, ignoring rule");
+
+      if (!peer_id_re.isValid())
+        LogMsg(msg_tmpl.arg("peer id").arg(peer_id_re.pattern()), Log::WARNING);
+
+      if (!client_re.isValid())
+        LogMsg(msg_tmpl.arg("client name").arg(client_re.pattern()), Log::WARNING);
+
+      if (peer_id_re.isValid() && client_re.isValid())
+        m_filters.append({peer_id_re, client_re});
+    }
+
+    if (m_filters.isEmpty())
+      LogMsg("whitelist: no rules were loaded, any connections will be dropped", Log::CRITICAL);
+  }
+
+  bool match_peer(const lt::peer_info& info, bool skip_name) const
+  {
+    QString peer_id = QString::fromLatin1(info.pid.data(), 8);
+    QString client = QString::fromStdString(info.client);
+    return std::any_of(m_filters.begin(), m_filters.end(),
+                       [&](const auto& filter) {
+                           return qregex_has_match(filter[0], peer_id) &&
+                               (skip_name || qregex_has_match(filter[1], client));
+                       });
+  }
+
+private:
+  QVector<QVector<QRegularExpression>> m_filters;
+};
+
+
+std::shared_ptr<lt::torrent_plugin> create_peer_whitelist_plugin(lt::torrent_handle const&, void*)
+{
+  QDir qbt_data_dir(specialFolderLocation(SpecialFolder::Data));
+
+  QString filter_file = qbt_data_dir.absoluteFilePath(QStringLiteral("peer_whitelist.txt"));
+  // do not create plugin if filter file doesn't exists
+  if (!QFile::exists(filter_file)) {
+    LogMsg("'peer_whitelist.txt' doesn't exist, do not enabling whitelist plugin", Log::INFO);
+    return nullptr;
+  }
+
+  // create filter object only once
+  static peer_filter filter(filter_file);
+
+  auto peer_not_in_list = [&](const lt::peer_info& info, bool handshake, bool* stop_filtering) {
+    bool matched = filter.match_peer(info, handshake);
+    *stop_filtering = !handshake && matched;
+    if (!matched)
+      peer_logger_singleton::instance().log_peer(info, "whitelist");
+    return !matched;
+  };
+
+  auto drop_connection = [](lt::peer_connection_handle p) {
+    p.disconnect(boost::asio::error::connection_refused,
+                 lt::operation_t::bittorrent,
+                 lt::disconnect_severity_t{0});
+  };
+
+  return std::make_shared<peer_action_plugin>(std::move(peer_not_in_list), std::move(drop_connection));
+}

src/base/bittorrent/session.cpp

@@ -97,6 +97,7 @@
 #include "magneturi.h"
 #include "nativesessionextension.h"
 #include "peer_blacklist.hpp"
+#include "peer_whitelist.hpp"
 #include "portforwarderimpl.h"
 #include "resumedatasavingmanager.h"
 #include "statistics.h"
@@ -1206,6 +1207,7 @@ void Session::initializeNativeSession()
     }
     if (isAutoBanBTPlayerPeerEnabled())
         m_nativeSession->add_extension(&create_drop_bittorrent_media_player_plugin);
+    m_nativeSession->add_extension(&create_peer_whitelist_plugin);
 
     m_nativeSession->add_extension(std::make_shared<NativeSessionExtension>());
 }