From ba0d2408f2c7bec69a5590e8e4ab4617f39718f0 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Fri, 1 Sep 2023 16:40:19 +0300
Subject: [PATCH] Improve handling of new_session_created.

---
 td/mtproto/PingConnection.cpp    |  2 +-
 td/mtproto/SessionConnection.cpp |  2 +-
 td/mtproto/SessionConnection.h   |  2 +-
 td/telegram/net/Session.cpp      | 19 ++++++++++++-------
 td/telegram/net/Session.h        |  2 +-
 5 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/td/mtproto/PingConnection.cpp b/td/mtproto/PingConnection.cpp
index d47b7dffc..3119123ec 100644
--- a/td/mtproto/PingConnection.cpp
+++ b/td/mtproto/PingConnection.cpp
@@ -122,7 +122,7 @@ class PingConnectionPingPong final
   void on_server_time_difference_updated(bool force) final {
   }
 
-  void on_session_created(uint64 unique_id, uint64 first_id) final {
+  void on_new_session_created(uint64 unique_id, uint64 first_id) final {
   }
   void on_session_failed(Status status) final {
   }
diff --git a/td/mtproto/SessionConnection.cpp b/td/mtproto/SessionConnection.cpp
index 54e4f2dea..885480fe0 100644
--- a/td/mtproto/SessionConnection.cpp
+++ b/td/mtproto/SessionConnection.cpp
@@ -306,7 +306,7 @@ Status SessionConnection::on_packet(const MsgInfo &info, const mtproto_api::new_
   VLOG(mtproto) << "NEW_SESSION_CREATED: [first_msg_id:" << format::as_hex(new_session_created.first_msg_id_)
                 << "] [unique_id:" << format::as_hex(new_session_created.unique_id_)
                 << "] [server_salt:" << format::as_hex(new_session_created.server_salt_) << "]";
-  callback_->on_session_created(new_session_created.unique_id_, new_session_created.first_msg_id_);
+  callback_->on_new_session_created(new_session_created.unique_id_, new_session_created.first_msg_id_);
   return Status::OK();
 }
 
diff --git a/td/mtproto/SessionConnection.h b/td/mtproto/SessionConnection.h
index 67a4821e9..2459ee94a 100644
--- a/td/mtproto/SessionConnection.h
+++ b/td/mtproto/SessionConnection.h
@@ -107,7 +107,7 @@ class SessionConnection final
     virtual void on_server_salt_updated() = 0;
     virtual void on_server_time_difference_updated(bool force) = 0;
 
-    virtual void on_session_created(uint64 unique_id, uint64 first_id) = 0;
+    virtual void on_new_session_created(uint64 unique_id, uint64 first_id) = 0;
     virtual void on_session_failed(Status status) = 0;
 
     virtual void on_container_sent(uint64 container_id, vector<uint64> msgs_id) = 0;
diff --git a/td/telegram/net/Session.cpp b/td/telegram/net/Session.cpp
index fe338d5bc..dc4b466ac 100644
--- a/td/telegram/net/Session.cpp
+++ b/td/telegram/net/Session.cpp
@@ -700,8 +700,7 @@ void Session::on_closed(Status status) {
   current_info_->state_ = ConnectionInfo::State::Empty;
 }
 
-void Session::on_session_created(uint64 unique_id, uint64 first_message_id) {
-  // TODO: use unique_id
+void Session::on_new_session_created(uint64 unique_id, uint64 first_message_id) {
   LOG(INFO) << "New session " << unique_id << " created with first message_id " << first_message_id;
   if (!use_pfs_ && !auth_data_.use_pfs()) {
     last_success_timestamp_ = Time::now();
@@ -713,16 +712,22 @@ void Session::on_session_created(uint64 unique_id, uint64 first_message_id) {
     last_activity_timestamp_ = Time::now();
     callback_->on_update(std::move(packet), auth_data_.get_auth_key().id());
   }
-
+  auto first_query_it = sent_queries_.find(first_message_id);
+  if (first_query_it != sent_queries_.end()) {
+    first_message_id = first_query_it->second.container_message_id_;
+    LOG(INFO) << "Update first_message_id to container's " << first_message_id;
+  } else {
+    LOG(ERROR) << "Failed to find query " << first_message_id << " from the new session";
+  }
   for (auto it = sent_queries_.begin(); it != sent_queries_.end();) {
     Query *query_ptr = &it->second;
     if (query_ptr->container_message_id_ < first_message_id) {
       // container vector leak otherwise
-      cleanup_container(it->first, &it->second);
-      mark_as_known(it->first, &it->second);
+      cleanup_container(it->first, query_ptr);
+      mark_as_known(it->first, query_ptr);
 
-      auto &query = it->second.net_query_;
-      VLOG(net_query) << "Resend query (on_session_created) " << query;
+      auto &query = query_ptr->net_query_;
+      VLOG(net_query) << "Resend from on_new_session_created " << query;
       query->set_message_id(0);
       query->cancel_slot_.clear_event();
       resend_query(std::move(query));
diff --git a/td/telegram/net/Session.h b/td/telegram/net/Session.h
index 3ac5f4e42..4bf381e68 100644
--- a/td/telegram/net/Session.h
+++ b/td/telegram/net/Session.h
@@ -214,7 +214,7 @@ class Session final
   void on_server_salt_updated() final;
   void on_server_time_difference_updated(bool force) final;
 
-  void on_session_created(uint64 unique_id, uint64 first_message_id) final;
+  void on_new_session_created(uint64 unique_id, uint64 first_message_id) final;
   void on_session_failed(Status status) final;
 
   void on_container_sent(uint64 container_message_id, vector<uint64> message_ids) final;