Handle -404 error in new ping

GitOrigin-RevId: 3cf88082f0fac9cb82027899c9d276e2d6c439fe

td/mtproto/PingConnection.cpp

@@ -89,11 +89,9 @@ class PingConnectionPingPong
   Status status_;
   void on_connected() override {
   }
-  void on_before_close() override {
-    Scheduler::unsubscribe_before_close(connection_->get_poll_info().get_pollable_fd_ref());
-  }
   void on_closed(Status status) override {
     is_closed_ = true;
+    CHECK(status.is_error());
     status_ = std::move(status);
   }
 
@@ -151,8 +149,10 @@ class PingConnectionPingPong
     if (was_pong()) {
       return Status::OK();
     }
+    CHECK(!is_closed_);
     connection_->flush(this);
     if (is_closed_) {
+      CHECK(status_.is_error());
       return std::move(status_);
     }
     return Status::OK();

td/mtproto/SessionConnection.cpp

@@ -735,8 +735,6 @@ void SessionConnection::set_online(bool online_flag, bool is_main) {
 
 void SessionConnection::do_close(Status status) {
   state_ = Closed;
-  callback_->on_before_close();
-  raw_connection_->close();
   // NB: this could be destroyed after on_closed
   callback_->on_closed(std::move(status));
 }

td/mtproto/SessionConnection.h

@@ -91,7 +91,6 @@ class SessionConnection
     virtual ~Callback() = default;
 
     virtual void on_connected() = 0;
-    virtual void on_before_close() = 0;
     virtual void on_closed(Status status) = 0;
 
     virtual void on_auth_key_updated() = 0;

td/telegram/net/ConnectionCreator.cpp

@@ -746,6 +746,7 @@ void ConnectionCreator::request_raw_connection(DcId dc_id, bool allow_media_only
     CHECK(client.is_media == is_media);
   }
   client.auth_data = std::move(auth_data);
+  client.auth_data_generation++;
   VLOG(connections) << "Request connection for " << tag("client", format::as_hex(client.hash)) << " to " << dc_id << " "
                     << tag("allow_media_only", allow_media_only);
   client.queries.push_back(std::move(promise));
@@ -1011,18 +1012,28 @@ void ConnectionCreator::client_loop(ClientInfo &client) {
 void ConnectionCreator::client_create_raw_connection(Result<ConnectionData> r_connection_data, bool check_mode,
                                                      mtproto::TransportType transport_type, size_t hash,
                                                      string debug_str, uint32 network_generation) {
-  auto promise = PromiseCreator::lambda([actor_id = actor_id(this), hash, check_mode,
+  unique_ptr<mtproto::AuthData> auth_data;
+  uint64 auth_data_generation{0};
+  if (check_mode) {
+    auto it = clients_.find(hash);
+    CHECK(it != clients_.end());
+    const auto &auth_data_ptr = it->second.auth_data;
+    if (auth_data_ptr && auth_data_ptr->use_pfs() && auth_data_ptr->has_auth_key(Time::now_cached())) {
+      auth_data = make_unique<mtproto::AuthData>(*auth_data_ptr);
+      auth_data_generation = it->second.auth_data_generation;
+    }
+  }
+  auto promise = PromiseCreator::lambda([actor_id = actor_id(this), hash, check_mode, auth_data_generation,
                                          debug_str](Result<unique_ptr<mtproto::RawConnection>> result) mutable {
     if (result.is_ok()) {
-      //FIXME
-      LOG(ERROR) << "RTT: " << result.ok()->rtt_;
       VLOG(connections) << "Ready connection (" << (check_mode ? "" : "un") << "checked) " << result.ok().get() << ' '
-                        << debug_str;
+                        << tag("rtt", format::as_time(result.ok()->rtt_)) << ' ' << debug_str;
     } else {
       VLOG(connections) << "Failed connection (" << (check_mode ? "" : "un") << "checked) " << result.error() << ' '
                         << debug_str;
     }
-    send_closure(std::move(actor_id), &ConnectionCreator::client_add_connection, hash, std::move(result), check_mode);
+    send_closure(std::move(actor_id), &ConnectionCreator::client_add_connection, hash, std::move(result), check_mode,
+                 auth_data_generation);
   });
 
   if (r_connection_data.is_error()) {
@@ -1038,26 +1049,8 @@ void ConnectionCreator::client_create_raw_connection(Result<ConnectionData> r_co
   raw_connection->debug_str_ = debug_str;
 
   if (check_mode) {
-    VLOG(connections) << "Start check: " << debug_str;
+    VLOG(connections) << "Start check: " << debug_str << " " << (auth_data ? "with" : "without") << " auth data";
     auto token = next_token();
-    auto it = clients_.find(hash);
-    CHECK(it != clients_.end());
-    const auto &auth_data_ptr = it->second.auth_data;
-    unique_ptr<mtproto::AuthData> auth_data;
-    if (auth_data_ptr && auth_data_ptr->use_pfs() && auth_data_ptr->has_auth_key(Time::now_cached())) {
-      auth_data = make_unique<mtproto::AuthData>(*auth_data_ptr);
-      // FIXME
-      LOG(ERROR) << "use auth_data";
-    } else {
-      if (auth_data_ptr) {
-        // FIXME
-        LOG(ERROR) << "do not use auth data " << !!auth_data_ptr << " " << auth_data_ptr->use_pfs() << " "
-                   << auth_data_ptr->has_auth_key(Time::now_cached());
-      } else {
-        // FIXME
-        LOG(ERROR) << "do not use auth data";
-      }
-    }
     children_[token] = {true, create_ping_actor(debug_str, std::move(raw_connection), std::move(auth_data),
                                                 std::move(promise), create_reference(token))};
   } else {
@@ -1075,7 +1068,7 @@ void ConnectionCreator::client_set_timeout_at(ClientInfo &client, double wakeup_
 }
 
 void ConnectionCreator::client_add_connection(size_t hash, Result<unique_ptr<mtproto::RawConnection>> r_raw_connection,
-                                              bool check_flag) {
+                                              bool check_flag, uint64 auth_data_generation) {
   auto &client = clients_[hash];
   CHECK(client.pending_connections > 0);
   client.pending_connections--;
@@ -1088,6 +1081,13 @@ void ConnectionCreator::client_add_connection(size_t hash, Result<unique_ptr<mtp
                       << tag("client", format::as_hex(hash));
     client.backoff.clear();
     client.ready_connections.push_back(std::make_pair(r_raw_connection.move_as_ok(), Time::now_cached()));
+  } else {
+    if (r_raw_connection.error().code() == -404 && client.auth_data &&
+        client.auth_data_generation == auth_data_generation) {
+      VLOG(connections) << "Drop auth data from " << tag("client", format::as_hex(hash));
+      client.auth_data = nullptr;
+      client.auth_data_generation++;
+    }
   }
   client_loop(client);
 }

td/telegram/net/ConnectionCreator.h

@@ -240,6 +240,7 @@ class ConnectionCreator : public NetQueryCallback {
     bool allow_media_only;
     bool is_media;
     unique_ptr<mtproto::AuthData> auth_data;
+    uint64 auth_data_generation{0};
   };
   std::map<size_t, ClientInfo> clients_;
 
@@ -301,7 +302,8 @@ class ConnectionCreator : public NetQueryCallback {
   void client_create_raw_connection(Result<ConnectionData> r_connection_data, bool check_mode,
                                     mtproto::TransportType transport_type, size_t hash, string debug_str,
                                     uint32 network_generation);
-  void client_add_connection(size_t hash, Result<unique_ptr<mtproto::RawConnection>> r_raw_connection, bool check_flag);
+  void client_add_connection(size_t hash, Result<unique_ptr<mtproto::RawConnection>> r_raw_connection, bool check_flag,
+                             uint64 auth_data_generation);
   void client_set_timeout_at(ClientInfo &client, double wakeup_at);
 
   void on_get_proxy_info(telegram_api::object_ptr<telegram_api::help_ProxyData> proxy_data_ptr);

td/telegram/net/Session.cpp

@@ -401,14 +401,13 @@ void Session::on_server_time_difference_updated() {
   shared_auth_data_->update_server_time_difference(auth_data_.get_server_time_difference());
 }
 
-void Session::on_before_close() {
-  Scheduler::unsubscribe_before_close(current_info_->connection->get_poll_info().get_pollable_fd_ref());
-}
-
 void Session::on_closed(Status status) {
   if (!close_flag_ && is_main_) {
     connection_token_.reset();
   }
+  auto raw_connection = current_info_->connection->move_as_raw_connection();
+  Scheduler::unsubscribe_before_close(raw_connection->get_poll_info().get_pollable_fd_ref());
+  raw_connection->close();
 
   if (status.is_error()) {
     LOG(WARNING) << "Session closed: " << status << " " << current_info_->connection->get_name();

td/telegram/net/Session.h

@@ -174,7 +174,6 @@ class Session final
 
   // mtproto::Connection::Callback
   void on_connected() override;
-  void on_before_close() override;
   void on_closed(Status status) override;
 
   Status on_pong() override;