From 03cf8b6ac2f8727bbeab0c4129a82bec8a3ad465 Mon Sep 17 00:00:00 2001
From: John Preston <johnprestonmail@gmail.com>
Date: Sun, 28 Jul 2019 18:00:42 +0200
Subject: [PATCH] Improve transport nonce generation.

---
 .../SourceFiles/mtproto/connection_tcp.cpp    | 20 ++-------------
 .../SourceFiles/mtproto/mtp_abstract_socket.h |  1 +
 .../SourceFiles/mtproto/mtp_tcp_socket.cpp    | 25 +++++++++++++++++++
 Telegram/SourceFiles/mtproto/mtp_tcp_socket.h |  1 +
 .../SourceFiles/mtproto/mtp_tls_socket.cpp    |  4 +++
 Telegram/SourceFiles/mtproto/mtp_tls_socket.h |  1 +
 6 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/Telegram/SourceFiles/mtproto/connection_tcp.cpp b/Telegram/SourceFiles/mtproto/connection_tcp.cpp
index d48eed980..2b73bce79 100644
--- a/Telegram/SourceFiles/mtproto/connection_tcp.cpp
+++ b/Telegram/SourceFiles/mtproto/connection_tcp.cpp
@@ -450,6 +450,7 @@ void TcpConnection::sendData(mtpBuffer &&buffer) {
 
 bytes::const_span TcpConnection::prepareConnectionStartPrefix(
 		bytes::span buffer) {
+	Expects(_socket != nullptr);
 	Expects(_protocol != nullptr);
 
 	if (_connectionStarted) {
@@ -460,26 +461,9 @@ bytes::const_span TcpConnection::prepareConnectionStartPrefix(
 	// prepare random part
 	char nonceBytes[64];
 	const auto nonce = bytes::make_span(nonceBytes);
-
-	const auto zero = reinterpret_cast<uchar*>(nonce.data());
-	const auto first = reinterpret_cast<uint32*>(nonce.data());
-	const auto second = first + 1;
-	const auto reserved01 = 0x000000EFU;
-	const auto reserved11 = 0x44414548U;
-	const auto reserved12 = 0x54534F50U;
-	const auto reserved13 = 0x20544547U;
-	const auto reserved14 = 0xEEEEEEEEU;
-	const auto reserved15 = 0xDDDDDDDDU;
-	const auto reserved21 = 0x00000000U;
 	do {
 		bytes::set_random(nonce);
-	} while (*zero == reserved01
-		|| *first == reserved11
-		|| *first == reserved12
-		|| *first == reserved13
-		|| *first == reserved14
-		|| *first == reserved15
-		|| *second == reserved21);
+	} while (!_socket->isGoodStartNonce(nonce));
 
 	// prepare encryption key/iv
 	_protocol->prepareKey(
diff --git a/Telegram/SourceFiles/mtproto/mtp_abstract_socket.h b/Telegram/SourceFiles/mtproto/mtp_abstract_socket.h
index 4ac12bb75..4d5c35108 100644
--- a/Telegram/SourceFiles/mtproto/mtp_abstract_socket.h
+++ b/Telegram/SourceFiles/mtproto/mtp_abstract_socket.h
@@ -42,6 +42,7 @@ public:
 	}
 
 	virtual void connectToHost(const QString &address, int port) = 0;
+	[[nodiscard]] virtual bool isGoodStartNonce(bytes::const_span nonce) = 0;
 	virtual void timedOut() = 0;
 	[[nodiscard]] virtual bool isConnected() = 0;
 	[[nodiscard]] virtual bool hasBytesAvailable() = 0;
diff --git a/Telegram/SourceFiles/mtproto/mtp_tcp_socket.cpp b/Telegram/SourceFiles/mtproto/mtp_tcp_socket.cpp
index 4c0731bc7..3349cbadd 100644
--- a/Telegram/SourceFiles/mtproto/mtp_tcp_socket.cpp
+++ b/Telegram/SourceFiles/mtproto/mtp_tcp_socket.cpp
@@ -47,6 +47,31 @@ void TcpSocket::connectToHost(const QString &address, int port) {
 	_socket.connectToHost(address, port);
 }
 
+bool TcpSocket::isGoodStartNonce(bytes::const_span nonce) {
+	Expects(nonce.size() >= 2 * sizeof(uint32));
+
+	const auto bytes = nonce.data();
+	const auto zero = *reinterpret_cast<const uchar*>(bytes);
+	const auto first = *reinterpret_cast<const uint32*>(bytes);
+	const auto second = *(reinterpret_cast<const uint32*>(bytes) + 1);
+	const auto reserved01 = 0x000000EFU;
+	const auto reserved11 = 0x44414548U;
+	const auto reserved12 = 0x54534F50U;
+	const auto reserved13 = 0x20544547U;
+	const auto reserved14 = 0xEEEEEEEEU;
+	const auto reserved15 = 0xDDDDDDDDU;
+	const auto reserved16 = 0x02010316U;
+	const auto reserved21 = 0x00000000U;
+	return (zero != reserved01)
+		&& (first != reserved11)
+		&& (first != reserved12)
+		&& (first != reserved13)
+		&& (first != reserved14)
+		&& (first != reserved15)
+		&& (first != reserved16)
+		&& (second != reserved21);
+}
+
 void TcpSocket::timedOut() {
 }
 
diff --git a/Telegram/SourceFiles/mtproto/mtp_tcp_socket.h b/Telegram/SourceFiles/mtproto/mtp_tcp_socket.h
index 231edd5ed..487b2dcc0 100644
--- a/Telegram/SourceFiles/mtproto/mtp_tcp_socket.h
+++ b/Telegram/SourceFiles/mtproto/mtp_tcp_socket.h
@@ -17,6 +17,7 @@ public:
 	TcpSocket(not_null<QThread*> thread, const QNetworkProxy &proxy);
 
 	void connectToHost(const QString &address, int port) override;
+	bool isGoodStartNonce(bytes::const_span nonce) override;
 	void timedOut() override;
 	bool isConnected() override;
 	bool hasBytesAvailable() override;
diff --git a/Telegram/SourceFiles/mtproto/mtp_tls_socket.cpp b/Telegram/SourceFiles/mtproto/mtp_tls_socket.cpp
index 4472cf6f0..7b5e1add1 100644
--- a/Telegram/SourceFiles/mtproto/mtp_tls_socket.cpp
+++ b/Telegram/SourceFiles/mtproto/mtp_tls_socket.cpp
@@ -575,6 +575,10 @@ void TlsSocket::connectToHost(const QString &address, int port) {
 	_socket.connectToHost(address, port);
 }
 
+bool TlsSocket::isGoodStartNonce(bytes::const_span nonce) {
+	return true;
+}
+
 void TlsSocket::timedOut() {
 	_syncTimeRequests.fire({});
 }
diff --git a/Telegram/SourceFiles/mtproto/mtp_tls_socket.h b/Telegram/SourceFiles/mtproto/mtp_tls_socket.h
index e130e633e..df8742b34 100644
--- a/Telegram/SourceFiles/mtproto/mtp_tls_socket.h
+++ b/Telegram/SourceFiles/mtproto/mtp_tls_socket.h
@@ -20,6 +20,7 @@ public:
 		const QNetworkProxy &proxy);
 
 	void connectToHost(const QString &address, int port) override;
+	bool isGoodStartNonce(bytes::const_span nonce) override;
 	void timedOut() override;
 	bool isConnected() override;
 	bool hasBytesAvailable() override;