Skip to content

Instantly share code, notes, and snippets.

@chisaato
Last active August 22, 2024 10:23
Show Gist options
  • Save chisaato/af02be29b826ee6545718dfc88174b31 to your computer and use it in GitHub Desktop.
Save chisaato/af02be29b826ee6545718dfc88174b31 to your computer and use it in GitHub Desktop.
Aria2 添加 IP 黑名单补丁
diff --git a/src/ActivePeerConnectionCommand.cc b/src/ActivePeerConnectionCommand.cc
index 7e566b42..bf044da9 100644
--- a/src/ActivePeerConnectionCommand.cc
+++ b/src/ActivePeerConnectionCommand.cc
@@ -133,6 +133,15 @@ void ActivePeerConnectionCommand::makeNewConnections(int num)
for (; num && peerStorage_->isPeerAvailable(); --num) {
cuid_t ncuid = e_->newCUID();
std::shared_ptr<Peer> peer = peerStorage_->checkoutPeer(ncuid);
+ // 在这里检查 IP
+ if (bittorrent::isIpInList(peer->getIPAddress(),requestGroup_->getOption()))
+ {
+ A2_LOG_INFO(
+ fmt("CUID#%" PRId64 " - 拒绝连接到 %s 因为在黑名单中(活动 Peer 连接)", getCuid(), peer->getIPAddress().c_str()));
+ peerStorage_->addBadPeer(peer->getIPAddress());
+ break;
+ }
+
// sanity check
if (!peer) {
break;
diff --git a/src/DefaultPeerStorage.cc b/src/DefaultPeerStorage.cc
index b4a9fc00..1ca1ede0 100644
--- a/src/DefaultPeerStorage.cc
+++ b/src/DefaultPeerStorage.cc
@@ -324,6 +324,32 @@ void DefaultPeerStorage::returnPeer(const std::shared_ptr<Peer>& peer)
}
}
+void DefaultPeerStorage::disconnectPeer(const std::shared_ptr<Peer>& peer)
+{
+ A2_LOG_DEBUG(fmt("Peer %s:%u returned from CUID#%" PRId64,
+ peer->getIPAddress().c_str(), peer->getOrigPort(),
+ peer->usedBy()));
+ if (usedPeers_.erase(peer)) {
+ if (peer->isActive()) {
+ if (peer->isDisconnectedGracefully() && !peer->isIncomingPeer()) {
+ peer->startDrop();
+ addDroppedPeer(peer);
+ }
+ // Execute choking algorithm if unchoked and interested peer is
+ // disconnected.
+ // if (!peer->amChoking() && peer->peerInterested()) {
+ // executeChoke();
+ // }
+ }
+ peer->usedBy(0);
+ onErasingPeer(peer);
+ }
+ else {
+ A2_LOG_WARN(fmt("Cannot find peer %s:%u in usedPeers_",
+ peer->getIPAddress().c_str(), peer->getOrigPort()));
+ }
+}
+
bool DefaultPeerStorage::chokeRoundIntervalElapsed()
{
constexpr auto CHOKE_ROUND_INTERVAL = 10_s;
diff --git a/src/DefaultPeerStorage.h b/src/DefaultPeerStorage.h
index 4ba716aa..3b3826fe 100644
--- a/src/DefaultPeerStorage.h
+++ b/src/DefaultPeerStorage.h
@@ -113,6 +113,7 @@ public:
virtual std::shared_ptr<Peer> checkoutPeer(cuid_t cuid) CXX11_OVERRIDE;
virtual void returnPeer(const std::shared_ptr<Peer>& peer) CXX11_OVERRIDE;
+ virtual void disconnectPeer(const std::shared_ptr<Peer>& peer) CXX11_OVERRIDE;
virtual bool chokeRoundIntervalElapsed() CXX11_OVERRIDE;
diff --git a/src/OptionHandlerFactory.cc b/src/OptionHandlerFactory.cc
index 000426b6..8917ee4d 100644
--- a/src/OptionHandlerFactory.cc
+++ b/src/OptionHandlerFactory.cc
@@ -1680,6 +1680,15 @@ std::vector<OptionHandler*> OptionHandlerFactory::createOptionHandlers()
op->setChangeOptionForReserved(true);
handlers.push_back(op);
}
+ {
+ OptionHandler* op(new DefaultOptionHandler(PREF_BT_BLACKLIST_IP, TEXT_BT_BLACKLIST_IP,
+ NO_DESCRIPTION, "IP,..."));
+ op->addTag(TAG_BITTORRENT);
+ op->setInitialOption(true);
+ op->setChangeGlobalOption(true);
+ op->setChangeOptionForReserved(true);
+ handlers.push_back(op);
+ }
{
OptionHandler* op(new NumberOptionHandler(PREF_BT_TRACKER_CONNECT_TIMEOUT,
TEXT_BT_TRACKER_CONNECT_TIMEOUT,
diff --git a/src/PeerAbstractCommand.cc b/src/PeerAbstractCommand.cc
index 6ab8aa77..6c2daf05 100644
--- a/src/PeerAbstractCommand.cc
+++ b/src/PeerAbstractCommand.cc
@@ -47,6 +47,8 @@
#include "wallclock.h"
#include "util.h"
+#include <bittorrent_helper.h>
+
namespace aria2 {
PeerAbstractCommand::PeerAbstractCommand(cuid_t cuid,
@@ -97,6 +99,15 @@ bool PeerAbstractCommand::execute()
if (checkPoint_.difference(global::wallclock()) >= timeout_) {
throw DL_ABORT_EX(EX_TIME_OUT);
}
+ // 在这里检查 IP
+ if (bittorrent::isIpInList(peer_->getIPAddress(),*e_->getOption())) {
+ A2_LOG_INFO(fmt("CUID#%" PRId64
+ " - 拒绝连接到 %s 因为在黑名单中 (在抽象命令中)",
+ getCuid(), peer_->getIPAddress().c_str()));
+ // peerStorage_->addBadPeer(peer->getIPAddress());
+ onAbort();
+ return prepareForNextPeer(0);
+ }
return executeInternal();
}
catch (DownloadFailureException& err) {
diff --git a/src/PeerStorage.h b/src/PeerStorage.h
index cc77c515..950b3aa1 100644
--- a/src/PeerStorage.h
+++ b/src/PeerStorage.h
@@ -120,6 +120,8 @@ public:
virtual bool chokeRoundIntervalElapsed() = 0;
virtual void executeChoke() = 0;
+ virtual void disconnectPeer(const std::shared_ptr<Peer>& peer)=0;
+
};
} // namespace aria2
diff --git a/src/RequestGroup.cc b/src/RequestGroup.cc
index fd6801fc..cf3ef1cf 100644
--- a/src/RequestGroup.cc
+++ b/src/RequestGroup.cc
@@ -974,6 +974,32 @@ int RequestGroup::getNumConnection() const
return numConnection;
}
+void RequestGroup::disconnectPeersByBlacklist(DownloadEngine* e) const
+{
+ A2_LOG_INFO(
+ fmt("GID#%" PRId64 " - 现在主动断开 Peer 列表中的黑名单 Peer", gid_))
+ // 从 PeerStorage 拿
+ for (auto& peer : peerStorage_->getUsedPeers()) {
+ if (!peer->isActive()) {
+ // A2_LOG_INFO(fmt("GID#%" PRId64 " - 不活跃,跳过", gid_));
+ continue;
+ }
+ // 拿出 IP
+ if (bittorrent::isIpInList(peer->getIPAddress(), getOption())) {
+ // A2_LOG_INFO(
+ // fmt("GID#%" PRId64 " - 现在调用 peerStorage 回收这个 peer", gid_));
+ // peer->releaseSessionResource();
+ // 记录 CUID
+ peer->usedBy(0);
+ // A2_LOG_INFO(fmt("GID#%" PRId64 " - 已将 Peer 置 0", gid_));
+
+ peerStorage_->disconnectPeer(peer);
+ // A2_LOG_INFO(fmt("GID#%" PRId64 " - 已断开并拉黑 Peer %s ", gid_,
+ // peer->getIPAddress().c_str()));
+ }
+ }
+}
+
void RequestGroup::increaseNumCommand() { ++numCommand_; }
void RequestGroup::decreaseNumCommand()
diff --git a/src/RequestGroup.h b/src/RequestGroup.h
index 6698f93d..27652b49 100644
--- a/src/RequestGroup.h
+++ b/src/RequestGroup.h
@@ -310,6 +310,8 @@ public:
int getNumConnection() const;
+ void disconnectPeersByBlacklist(DownloadEngine* e) const;
+
void increaseNumCommand();
void decreaseNumCommand();
diff --git a/src/RpcMethodFactory.cc b/src/RpcMethodFactory.cc
index bafbffd8..1394164c 100644
--- a/src/RpcMethodFactory.cc
+++ b/src/RpcMethodFactory.cc
@@ -146,7 +146,11 @@ std::unique_ptr<RpcMethod> createMethod(const std::string& methodName)
if (methodName == ForcePauseRpcMethod::getMethodName()) {
return make_unique<ForcePauseRpcMethod>();
}
-
+
+ if (methodName == DisconnectPeerByBlacklistRpcMethod::getMethodName()) {
+ return make_unique<DisconnectPeerByBlacklistRpcMethod>();
+ }
+
if (methodName == PauseAllRpcMethod::getMethodName()) {
return make_unique<PauseAllRpcMethod>();
}
diff --git a/src/RpcMethodImpl.cc b/src/RpcMethodImpl.cc
index c2249202..2f3bf527 100644
--- a/src/RpcMethodImpl.cc
+++ b/src/RpcMethodImpl.cc
@@ -465,6 +465,25 @@ std::unique_ptr<ValueBase> ForcePauseRpcMethod::process(const RpcRequest& req,
return pauseDownload(req, e, true);
}
+std::unique_ptr<ValueBase>
+DisconnectPeerByBlacklistRpcMethod::process(const RpcRequest& req,
+ DownloadEngine* e)
+{
+ // 先取 pid
+ const String* gidParam = checkRequiredParam<String>(req, 0);
+ a2_gid_t gid = str2Gid(gidParam);
+ auto group = e->getRequestGroupMan()->findGroup(gid);
+ if (group) {
+ // group->setForceHaltRequested(true, RequestGroup::USER_REQUEST);
+ if (group->getState() == RequestGroup::STATE_ACTIVE) {
+ group->disconnectPeersByBlacklist(e);
+ }
+ return createGIDResponse(gid);
+ }
+ // 如果因为各种奇怪原因不能断开,也不用报错
+ return createOKResponse();
+}
+
namespace {
template <typename InputIterator>
void pauseRequestGroups(InputIterator first, InputIterator last, bool reserved,
@@ -1685,6 +1704,14 @@ void changeGlobalOption(const Option& option, DownloadEngine* e)
auto& openedFileCounter = e->getRequestGroupMan()->getOpenedFileCounter();
openedFileCounter->setMaxOpenFiles(option.getAsInt(PREF_BT_MAX_OPEN_FILES));
}
+ // 向每一个下载更新 PREF_BT_BLACKLIST_IP
+ if (option.defined(PREF_BT_BLACKLIST_IP)) {
+ auto& groups = e->getRequestGroupMan()->getRequestGroups();
+ for (auto& group : groups) {
+ group->getOption()->put(PREF_BT_BLACKLIST_IP,
+ option.get(PREF_BT_BLACKLIST_IP));
+ }
+ }
}
} // namespace aria2
diff --git a/src/RpcMethodImpl.h b/src/RpcMethodImpl.h
index de6d7685..84ba6eb4 100644
--- a/src/RpcMethodImpl.h
+++ b/src/RpcMethodImpl.h
@@ -180,6 +180,14 @@ protected:
public:
static const char* getMethodName() { return "aria2.pauseAll"; }
};
+class DisconnectPeerByBlacklistRpcMethod : public RpcMethod {
+protected:
+ virtual std::unique_ptr<ValueBase> process(const RpcRequest& req,
+ DownloadEngine* e) CXX11_OVERRIDE;
+
+public:
+ static const char* getMethodName() { return "aria2.disconnectPeerByBlacklist"; }
+};
class ForcePauseAllRpcMethod : public RpcMethod {
protected:
diff --git a/src/bittorrent_helper.cc b/src/bittorrent_helper.cc
index d3df0b4b..814fdcd9 100644
--- a/src/bittorrent_helper.cc
+++ b/src/bittorrent_helper.cc
@@ -63,6 +63,8 @@
#include "DownloadFailureException.h"
#include "ValueBaseBencodeParser.h"
+#include <LogFactory.h>
+
namespace aria2 {
namespace bittorrent {
@@ -1081,6 +1083,29 @@ void adjustAnnounceUri(TorrentAttribute* attrs,
addAnnounceUri(attrs, addUris);
}
+// 来一个判断 IP 是否在列表中的函数
+bool isIpInList(const std::string& ip, const std::shared_ptr<Option>& option)
+{
+ A2_LOG_WARN(fmt("现在检测 IP %s 是否在黑名单", ip.c_str()))
+ // 获取黑名单列表
+ std::vector<std::string> ipList;
+ util::split(option->get(PREF_BT_BLACKLIST_IP).begin(),
+ option->get(PREF_BT_BLACKLIST_IP).end(),
+ std::back_inserter(ipList), ',', true);
+ return std::find(ipList.begin(), ipList.end(), ip) != ipList.end();
+}
+
+// 判断 IP 但是传参改为 const std::string&
+bool isIpInList(const std::string& ip, const Option& option)
+{
+ A2_LOG_WARN(fmt("现在检测 IP %s 是否在黑名单", ip.c_str()))
+ // 获取黑名单列表
+ std::vector<std::string> ipList;
+ util::split(option.get(PREF_BT_BLACKLIST_IP).begin(),
+ option.get(PREF_BT_BLACKLIST_IP).end(),
+ std::back_inserter(ipList), ',', true);
+ return std::find(ipList.begin(), ipList.end(), ip) != ipList.end();
+}
const char* getModeString(BtFileMode mode)
{
switch (mode) {
diff --git a/src/bittorrent_helper.h b/src/bittorrent_helper.h
index ea1a6277..1f9b242b 100644
--- a/src/bittorrent_helper.h
+++ b/src/bittorrent_helper.h
@@ -337,6 +337,12 @@ int getCompactLength(int family);
// Returns textual representation of the |mode|.
const char* getModeString(BtFileMode mode);
+// 返回这个 IP 是否在黑名单
+bool isIpInList(const std::string& ip, const std::shared_ptr<Option>& option);
+
+// 返回该 IP 是否在黑名单,但是传参使用字符串
+bool isIpInList(const std::string& ip, const Option& option);
+
// Writes the detailed information about torrent loaded in dctx.
template <typename Output>
void print(Output& o, const std::shared_ptr<DownloadContext>& dctx)
diff --git a/src/prefs.cc b/src/prefs.cc
index 2591b9f0..1c23197d 100644
--- a/src/prefs.cc
+++ b/src/prefs.cc
@@ -558,6 +558,8 @@ PrefPtr PREF_ON_BT_DOWNLOAD_COMPLETE = makePref("on-bt-download-complete");
// values: string
PrefPtr PREF_BT_TRACKER = makePref("bt-tracker");
// values: string
+PrefPtr PREF_BT_BLACKLIST_IP = makePref("bt-blacklist-ip");
+// values: string
PrefPtr PREF_BT_EXCLUDE_TRACKER = makePref("bt-exclude-tracker");
// values: true | false
PrefPtr PREF_BT_REMOVE_UNSELECTED_FILE = makePref("bt-remove-unselected-file");
diff --git a/src/prefs.h b/src/prefs.h
index 338fd6e6..2235160a 100644
--- a/src/prefs.h
+++ b/src/prefs.h
@@ -508,6 +508,8 @@ extern PrefPtr PREF_ON_BT_DOWNLOAD_COMPLETE;
// values: string
extern PrefPtr PREF_BT_TRACKER;
// values: string
+extern PrefPtr PREF_BT_BLACKLIST_IP;
+// values: string
extern PrefPtr PREF_BT_EXCLUDE_TRACKER;
// values: true | false
extern PrefPtr PREF_BT_REMOVE_UNSELECTED_FILE;
diff --git a/src/usage_text.h b/src/usage_text.h
index 08715436..4f837bb6 100644
--- a/src/usage_text.h
+++ b/src/usage_text.h
@@ -787,6 +787,9 @@
" affected by --bt-exclude-tracker option because\n" \
" they are added after URIs in --bt-exclude-tracker\n" \
" option are removed.")
+#define TEXT_BT_BLACKLIST_IP \
+ _(" --bt-blacklist-ip=IP[,...] Comma separated list of additional BitTorrent\n" \
+ " blacklist IP.")
#define TEXT_BT_EXCLUDE_TRACKER \
_(" --bt-exclude-tracker=URI[,...] Comma separated list of BitTorrent tracker's\n" \
" announce URI to remove. You can use special value\n" \

稍后补充 RPC 调用方法
这是通过添加一个全局选项来实现的,同时这个全局选项更新的时候会联动每一项下载的选项进行更新(我知道这不怎么符合设计)
然后添加一个RPC调用来实现主动断开Peer.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment