From 83d9f84c5b88b3c9f3ab0e889f1b0bc3f0fa9b73 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Sun, 3 Sep 2023 02:07:12 +0300
Subject: [PATCH] Improve new_session_created handling.

---
 td/mtproto/SessionConnection.cpp | 22 ++++++++++++++--------
 td/mtproto/SessionConnection.h   |  1 +
 td/telegram/net/Session.cpp      |  6 +++---
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/td/mtproto/SessionConnection.cpp b/td/mtproto/SessionConnection.cpp
index 1ba8a6973..da51e9655 100644
--- a/td/mtproto/SessionConnection.cpp
+++ b/td/mtproto/SessionConnection.cpp
@@ -303,10 +303,18 @@ Status SessionConnection::on_packet(const MsgInfo &info, const mtproto_api::rpc_
 }
 
 Status SessionConnection::on_packet(const MsgInfo &info, const mtproto_api::new_session_created &new_session_created) {
-  VLOG(mtproto) << "NEW_SESSION_CREATED: [first_msg_id:" << format::as_hex(new_session_created.first_msg_id_)
+  auto first_message_id = new_session_created.first_msg_id_;
+  VLOG(mtproto) << "NEW_SESSION_CREATED: [first_msg_id:" << format::as_hex(first_message_id)
                 << "] [unique_id:" << format::as_hex(new_session_created.unique_id_)
                 << "] [server_salt:" << format::as_hex(new_session_created.server_salt_) << "]";
-  callback_->on_new_session_created(new_session_created.unique_id_, new_session_created.first_msg_id_);
+
+  auto it = service_queries_.find(first_message_id);
+  if (it != service_queries_.end()) {
+    first_message_id = it->second.container_message_id;
+    LOG(INFO) << "Update first_message_id to container's " << format::as_hex(first_message_id);
+  }
+
+  callback_->on_new_session_created(new_session_created.unique_id_, first_message_id);
   return Status::OK();
 }
 
@@ -609,18 +617,16 @@ void SessionConnection::on_message_failed_inner(uint64 id) {
   service_queries_.erase(it);
 
   switch (query.type) {
-    case ServiceQuery::ResendAnswer: {
+    case ServiceQuery::ResendAnswer:
       for (auto message_id : query.message_ids) {
         resend_answer(message_id);
       }
       break;
-    }
-    case ServiceQuery::GetStateInfo: {
+    case ServiceQuery::GetStateInfo:
       for (auto message_id : query.message_ids) {
         get_state_info(message_id);
       }
       break;
-    }
     default:
       UNREACHABLE();
   }
@@ -1005,11 +1011,11 @@ void SessionConnection::flush_packet() {
 
   if (resend_answer_message_id) {
     service_queries_.emplace(resend_answer_message_id,
-                             ServiceQuery{ServiceQuery::ResendAnswer, std::move(to_resend_answer)});
+                             ServiceQuery{ServiceQuery::ResendAnswer, container_id, std::move(to_resend_answer)});
   }
   if (get_state_info_message_id) {
     service_queries_.emplace(get_state_info_message_id,
-                             ServiceQuery{ServiceQuery::GetStateInfo, std::move(to_get_state_info)});
+                             ServiceQuery{ServiceQuery::GetStateInfo, container_id, std::move(to_get_state_info)});
   }
   if (ping_id != 0) {
     last_ping_container_id_ = container_id;
diff --git a/td/mtproto/SessionConnection.h b/td/mtproto/SessionConnection.h
index 2459ee94a..f14971463 100644
--- a/td/mtproto/SessionConnection.h
+++ b/td/mtproto/SessionConnection.h
@@ -170,6 +170,7 @@ class SessionConnection final
 
   struct ServiceQuery {
     enum Type { GetStateInfo, ResendAnswer } type;
+    uint64 container_message_id;
     vector<int64> message_ids;
   };
   vector<int64> to_resend_answer_;
diff --git a/td/telegram/net/Session.cpp b/td/telegram/net/Session.cpp
index c9acec0bf..a07b51725 100644
--- a/td/telegram/net/Session.cpp
+++ b/td/telegram/net/Session.cpp
@@ -698,7 +698,7 @@ void Session::on_closed(Status status) {
 }
 
 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;
+  LOG(INFO) << "New session " << unique_id << " created with first message_id " << format::as_hex(first_message_id);
   if (!use_pfs_ && !auth_data_.use_pfs()) {
     last_success_timestamp_ = Time::now();
   }
@@ -712,9 +712,9 @@ void Session::on_new_session_created(uint64 unique_id, uint64 first_message_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;
+    LOG(INFO) << "Update first_message_id to container's " << format::as_hex(first_message_id);
   } else {
-    LOG(ERROR) << "Failed to find query " << first_message_id << " from the new session";
+    LOG(INFO) << "Failed to find query " << format::as_hex(first_message_id) << " from the new session";
   }
   for (auto it = sent_queries_.begin(); it != sent_queries_.end();) {
     Query *query_ptr = &it->second;