From dd1e0af5dcb159acbb88a5d1f4eae938c3c32f1c Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Thu, 6 Jan 2022 21:19:07 +0300
Subject: [PATCH] Keep connection header for 3 seconds after successful
 response.

---
 td/mtproto/AuthData.h |  6 ++++--
 td/mtproto/AuthKey.h  | 15 ++++++++++-----
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/td/mtproto/AuthData.h b/td/mtproto/AuthData.h
index 31323e0e7..93cc70787 100644
--- a/td/mtproto/AuthData.h
+++ b/td/mtproto/AuthData.h
@@ -162,17 +162,19 @@ class AuthData {
       return main_auth_key_.need_header() ? Slice(header_) : Slice();
     }
   }
+
   void set_header(std::string header) {
     header_ = std::move(header);
   }
+
   void on_api_response() {
     if (use_pfs()) {
       if (tmp_auth_key_.auth_flag()) {
-        tmp_auth_key_.set_need_header(false);
+        tmp_auth_key_.remove_header();
       }
     } else {
       if (main_auth_key_.auth_flag()) {
-        main_auth_key_.set_need_header(false);
+        main_auth_key_.remove_header();
       }
     }
   }
diff --git a/td/mtproto/AuthKey.h b/td/mtproto/AuthKey.h
index 7a50c8a71..98e9d9dad 100644
--- a/td/mtproto/AuthKey.h
+++ b/td/mtproto/AuthKey.h
@@ -39,11 +39,15 @@ class AuthKey {
   }
 
   bool need_header() const {
-    return need_header_;
+    return have_header_ || Time::now() < header_expires_at_;
   }
-  void set_need_header(bool need_header) {
-    need_header_ = need_header;
+  void remove_header() {
+    if (have_header_) {
+      have_header_ = false;
+      header_expires_at_ = Time::now() + 3;
+    }
   }
+
   double expires_at() const {
     return expires_at_;
   }
@@ -86,14 +90,15 @@ class AuthKey {
       created_at_ = parser.fetch_double();
     }
     // just in case
-    need_header_ = true;
+    have_header_ = true;
   }
 
  private:
   uint64 auth_key_id_{0};
   string auth_key_;
   bool auth_flag_{false};
-  bool need_header_{true};
+  bool have_header_{true};
+  double header_expires_at_{0};
   double expires_at_{0};
   double created_at_{0};
 };