QR code authentication fixes.

GitOrigin-RevId: 29108234fa8afc589ac7248ad1fdc92ec1d0ffaf

td/telegram/AuthManager.cpp

@@ -543,14 +543,14 @@ void AuthManager::on_get_login_token(tl_object_ptr<telegram_api::auth_LoginToken
 
 void AuthManager::on_get_password_result(NetQueryPtr &result) {
   auto r_password = fetch_result<telegram_api::account_getPassword>(result->ok());
-  if (r_password.is_error()) {
+  if (r_password.is_error() && query_id_ != 0) {
     return on_query_error(r_password.move_as_error());
   }
-  auto password = r_password.move_as_ok();
+  auto password = r_password.is_ok() ? r_password.move_as_ok() : nullptr;
   LOG(INFO) << "Receive password info: " << to_string(password);
 
   wait_password_state_ = WaitPasswordState();
-  if (password->current_algo_ != nullptr) {
+  if (password != nullptr && password->current_algo_ != nullptr) {
     switch (password->current_algo_->get_id()) {
       case telegram_api::passwordKdfAlgoUnknown::ID:
         return on_query_error(Status::Error(400, "Application update is needed to log in"));
@@ -572,7 +572,10 @@ void AuthManager::on_get_password_result(NetQueryPtr &result) {
         UNREACHABLE();
     }
   } else if (was_qr_code_request_) {
-    send_export_login_token_query();
+    imported_dc_id_ = -1;
+    login_code_retry_delay_ = td::clamp(2 * login_code_retry_delay_, 1, 60);
+    set_login_token_expires_at(Time::now() + login_code_retry_delay_);
+    return;
   } else {
     start_net_query(NetQueryType::SignIn,
                     G()->net_query_creator().create(
@@ -582,6 +585,11 @@ void AuthManager::on_get_password_result(NetQueryPtr &result) {
     return;
   }
 
+  if (imported_dc_id_ != -1) {
+    G()->net_query_dispatcher().set_main_dc_id(imported_dc_id_);
+    imported_dc_id_ = -1;
+  }
+
   if (state_ == State::WaitPassword) {
     LOG(INFO) << "Have SRP id " << wait_password_state_.srp_id_;
     auto hash = PasswordManager::get_input_check_password(password_, wait_password_state_.current_client_salt_,
@@ -594,7 +602,9 @@ void AuthManager::on_get_password_result(NetQueryPtr &result) {
                                                     DcId::main(), NetQuery::Type::Common, NetQuery::AuthFlag::Off));
   } else {
     update_state(State::WaitPassword);
-    on_query_ok();
+    if (query_id_ != 0) {
+      on_query_ok();
+    }
   }
 }
 
@@ -759,9 +769,14 @@ void AuthManager::on_result(NetQueryPtr result) {
     if (result->is_error()) {
       if ((type == NetQueryType::SignIn || type == NetQueryType::RequestQrCode || type == NetQueryType::ImportQrCode) &&
           result->error().code() == 401 && result->error().message() == CSlice("SESSION_PASSWORD_NEEDED")) {
+        auto dc_id = DcId::main();
+        if (type == NetQueryType::ImportQrCode) {
+          CHECK(DcId::is_valid(imported_dc_id_));
+          dc_id = DcId::internal(imported_dc_id_);
+        }
         start_net_query(NetQueryType::GetPassword,
-                        G()->net_query_creator().create(create_storer(telegram_api::account_getPassword()),
+                        G()->net_query_creator().create(create_storer(telegram_api::account_getPassword()), dc_id,
-                                                        DcId::main(), NetQuery::Type::Common, NetQuery::AuthFlag::Off));
+                                                        NetQuery::Type::Common, NetQuery::AuthFlag::Off));
         return;
       }
       if (result->error().message() == CSlice("PHONE_NUMBER_BANNED")) {

td/telegram/ConfigManager.cpp

@@ -488,16 +488,17 @@ ActorOwn<> get_full_config(DcOption option, Promise<FullConfig> promise, ActorSh
    private:
     void start_up() override {
       auto auth_data = std::make_shared<SimpleAuthData>(option_.get_dc_id());
-      int32 int_dc_id = option_.get_dc_id().get_raw_id();
+      int32 raw_dc_id = option_.get_dc_id().get_raw_id();
       auto session_callback = make_unique<SessionCallback>(actor_shared(this, 1), std::move(option_));
 
+      int32 int_dc_id = raw_dc_id;
       if (G()->is_test_dc()) {
         int_dc_id += 10000;
       }
-      session_ =
+      session_ = create_actor<Session>("ConfigSession", std::move(session_callback), std::move(auth_data), raw_dc_id,
-          create_actor<Session>("ConfigSession", std::move(session_callback), std::move(auth_data), int_dc_id,
+                                       int_dc_id, false /*is_main*/, true /*use_pfs*/, false /*is_cdn*/,
-                                false /*is_main*/, true /*use_pfs*/, false /*is_cdn*/, false /*need_destroy_auth_key*/,
+                                       false /*need_destroy_auth_key*/, mtproto::AuthKey(),
-                                mtproto::AuthKey(), std::vector<mtproto::ServerSalt>());
+                                       std::vector<mtproto::ServerSalt>());
       auto query = G()->net_query_creator().create(create_storer(telegram_api::help_getConfig()), DcId::empty(),
                                                    NetQuery::Type::Common, NetQuery::AuthFlag::Off,
                                                    NetQuery::GzipFlag::On, 60 * 60 * 24);

td/telegram/cli.cpp

@@ -1341,7 +1341,7 @@ class CliClient final : public Actor {
       send_request(td_api::make_object<td_api::recoverAuthenticationPassword>(args));
     } else if (op == "lo" || op == "LogOut" || op == "logout") {
       send_request(td_api::make_object<td_api::logOut>());
-    } else if (op == "ra" || op == "destroy") {
+    } else if (op == "destroy") {
       send_request(td_api::make_object<td_api::destroy>());
     } else if (op == "reset") {
       td_client_.reset();

td/telegram/net/DcAuthManager.cpp

@@ -34,8 +34,9 @@ DcAuthManager::DcAuthManager(ActorShared<> parent) {
     auto main_dc_id = to_integer<int32>(s_main_dc_id);
     if (DcId::is_valid(main_dc_id)) {
       main_dc_id_ = DcId::internal(main_dc_id);
+      VLOG(dc) << "Init main DcId to " << main_dc_id_;
     } else {
-      LOG(ERROR) << "Receive invalid main dc id " << main_dc_id;
+      LOG(ERROR) << "Receive invalid main DcId " << main_dc_id;
     }
   }
 }
@@ -64,9 +65,12 @@ void DcAuthManager::add_dc(std::shared_ptr<AuthDataShared> auth_data) {
   info.shared_auth_data = std::move(auth_data);
   auto state_was_auth = info.shared_auth_data->get_auth_key_state();
   info.auth_key_state = state_was_auth.first;
+  VLOG(dc) << "Add " << info.dc_id << " with auth key state " << info.auth_key_state
+           << " and was_auth = " << state_was_auth.second;
   was_auth_ |= state_was_auth.second;
   if (!main_dc_id_.is_exact()) {
     main_dc_id_ = info.dc_id;
+    VLOG(dc) << "Set main DcId to " << main_dc_id_;
   }
   info.shared_auth_data->add_auth_key_listener(make_unique<Listener>(actor_shared(this, info.dc_id.get_raw_id())));
   dcs_.emplace_back(std::move(info));
@@ -75,6 +79,7 @@ void DcAuthManager::add_dc(std::shared_ptr<AuthDataShared> auth_data) {
 
 void DcAuthManager::update_main_dc(DcId new_main_dc_id) {
   main_dc_id_ = new_main_dc_id;
+  VLOG(dc) << "Update main DcId to " << main_dc_id_;
   loop();
 }
 
@@ -95,8 +100,8 @@ void DcAuthManager::update_auth_key_state() {
   int32 dc_id = narrow_cast<int32>(get_link_token());
   auto &dc = get_dc(dc_id);
   auto state_was_auth = dc.shared_auth_data->get_auth_key_state();
-  VLOG(dc) << "Update DC auth key state " << tag("dc_id", dc_id) << tag("old_auth_key_state", dc.auth_key_state)
+  VLOG(dc) << "Update " << dc_id << " auth key state from " << dc.auth_key_state << " to " << state_was_auth.first
-           << tag("new_auth_key_state", state_was_auth.first);
+           << " with was_auth = " << state_was_auth.second;
   dc.auth_key_state = state_was_auth.first;
   was_auth_ |= state_was_auth.second;
 
@@ -213,10 +218,10 @@ void DcAuthManager::destroy_loop() {
   }
 
   if (is_ready) {
-    LOG(INFO) << "Destroy auth keys loop is ready, all keys are destroyed";
+    VLOG(dc) << "Destroy auth keys loop is ready, all keys are destroyed";
     destroy_promise_.set_value(Unit());
   } else {
-    LOG(INFO) << "DC is not ready for destroying auth key";
+    VLOG(dc) << "DC is not ready for destroying auth key";
   }
 }
 
@@ -232,11 +237,13 @@ void DcAuthManager::loop() {
   }
   auto main_dc = find_dc(main_dc_id_.get_raw_id());
   if (!main_dc || main_dc->auth_key_state != AuthKeyState::OK) {
+    VLOG(dc) << "Main is " << main_dc_id_ << ", main auth key state is "
+             << (main_dc ? main_dc->auth_key_state : AuthKeyState::Empty) << ", was_auth = " << was_auth_;
     if (was_auth_) {
       G()->shared_config().set_option_boolean("auth", false);
       destroy_loop();
     }
-    VLOG(dc) << "Skip loop because auth state of main dc " << main_dc_id_.get_raw_id() << " is "
+    VLOG(dc) << "Skip loop because auth state of main DcId " << main_dc_id_.get_raw_id() << " is "
              << (main_dc != nullptr ? (PSTRING() << main_dc->auth_key_state) : "unknown");
 
     return;

td/telegram/net/NetQueryDelayer.cpp

@@ -36,7 +36,7 @@ void NetQueryDelayer::delay(NetQueryPtr query) {
     for (auto prefix :
          {Slice("FLOOD_WAIT_"), Slice("SLOWMODE_WAIT_"), Slice("2FA_CONFIRM_WAIT_"), Slice("TAKEOUT_INIT_DELAY_")}) {
       if (begins_with(msg, prefix)) {
-        timeout = clamp(to_integer<int>(msg.substr(prefix.size())), 0, 14 * 24 * 60 * 60);
+        timeout = clamp(to_integer<int>(msg.substr(prefix.size())), 1, 14 * 24 * 60 * 60);
         break;
       }
     }

td/telegram/net/Session.cpp

@@ -115,10 +115,10 @@ class GenAuthKeyActor : public Actor {
 
 }  // namespace detail
 
-Session::Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 dc_id,
+Session::Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 raw_dc_id,
-                 bool is_main, bool use_pfs, bool is_cdn, bool need_destroy, const mtproto::AuthKey &tmp_auth_key,
+                 int32 dc_id, bool is_main, bool use_pfs, bool is_cdn, bool need_destroy,
-                 std::vector<mtproto::ServerSalt> server_salts)
+                 const mtproto::AuthKey &tmp_auth_key, std::vector<mtproto::ServerSalt> server_salts)
-    : dc_id_(dc_id), is_main_(is_main), is_cdn_(is_cdn) {
+    : raw_dc_id_(raw_dc_id), dc_id_(dc_id), is_main_(is_main), is_cdn_(is_cdn) {
   VLOG(dc) << "Start connection";
   need_destroy_ = need_destroy;
   if (need_destroy) {
@@ -141,7 +141,8 @@ Session::Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared>
   } while (session_id == 0);
   auth_data_.set_session_id(session_id);
   LOG(WARNING) << "Generate new session_id " << session_id << " for " << (use_pfs ? "temp " : "")
-               << (is_cdn ? "CDN " : "") << "auth key " << auth_data_.get_auth_key().id() << " for DC" << dc_id;
+               << (is_cdn ? "CDN " : "") << "auth key " << auth_data_.get_auth_key().id() << " for "
+               << (is_main_ ? "main " : "") << "DC" << dc_id;
 
   callback_ = std::shared_ptr<Callback>(callback.release());
 
@@ -466,8 +467,7 @@ void Session::on_closed(Status status) {
 void Session::on_session_created(uint64 unique_id, uint64 first_id) {
   // TODO: use unique_id
   // send updatesTooLong to force getDifference
-  LOG(INFO) << "New session " << unique_id << " created "
+  LOG(INFO) << "New session " << unique_id << " created with first message_id " << first_id;
-            << " with first message_id " << first_id;
   if (is_main_) {
     LOG(DEBUG) << "Sending updatesTooLong to force getDifference";
     telegram_api::updatesTooLong too_long_;
@@ -608,16 +608,6 @@ void Session::mark_as_unknown(uint64 id, Query *query) {
 Status Session::on_message_result_ok(uint64 id, BufferSlice packet, size_t original_size) {
   // Steal authorization information.
   // It is a dirty hack, yep.
-  TlParser parser(packet.as_slice());
-  int32 ID = parser.fetch_int();
-  if (!parser.get_error()) {
-    if (ID == telegram_api::auth_authorization::ID) {
-      LOG(INFO) << "GOT AUTHORIZATION!";
-      auth_data_.set_auth_flag(true);
-      shared_auth_data_->set_auth_key(auth_data_.get_main_auth_key());
-    }
-  }
-
   if (id == 0) {
     if (is_cdn_) {
       return Status::Error("Got update from CDN connection");
@@ -625,9 +615,13 @@ Status Session::on_message_result_ok(uint64 id, BufferSlice packet, size_t origi
     return_query(G()->net_query_creator().create_result(0, std::move(packet)));
     return Status::OK();
   }
+
+  TlParser parser(packet.as_slice());
+  int32 ID = parser.fetch_int();
+
   auto it = sent_queries_.find(id);
   if (it == sent_queries_.end()) {
-    LOG(DEBUG) << "DROP result to " << tag("request_id", format::as_hex(id)) << tag("tl", format::as_hex(ID));
+    LOG(DEBUG) << "Drop result to " << tag("request_id", format::as_hex(id)) << tag("tl", format::as_hex(ID));
 
     if (packet.size() > 16 * 1024) {
       dropped_size_ += packet.size();
@@ -640,10 +634,21 @@ Status Session::on_message_result_ok(uint64 id, BufferSlice packet, size_t origi
     }
     return Status::OK();
   }
+
   auth_data_.on_api_response();
   Query *query_ptr = &it->second;
   VLOG(net_query) << "Return query result " << query_ptr->query;
 
+  if (!parser.get_error()) {
+    if (ID == telegram_api::auth_authorization::ID || ID == telegram_api::auth_loginTokenSuccess::ID) {
+      if (query_ptr->query->tl_constructor() != telegram_api::auth_importAuthorization::ID) {
+        G()->net_query_dispatcher().set_main_dc_id(raw_dc_id_);
+      }
+      auth_data_.set_auth_flag(true);
+      shared_auth_data_->set_auth_key(auth_data_.get_main_auth_key());
+    }
+  }
+
   cleanup_container(id, query_ptr);
   mark_as_known(id, query_ptr);
   query_ptr->query->on_net_read(original_size);

td/telegram/net/Session.h

@@ -65,8 +65,8 @@ class Session final
     virtual void on_result(NetQueryPtr net_query) = 0;
   };
 
-  Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 dc_id, bool is_main,
+  Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 raw_dc_id, int32 dc_id,
-          bool use_pfs, bool is_cdn, bool need_destroy, const mtproto::AuthKey &tmp_auth_key,
+          bool is_main, bool use_pfs, bool is_cdn, bool need_destroy, const mtproto::AuthKey &tmp_auth_key,
           std::vector<mtproto::ServerSalt> server_salts);
   void send(NetQueryPtr &&query);
   void on_network(bool network_flag, uint32 network_generation);
@@ -101,6 +101,7 @@ class Session final
   // Just re-ask answer_id each time we get information about it.
   // Though mtproto::Connection must ensure delivery of such query.
 
+  int32 raw_dc_id_;
   int32 dc_id_;
   enum class Mode : int8 { Tcp, Http } mode_ = Mode::Tcp;
   bool is_main_;

td/telegram/net/SessionProxy.cpp

@@ -192,7 +192,8 @@ void SessionProxy::open_session(bool force) {
   string name = PSTRING() << "Session" << get_name().substr(Slice("SessionProxy").size());
   string hash_string = PSTRING() << name << " " << dc_id.get_raw_id() << " " << allow_media_only_;
   auto hash = std::hash<std::string>()(hash_string);
-  int32 int_dc_id = dc_id.get_raw_id();
+  int32 raw_dc_id = dc_id.get_raw_id();
+  int32 int_dc_id = raw_dc_id;
   if (G()->is_test_dc()) {
     int_dc_id += 10000;
   }
@@ -202,7 +203,7 @@ 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_, int_dc_id, is_main_, use_pfs_, is_cdn_, need_destroy_, tmp_auth_key_, server_salts_);
+      auth_data_, raw_dc_id, int_dc_id, is_main_, use_pfs_, is_cdn_, need_destroy_, tmp_auth_key_, server_salts_);
 }
 
 void SessionProxy::update_auth_key_state() {