From 746c70900b05df664e19000b9fd19e49806c0117 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Sat, 27 Nov 2021 22:22:51 +0300
Subject: [PATCH] Use different timeouts for different handshake parts.

---
 td/mtproto/Handshake.cpp | 21 ++++++++++++++-------
 td/mtproto/Handshake.h   |  3 ++-
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/td/mtproto/Handshake.cpp b/td/mtproto/Handshake.cpp
index f47d2f8f1..421ad5788 100644
--- a/td/mtproto/Handshake.cpp
+++ b/td/mtproto/Handshake.cpp
@@ -50,17 +50,20 @@ AuthKeyHandshake::AuthKeyHandshake(int32 dc_id, int32 expires_in)
     : mode_(expires_in == 0 ? Mode::Main : Mode::Temp)
     , dc_id_(dc_id)
     , expires_in_(expires_in)
-    , timeout_at_(Time::now() + 1e9) {
+    , start_time_(Time::now())
+    , timeout_in_(1e9) {
 }
 
 void AuthKeyHandshake::set_timeout_in(double timeout_in) {
-  timeout_at_ = Time::now() + timeout_in;
+  start_time_ = Time::now();
+  timeout_in_ = timeout_in;
 }
 
 void AuthKeyHandshake::clear() {
   last_query_ = BufferSlice();
   state_ = Start;
-  timeout_at_ = Time::now() + 1e9;
+  start_time_ = Time::now();
+  timeout_in_ = 1e9;
 }
 
 bool AuthKeyHandshake::is_ready_for_finish() const {
@@ -81,6 +84,10 @@ string AuthKeyHandshake::store_object(const mtproto_api::Object &object) {
 }
 
 Status AuthKeyHandshake::on_res_pq(Slice message, Callback *connection, PublicRsaKeyInterface *public_rsa_key) {
+  if (Time::now() >= start_time_ + timeout_in_ * 0.6) {
+    return Status::Error("Handshake ResPQ timeout expired");
+  }
+
   TRY_RESULT(res_pq, fetch_result<mtproto_api::req_pq_multi>(message, false));
   if (res_pq->nonce_ != nonce_) {
     return Status::Error("Nonce mismatch");
@@ -155,6 +162,10 @@ Status AuthKeyHandshake::on_res_pq(Slice message, Callback *connection, PublicRs
 }
 
 Status AuthKeyHandshake::on_server_dh_params(Slice message, Callback *connection, DhCallback *dh_callback) {
+  if (Time::now() >= start_time_ + timeout_in_ * 0.8) {
+    return Status::Error("Handshake DH params timeout expired");
+  }
+
   TRY_RESULT(dh_params, fetch_result<mtproto_api::req_DH_params>(message, false));
 
   // server_DH_params_ok#d0e8075c nonce:int128 server_nonce:int128 encrypted_answer:string = Server_DH_Params;
@@ -318,10 +329,6 @@ Status AuthKeyHandshake::on_start(Callback *connection) {
 
 Status AuthKeyHandshake::on_message(Slice message, Callback *connection, AuthKeyHandshakeContext *context) {
   Status status = [&] {
-    if (Time::now() >= timeout_at_) {
-      return Status::Error("Handshake timeout expired");
-    }
-
     switch (state_) {
       case ResPQ:
         return on_res_pq(message, connection, context->get_public_rsa_key_interface());
diff --git a/td/mtproto/Handshake.h b/td/mtproto/Handshake.h
index b3979ed20..5baf7657d 100644
--- a/td/mtproto/Handshake.h
+++ b/td/mtproto/Handshake.h
@@ -82,7 +82,8 @@ class AuthKeyHandshake {
   int32 expires_in_ = 0;
   double expires_at_ = 0;
 
-  double timeout_at_ = 0;
+  double start_time_ = 0;
+  double timeout_in_ = 0;
 
   AuthKey auth_key_;
   double server_time_diff_ = 0;