From 8e442563f2e9f4da2535a4d8cd8e5a3ab25a0b22 Mon Sep 17 00:00:00 2001
From: John Preston <johnprestonmail@gmail.com>
Date: Sun, 20 May 2018 21:23:50 +0300
Subject: [PATCH] Enable checked auth key creation.

---
 Telegram/SourceFiles/mtproto/connection.cpp   | 21 ++++----------
 .../mtproto/connection_resolving.cpp          | 28 ++++++++++---------
 .../mtproto/connection_resolving.h            |  6 ++--
 3 files changed, 23 insertions(+), 32 deletions(-)

diff --git a/Telegram/SourceFiles/mtproto/connection.cpp b/Telegram/SourceFiles/mtproto/connection.cpp
index 33c43a099..38dfa6410 100644
--- a/Telegram/SourceFiles/mtproto/connection.cpp
+++ b/Telegram/SourceFiles/mtproto/connection.cpp
@@ -2572,22 +2572,14 @@ void ConnectionPrivate::pqAnswered() {
 		return restart();
 	}
 
-	// #TODO checked key creation
-	//auto p_q_inner = MTP_p_q_inner_data_dc(
-	//	res_pq_data.vpq,
-	//	MTP_bytes(std::move(p)),
-	//	MTP_bytes(std::move(q)),
-	//	_authKeyData->nonce,
-	//	_authKeyData->server_nonce,
-	//	_authKeyData->new_nonce,
-	//	MTP_int(getProtocolDcId()));
-	auto p_q_inner = MTP_p_q_inner_data(
+	auto p_q_inner = MTP_p_q_inner_data_dc(
 		res_pq_data.vpq,
 		MTP_bytes(std::move(p)),
 		MTP_bytes(std::move(q)),
 		_authKeyData->nonce,
 		_authKeyData->server_nonce,
-		_authKeyData->new_nonce);
+		_authKeyData->new_nonce,
+		MTP_int(getProtocolDcId()));
 	auto dhEncString = encryptPQInnerRSA(p_q_inner, rsaKey);
 	if (dhEncString.empty()) {
 		return restart();
@@ -2603,11 +2595,8 @@ void ConnectionPrivate::pqAnswered() {
 	req_DH_params.vnonce = _authKeyData->nonce;
 	req_DH_params.vserver_nonce = _authKeyData->server_nonce;
 	req_DH_params.vpublic_key_fingerprint = MTP_long(rsaKey.getFingerPrint());
-	// #TODO checked key creation
-	//req_DH_params.vp = p_q_inner.c_p_q_inner_data_dc().vp;
-	//req_DH_params.vq = p_q_inner.c_p_q_inner_data_dc().vq;
-	req_DH_params.vp = p_q_inner.c_p_q_inner_data().vp;
-	req_DH_params.vq = p_q_inner.c_p_q_inner_data().vq;
+	req_DH_params.vp = p_q_inner.c_p_q_inner_data_dc().vp;
+	req_DH_params.vq = p_q_inner.c_p_q_inner_data_dc().vq;
 	req_DH_params.vencrypted_data = MTP_bytes(dhEncString);
 	sendRequestNotSecure(req_DH_params);
 }
diff --git a/Telegram/SourceFiles/mtproto/connection_resolving.cpp b/Telegram/SourceFiles/mtproto/connection_resolving.cpp
index 4d9d535c4..1cb030af3 100644
--- a/Telegram/SourceFiles/mtproto/connection_resolving.cpp
+++ b/Telegram/SourceFiles/mtproto/connection_resolving.cpp
@@ -22,7 +22,7 @@ ResolvingConnection::ResolvingConnection(
 	ConnectionPointer &&child)
 : AbstractConnection(thread, proxy)
 , _instance(instance)
-, _timeoutTimer([=] { handleError(); }) {
+, _timeoutTimer([=] { handleError(kErrorCodeOther); }) {
 	setChild(std::move(child));
 	if (proxy.resolvedExpireAt < getms(true)) {
 		const auto host = proxy.host;
@@ -104,7 +104,7 @@ void ResolvingConnection::domainResolved(
 	if (index < _proxy.resolvedIPs.size()) {
 		_proxy.resolvedIPs.resize(index);
 		if (_ipIndex >= index) {
-			emitError();
+			emitError(kErrorCodeOther);
 		}
 	}
 	if (_ipIndex < 0) {
@@ -112,28 +112,30 @@ void ResolvingConnection::domainResolved(
 	}
 }
 
-void ResolvingConnection::refreshChild() {
+bool ResolvingConnection::refreshChild() {
 	if (!_child) {
-		return;
+		return true;
 	} else if (++_ipIndex >= _proxy.resolvedIPs.size()) {
-		emitError();
-		return;
+		return false;
 	}
 	setChild(_child->clone(ToDirectIpProxy(_proxy, _ipIndex)));
 	_timeoutTimer.callOnce(kOneConnectionTimeout);
+	return true;
 }
 
-void ResolvingConnection::emitError() {
+void ResolvingConnection::emitError(int errorCode) {
 	_ipIndex = -1;
 	_child = nullptr;
-	emit error(kErrorCodeOther);
+	emit error(errorCode);
 }
 
-void ResolvingConnection::handleError() {
+void ResolvingConnection::handleError(int errorCode) {
 	if (_connected) {
-		emitError();
+		emitError(errorCode);
 	} else if (!_proxy.resolvedIPs.empty()) {
-		refreshChild();
+		if (!refreshChild()) {
+			emitError(errorCode);
+		}
 	} else {
 		// Wait for the domain to be resolved.
 	}
@@ -143,7 +145,7 @@ void ResolvingConnection::handleDisconnected() {
 	if (_connected) {
 		emit disconnected();
 	} else {
-		handleError();
+		handleError(kErrorCodeOther);
 	}
 }
 
@@ -204,7 +206,7 @@ void ResolvingConnection::connectToServer(
 		const bytes::vector &protocolSecret,
 		int16 protocolDcId) {
 	if (!_child) {
-		InvokeQueued(this, [=] { emitError(); });
+		InvokeQueued(this, [=] { emitError(kErrorCodeOther); });
 		return;
 	}
 	_address = address;
diff --git a/Telegram/SourceFiles/mtproto/connection_resolving.h b/Telegram/SourceFiles/mtproto/connection_resolving.h
index cc70c7da4..ed80ba3d5 100644
--- a/Telegram/SourceFiles/mtproto/connection_resolving.h
+++ b/Telegram/SourceFiles/mtproto/connection_resolving.h
@@ -42,14 +42,14 @@ public:
 
 private:
 	void setChild(ConnectionPointer &&child);
-	void refreshChild();
-	void emitError();
+	bool refreshChild();
+	void emitError(int errorCode);
 
 	void domainResolved(
 		const QString &host,
 		const QStringList &ips,
 		qint64 expireAt);
-	void handleError();
+	void handleError(int errorCode);
 	void handleConnected();
 	void handleDisconnected();
 	void handleReceivedData();