From 723d09ddee66601fee0d3ce5d5d748ee8575f509 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Mon, 17 Jul 2023 19:21:18 +0300
Subject: [PATCH] Reduce refresh margin for persistent temp auth keys to speed
 up recover when the old key reregisters after auth.dropTempAuthKeys.

---
 td/mtproto/AuthData.h            | 4 ++--
 td/telegram/ConfigManager.cpp    | 8 ++++----
 td/telegram/net/Session.cpp      | 8 +++++---
 td/telegram/net/Session.h        | 3 ++-
 td/telegram/net/SessionProxy.cpp | 4 ++--
 5 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/td/mtproto/AuthData.h b/td/mtproto/AuthData.h
index 251cecd40..bc2403508 100644
--- a/td/mtproto/AuthData.h
+++ b/td/mtproto/AuthData.h
@@ -89,14 +89,14 @@ class AuthData {
   bool was_tmp_auth_key() const {
     return use_pfs() && !tmp_auth_key_.empty();
   }
-  bool need_tmp_auth_key(double now) const {
+  bool need_tmp_auth_key(double now, double refresh_margin) const {
     if (!use_pfs()) {
       return false;
     }
     if (tmp_auth_key_.empty()) {
       return true;
     }
-    if (now > tmp_auth_key_.expires_at() - 60 * 60 * 2 /*2 hours*/) {
+    if (now > tmp_auth_key_.expires_at() - refresh_margin) {
       return true;
     }
     return false;
diff --git a/td/telegram/ConfigManager.cpp b/td/telegram/ConfigManager.cpp
index 2d1c1f331..5a22f2a88 100644
--- a/td/telegram/ConfigManager.cpp
+++ b/td/telegram/ConfigManager.cpp
@@ -523,10 +523,10 @@ static ActorOwn<> get_full_config(DcOption option, Promise<tl_object_ptr<telegra
       if (G()->is_test_dc()) {
         int_dc_id += 10000;
       }
-      session_ = create_actor<Session>("ConfigSession", std::move(session_callback), std::move(auth_data), raw_dc_id,
-                                       int_dc_id, false /*is_primary*/, false /*is_main*/, true /*use_pfs*/,
-                                       false /*is_cdn*/, false /*need_destroy_auth_key*/, mtproto::AuthKey(),
-                                       std::vector<mtproto::ServerSalt>());
+      session_ = create_actor<Session>(
+          "ConfigSession", std::move(session_callback), std::move(auth_data), raw_dc_id, int_dc_id,
+          false /*is_primary*/, false /*is_main*/, true /*use_pfs*/, false /*persist_tmp_auth_key*/, false /*is_cdn*/,
+          false /*need_destroy_auth_key*/, mtproto::AuthKey(), std::vector<mtproto::ServerSalt>());
       auto query = G()->net_query_creator().create_unauth(telegram_api::help_getConfig(), DcId::empty());
       query->total_timeout_limit_ = 60 * 60 * 24;
       query->set_callback(actor_shared(this));
diff --git a/td/telegram/net/Session.cpp b/td/telegram/net/Session.cpp
index 0f4602d00..c374ac0e1 100644
--- a/td/telegram/net/Session.cpp
+++ b/td/telegram/net/Session.cpp
@@ -232,12 +232,14 @@ bool Session::PriorityQueue::empty() const {
 }
 
 Session::Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 raw_dc_id,
-                 int32 dc_id, bool is_primary, bool is_main, bool use_pfs, bool is_cdn, bool need_destroy,
-                 const mtproto::AuthKey &tmp_auth_key, const vector<mtproto::ServerSalt> &server_salts)
+                 int32 dc_id, bool is_primary, bool is_main, bool use_pfs, bool persist_tmp_auth_key, bool is_cdn,
+                 bool need_destroy, const mtproto::AuthKey &tmp_auth_key,
+                 const vector<mtproto::ServerSalt> &server_salts)
     : raw_dc_id_(raw_dc_id)
     , dc_id_(dc_id)
     , is_primary_(is_primary)
     , is_main_(is_main)
+    , persist_tmp_auth_key_(use_pfs && persist_tmp_auth_key)
     , is_cdn_(is_cdn)
     , need_destroy_(need_destroy) {
   VLOG(dc) << "Start connection " << tag("need_destroy", need_destroy_);
@@ -1504,7 +1506,7 @@ void Session::auth_loop(double now) {
   if (auth_data_.need_main_auth_key()) {
     create_gen_auth_key_actor(MainAuthKeyHandshake);
   }
-  if (auth_data_.need_tmp_auth_key(now)) {
+  if (auth_data_.need_tmp_auth_key(now, persist_tmp_auth_key_ ? 2 * 60 : 60 * 60)) {
     create_gen_auth_key_actor(TmpAuthKeyHandshake);
   }
 }
diff --git a/td/telegram/net/Session.h b/td/telegram/net/Session.h
index b7f0f437b..ff68a634f 100644
--- a/td/telegram/net/Session.h
+++ b/td/telegram/net/Session.h
@@ -67,7 +67,7 @@ class Session final
   };
 
   Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 raw_dc_id, int32 dc_id,
-          bool is_primary, bool is_main, bool use_pfs, bool is_cdn, bool need_destroy,
+          bool is_primary, bool is_main, bool use_pfs, bool persist_tmp_auth_key, bool is_cdn, bool need_destroy,
           const mtproto::AuthKey &tmp_auth_key, const vector<mtproto::ServerSalt> &server_salts);
 
   void send(NetQueryPtr &&query);
@@ -108,6 +108,7 @@ class Session final
   const int32 dc_id_;      // unique datacenter ID, i.e. -10002
   const bool is_primary_;  // true for primary Sessions to all DCs
   const bool is_main_;     // true only for the primary Session(s) to the main DC
+  const bool persist_tmp_auth_key_;
   const bool is_cdn_;
   const bool need_destroy_;
   bool was_on_network_ = false;
diff --git a/td/telegram/net/SessionProxy.cpp b/td/telegram/net/SessionProxy.cpp
index c4283e114..36f4e702e 100644
--- a/td/telegram/net/SessionProxy.cpp
+++ b/td/telegram/net/SessionProxy.cpp
@@ -236,8 +236,8 @@ void SessionProxy::open_session(bool force) {
   session_ = create_actor<Session>(
       name,
       make_unique<SessionCallback>(actor_shared(this, session_generation_), dc_id, allow_media_only_, is_media_, hash),
-      auth_data_, raw_dc_id, int_dc_id, is_primary_, is_main_, use_pfs_, is_cdn_, need_destroy_, tmp_auth_key_,
-      server_salts_);
+      auth_data_, raw_dc_id, int_dc_id, is_primary_, is_main_, use_pfs_, persist_tmp_auth_key_, is_cdn_, need_destroy_,
+      tmp_auth_key_, server_salts_);
 }
 
 void SessionProxy::update_auth_key_state() {