From 41cc287d666b7bad6d864c1c955e5f9b3291edbd Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 20 Oct 2021 00:54:49 +0300
Subject: [PATCH] Improve checks.

---
 td/mtproto/Handshake.cpp    | 21 ++++++++++++++++++---
 tdutils/td/utils/crypto.cpp |  6 ++++++
 tdutils/td/utils/crypto.h   |  2 ++
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/td/mtproto/Handshake.cpp b/td/mtproto/Handshake.cpp
index 79fd9a716..f47d2f8f1 100644
--- a/td/mtproto/Handshake.cpp
+++ b/td/mtproto/Handshake.cpp
@@ -247,17 +247,32 @@ Status AuthKeyHandshake::on_server_dh_params(Slice message, Callback *connection
 Status AuthKeyHandshake::on_dh_gen_response(Slice message, Callback *connection) {
   TRY_RESULT(answer, fetch_result<mtproto_api::set_client_DH_params>(message, false));
   switch (answer->get_id()) {
-    case mtproto_api::dh_gen_ok::ID:
+    case mtproto_api::dh_gen_ok::ID: {
+      auto dh_gen_ok = move_tl_object_as<mtproto_api::dh_gen_ok>(answer);
+      if (dh_gen_ok->nonce_ != nonce_) {
+        return Status::Error("Nonce mismatch");
+      }
+      if (dh_gen_ok->server_nonce_ != server_nonce_) {
+        return Status::Error("Server nonce mismatch");
+      }
+
+      UInt<160> auth_key_sha1;
+      sha1(auth_key_.key(), auth_key_sha1.raw);
+      auto new_nonce_hash = sha1(PSLICE() << new_nonce_.as_slice() << '\x01' << auth_key_sha1.as_slice().substr(0, 8));
+      if (dh_gen_ok->new_nonce_hash1_.as_slice() != Slice(new_nonce_hash).substr(4)) {
+        return Status::Error("New nonce hash mismatch");
+      }
       state_ = Finish;
-      break;
+      return Status::OK();
+    }
     case mtproto_api::dh_gen_fail::ID:
       return Status::Error("DhGenFail");
     case mtproto_api::dh_gen_retry::ID:
       return Status::Error("DhGenRetry");
     default:
+      UNREACHABLE();
       return Status::Error("Unknown set_client_DH_params response");
   }
-  return Status::OK();
 }
 
 void AuthKeyHandshake::send(Callback *connection, const Storer &storer) {
diff --git a/tdutils/td/utils/crypto.cpp b/tdutils/td/utils/crypto.cpp
index c023f8910..377be86ef 100644
--- a/tdutils/td/utils/crypto.cpp
+++ b/tdutils/td/utils/crypto.cpp
@@ -723,6 +723,12 @@ void sha512(Slice data, MutableSlice output) {
 #endif
 }
 
+string sha1(Slice data) {
+  string result(20, '\0');
+  sha1(data, MutableSlice(result).ubegin());
+  return result;
+}
+
 string sha256(Slice data) {
   string result(32, '\0');
   sha256(data, result);
diff --git a/tdutils/td/utils/crypto.h b/tdutils/td/utils/crypto.h
index 714620060..071c0fc46 100644
--- a/tdutils/td/utils/crypto.h
+++ b/tdutils/td/utils/crypto.h
@@ -122,6 +122,8 @@ void sha256(Slice data, MutableSlice output);
 
 void sha512(Slice data, MutableSlice output);
 
+string sha1(Slice data) TD_WARN_UNUSED_RESULT;
+
 string sha256(Slice data) TD_WARN_UNUSED_RESULT;
 
 string sha512(Slice data) TD_WARN_UNUSED_RESULT;