From 7471e62740598a2e0e55b90587926c86aabc7601 Mon Sep 17 00:00:00 2001
From: Arseny Smirnov <arseny30@gmail.com>
Date: Tue, 6 Nov 2018 14:06:08 +0300
Subject: [PATCH] Destroy authorization keys: fixes

GitOrigin-RevId: 9538efa2deadda91c18c1c97aa22281597a44d4b
---
 td/telegram/Td.cpp                    |  2 +-
 td/telegram/net/Session.cpp           | 12 +++---------
 td/telegram/net/Session.h             |  1 -
 td/telegram/net/SessionMultiProxy.cpp |  6 ++----
 td/telegram/net/SessionProxy.cpp      | 13 +++++++------
 5 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/td/telegram/Td.cpp b/td/telegram/Td.cpp
index 69da0c15..aaa5a973 100644
--- a/td/telegram/Td.cpp
+++ b/td/telegram/Td.cpp
@@ -3521,7 +3521,7 @@ void Td::on_connection_state_changed(StateManager::State new_state) {
 
 void Td::on_authorization_lost() {
   LOG(WARNING) << "on_authorization_lost";
-  destroy();
+  send_closure(auth_manager_actor_, &AuthManager::on_authorization_lost);
 }
 
 void Td::start_up() {
diff --git a/td/telegram/net/Session.cpp b/td/telegram/net/Session.cpp
index 40a7f87f..42cdfb7d 100644
--- a/td/telegram/net/Session.cpp
+++ b/td/telegram/net/Session.cpp
@@ -111,6 +111,7 @@ Session::Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared>
                  bool is_main, bool use_pfs, bool is_cdn, bool need_destroy, const mtproto::AuthKey &tmp_auth_key,
                  std::vector<mtproto::ServerSalt> server_salts)
     : dc_id_(dc_id), is_main_(is_main), is_cdn_(is_cdn) {
+  LOG(INFO) << "Open session: ";
   VLOG(dc) << "Start connection";
   need_destroy_ = need_destroy;
   if (need_destroy) {
@@ -287,6 +288,7 @@ void Session::flush_pending_invoke_after_queries() {
 }
 
 void Session::close() {
+  LOG(INFO) << "Close session (external)";
   close_flag_ = true;
   connection_close(&main_connection_);
   connection_close(&long_poll_connection_);
@@ -669,7 +671,7 @@ void Session::on_message_result_error(uint64 id, int error_code, BufferSlice mes
       LOG(WARNING) << "Lost authorization due to " << tag("msg", message.as_slice());
       auth_data_.set_auth_flag(false);
       shared_auth_data_->set_auth_key(auth_data_.get_main_auth_key());
-      auth_lost_flag_ = true;
+      on_session_failed(Status::OK());
     }
   }
 
@@ -1166,7 +1168,6 @@ void Session::loop() {
   main_connection_.wakeup_at = 0;
   long_poll_connection_.wakeup_at = 0;
 
-  auth_lost_flag_ = false;
   // NB: order is crucial. First long_poll_connection, then main_connection
   // Otherwise queries could be sent with big delay
 
@@ -1219,13 +1220,6 @@ void Session::loop() {
     connection_open(&main_connection_, true /*send ask_info*/);
   }
 
-  if (auth_lost_flag_) {
-    connection_close(&main_connection_);
-    connection_close(&long_poll_connection_);
-    auth_lost_flag_ = false;
-    relax_timeout_at(&wakeup_at, Time::now_cached() + 0.1);
-  }
-
   relax_timeout_at(&wakeup_at, main_connection_.wakeup_at);
 
   double wakeup_in = 0;
diff --git a/td/telegram/net/Session.h b/td/telegram/net/Session.h
index da042fa7..5323c04f 100644
--- a/td/telegram/net/Session.h
+++ b/td/telegram/net/Session.h
@@ -136,7 +136,6 @@ class Session final
   ConnectionInfo *current_info_;
   ConnectionInfo main_connection_;
   ConnectionInfo long_poll_connection_;
-  bool auth_lost_flag_ = false;
   StateManager::ConnectionToken connection_token_;
 
   double cached_connection_timestamp_ = 0;
diff --git a/td/telegram/net/SessionMultiProxy.cpp b/td/telegram/net/SessionMultiProxy.cpp
index 4eb72e5d..a2de3d30 100644
--- a/td/telegram/net/SessionMultiProxy.cpp
+++ b/td/telegram/net/SessionMultiProxy.cpp
@@ -57,9 +57,7 @@ void SessionMultiProxy::update_main_flag(bool is_main) {
 
 void SessionMultiProxy::update_destroy_auth_key(bool need_destroy_auth_key) {
   need_destroy_auth_key_ = need_destroy_auth_key;
-  for (auto &session : sessions_) {
-    send_closure(session, &SessionProxy::update_destroy, need_destroy_auth_key_);
-  }
+  send_closure(sessions_[0], &SessionProxy::update_destroy, need_destroy_auth_key_);
 }
 void SessionMultiProxy::update_session_count(int32 session_count) {
   update_options(session_count, use_pfs_);
@@ -120,7 +118,7 @@ void SessionMultiProxy::init() {
                             << format::cond(session_count_ > 1, format::concat("#", i));
     sessions_.push_back(create_actor<SessionProxy>(name, auth_data_, is_main_, allow_media_only_, is_media_,
                                                    get_pfs_flag(), is_main_ && i != 0, is_cdn_,
-                                                   need_destroy_auth_key_));
+                                                   need_destroy_auth_key_ && i == 0));
   }
 }
 
diff --git a/td/telegram/net/SessionProxy.cpp b/td/telegram/net/SessionProxy.cpp
index 58d34d59..3448544a 100644
--- a/td/telegram/net/SessionProxy.cpp
+++ b/td/telegram/net/SessionProxy.cpp
@@ -165,13 +165,10 @@ void SessionProxy::open_session(bool force) {
     if (need_destroy_) {
       return auth_state_ != AuthState::Empty;
     }
-    if (is_main_) {  // always open main for ordinary queries
-      return true;
+    if (auth_state_ != AuthState::OK) {
+      return false;
     }
-    if (!pending_queries_.empty() && auth_state_ == AuthState::OK) {
-      return true;
-    }
-    return false;
+    return is_main_ || !pending_queries_.empty();
   }();
   if (!should_open) {
     return;
@@ -196,7 +193,11 @@ void SessionProxy::open_session(bool force) {
 }
 
 void SessionProxy::update_auth_state() {
+  auto old_auth_state = auth_state_;
   auth_state_ = auth_data_->get_auth_state().first;
+  if (auth_state_ != old_auth_state && old_auth_state == AuthState::OK) {
+    close_session();
+  }
   open_session();
   if (session_.empty() || auth_state_ != AuthState::OK) {
     return;