From 0450b05757dafae634f4a0949afa976273de4415 Mon Sep 17 00:00:00 2001
From: Arseny Smirnov <arseny30@gmail.com>
Date: Tue, 3 May 2022 16:33:09 +0400
Subject: [PATCH] Session: improve key immunity condition

---
 td/telegram/net/Session.cpp | 9 +++++++--
 td/telegram/net/Session.h   | 3 ++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/td/telegram/net/Session.cpp b/td/telegram/net/Session.cpp
index d7877f3a8..7553e957a 100644
--- a/td/telegram/net/Session.cpp
+++ b/td/telegram/net/Session.cpp
@@ -182,6 +182,7 @@ Session::Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared>
   }
   last_activity_timestamp_ = Time::now();
   last_success_timestamp_ = Time::now() - 366 * 86400;
+  last_bind_success_timestamp_ = Time::now() - 366 * 86400;
 }
 
 bool Session::can_destroy_auth_key() const {
@@ -286,7 +287,8 @@ void Session::on_bind_result(NetQueryPtr query) {
     if (status.code() == 400 && status.message() == "ENCRYPTED_MESSAGE_INVALID") {
       auto auth_key_age = G()->server_time() - auth_data_.get_main_auth_key().created_at();
       bool has_immunity = !G()->is_server_time_reliable() || auth_key_age < 60 ||
-                          (auth_key_age > 86400 && last_success_timestamp_ > Time::now() - 86400);
+                          (auth_key_age > 86400 &&
+                           (use_pfs_ ? last_bind_success_timestamp_ : last_success_timestamp_) > Time::now() - 86400);
       if (!use_pfs_) {
         if (has_immunity) {
           LOG(WARNING) << "Do not drop main key, because it was created too recently";
@@ -317,6 +319,7 @@ void Session::on_bind_result(NetQueryPtr query) {
   if (status.is_ok()) {
     LOG(INFO) << "Bound temp auth key " << auth_data_.get_tmp_auth_key().id();
     auth_data_.on_bind();
+    last_bind_success_timestamp_ = td::Time::now();
     on_tmp_auth_key_updated();
   } else if (status.error().message() == "DispatchTtlError") {
     LOG(INFO) << "Resend bind auth key " << auth_data_.get_tmp_auth_key().id() << " request after DispatchTtlError";
@@ -580,7 +583,9 @@ void Session::on_closed(Status status) {
 void Session::on_session_created(uint64 unique_id, uint64 first_id) {
   // TODO: use unique_id
   LOG(INFO) << "New session " << unique_id << " created with first message_id " << first_id;
-  last_success_timestamp_ = Time::now();
+  if (!use_pfs_) {
+    last_success_timestamp_ = Time::now();
+  }
   if (is_main_) {
     LOG(DEBUG) << "Sending updatesTooLong to force getDifference";
     BufferSlice packet(4);
diff --git a/td/telegram/net/Session.h b/td/telegram/net/Session.h
index 259307d02..16ce9a5b1 100644
--- a/td/telegram/net/Session.h
+++ b/td/telegram/net/Session.h
@@ -119,7 +119,8 @@ class Session final
   uint64 last_bind_query_id_ = 0;
   uint64 last_check_query_id_ = 0;
   double last_activity_timestamp_ = 0;
-  double last_success_timestamp_ = 0;  // time when auth_key and Session definitely was valid
+  double last_success_timestamp_ = 0;       // time when auth_key and Session definitely was valid
+  double last_bind_success_timestamp_ = 0;  // time when auth_key and Session definitely was valid and authorized
   size_t dropped_size_ = 0;
 
   FlatHashSet<uint64> unknown_queries_;