From 1176421bf2fdf9a645452c482d58792ef63be2bb Mon Sep 17 00:00:00 2001
From: John Preston <johnprestonmail@gmail.com>
Date: Tue, 10 Jan 2023 15:03:37 +0400
Subject: [PATCH] Limit special config request types.

---
 .../SourceFiles/mtproto/config_loader.cpp     |  5 ++--
 Telegram/SourceFiles/mtproto/mtp_instance.cpp |  2 +-
 .../mtproto/special_config_request.cpp        | 29 ++++++++++++-------
 .../mtproto/special_config_request.h          |  3 ++
 4 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/Telegram/SourceFiles/mtproto/config_loader.cpp b/Telegram/SourceFiles/mtproto/config_loader.cpp
index 3529dd33e..5d07404ba 100644
--- a/Telegram/SourceFiles/mtproto/config_loader.cpp
+++ b/Telegram/SourceFiles/mtproto/config_loader.cpp
@@ -116,7 +116,7 @@ void ConfigLoader::enumerate() {
 }
 
 void ConfigLoader::refreshSpecialLoader() {
-	if (_proxyEnabled) {
+	if (_proxyEnabled || _instance->isKeysDestroyer()) {
 		_specialLoader.reset();
 		return;
 	}
@@ -136,6 +136,7 @@ void ConfigLoader::setPhone(const QString &phone) {
 }
 
 void ConfigLoader::createSpecialLoader() {
+	const auto testMode = _instance->isTestMode();
 	_triedSpecialEndpoints.clear();
 	_specialLoader = std::make_unique<SpecialConfigRequest>([=](
 			DcId dcId,
@@ -147,7 +148,7 @@ void ConfigLoader::createSpecialLoader() {
 		} else {
 			addSpecialEndpoint(dcId, ip, port, secret);
 		}
-	}, _instance->configValues().txtDomainString, _phone);
+	}, testMode, _instance->configValues().txtDomainString, _phone);
 }
 
 void ConfigLoader::addSpecialEndpoint(
diff --git a/Telegram/SourceFiles/mtproto/mtp_instance.cpp b/Telegram/SourceFiles/mtproto/mtp_instance.cpp
index e24713a03..1794fe8b2 100644
--- a/Telegram/SourceFiles/mtproto/mtp_instance.cpp
+++ b/Telegram/SourceFiles/mtproto/mtp_instance.cpp
@@ -542,7 +542,7 @@ void Instance::Private::syncHttpUnixtime() {
 		InvokeQueued(_instance, [=] {
 			_httpUnixtimeLoader = nullptr;
 		});
-	}, configValues().txtDomainString);
+	}, isTestMode(), configValues().txtDomainString);
 }
 
 void Instance::Private::restartedByTimeout(ShiftedDcId shiftedDcId) {
diff --git a/Telegram/SourceFiles/mtproto/special_config_request.cpp b/Telegram/SourceFiles/mtproto/special_config_request.cpp
index c6743fcd9..8bbfab698 100644
--- a/Telegram/SourceFiles/mtproto/special_config_request.cpp
+++ b/Telegram/SourceFiles/mtproto/special_config_request.cpp
@@ -188,6 +188,7 @@ SpecialConfigRequest::SpecialConfigRequest(
 		int port,
 		bytes::const_span secret)> callback,
 	Fn<void()> timeDoneCallback,
+	bool isTestMode,
 	const QString &domainString,
 	const QString &phone)
 : _callback(std::move(callback))
@@ -219,14 +220,9 @@ SpecialConfigRequest::SpecialConfigRequest(
 
 	_attempts = {};
 	_attempts.push_back({ Type::Google, "dns.google.com" });
-	_attempts.push_back({ Type::Google, takeDomain(), "dns" });
 	_attempts.push_back({ Type::Mozilla, "mozilla.cloudflare-dns.com" });
 	_attempts.push_back({ Type::RemoteConfig, "firebaseremoteconfig" });
-	while (!domains.empty()) {
-		_attempts.push_back({ Type::Google, takeDomain(), "dns" });
-	}
 	if (!_timeDoneCallback) {
-		_attempts.push_back({ Type::Realtime, "firebaseio.com" });
 		_attempts.push_back({ Type::FireStore, "firestore" });
 		for (const auto &domain : DnsDomains()) {
 			_attempts.push_back({ Type::FireStore, domain, "firestore" });
@@ -234,12 +230,15 @@ SpecialConfigRequest::SpecialConfigRequest(
 	}
 
 	shuffle(0, 2);
-	shuffle(2, 4);
 	if (!_timeDoneCallback) {
-		shuffle(
-			_attempts.size() - (2 + domainsCount),
-			_attempts.size() - domainsCount);
-		shuffle(_attempts.size() - domainsCount, _attempts.size());
+		shuffle(_attempts.size() - (domainsCount + 1), _attempts.size());
+	}
+	if (isTestMode) {
+		_attempts.erase(ranges::remove_if(_attempts, [](
+				const Attempt &attempt) {
+			return (attempt.type != Type::Google)
+				&& (attempt.type != Type::Mozilla);
+		}), _attempts.end());
 	}
 	ranges::reverse(_attempts); // We go from last to first.
 
@@ -252,17 +251,25 @@ SpecialConfigRequest::SpecialConfigRequest(
 		const std::string &ip,
 		int port,
 		bytes::const_span secret)> callback,
+	bool isTestMode,
 	const QString &domainString,
 	const QString &phone)
-: SpecialConfigRequest(std::move(callback), nullptr, domainString, phone) {
+: SpecialConfigRequest(
+	std::move(callback),
+	nullptr,
+	isTestMode,
+	domainString,
+	phone) {
 }
 
 SpecialConfigRequest::SpecialConfigRequest(
 	Fn<void()> timeDoneCallback,
+	bool isTestMode,
 	const QString &domainString)
 : SpecialConfigRequest(
 	nullptr,
 	std::move(timeDoneCallback),
+	isTestMode,
 	domainString,
 	QString()) {
 }
diff --git a/Telegram/SourceFiles/mtproto/special_config_request.h b/Telegram/SourceFiles/mtproto/special_config_request.h
index 6962c04fe..91901489e 100644
--- a/Telegram/SourceFiles/mtproto/special_config_request.h
+++ b/Telegram/SourceFiles/mtproto/special_config_request.h
@@ -25,10 +25,12 @@ public:
 			const std::string &ip,
 			int port,
 			bytes::const_span secret)> callback,
+		bool isTestMode,
 		const QString &domainString,
 		const QString &phone);
 	SpecialConfigRequest(
 		Fn<void()> timeDoneCallback,
+		bool isTestMode,
 		const QString &domainString);
 
 private:
@@ -52,6 +54,7 @@ private:
 			int port,
 			bytes::const_span secret)> callback,
 		Fn<void()> timeDoneCallback,
+		bool isTestMode,
 		const QString &domainString,
 		const QString &phone);