From 5ce5ef62349c61630bca23ec1896a3f4872f888f Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Fri, 1 Sep 2023 16:11:21 +0300
Subject: [PATCH] Save all state flags to ack_state_.

---
 td/mtproto/PingConnection.cpp    | 2 +-
 td/mtproto/SessionConnection.cpp | 6 +++---
 td/mtproto/SessionConnection.h   | 2 +-
 td/telegram/net/NetQueryStats.h  | 2 +-
 td/telegram/net/Session.cpp      | 6 ++++--
 td/telegram/net/Session.h        | 2 +-
 6 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/td/mtproto/PingConnection.cpp b/td/mtproto/PingConnection.cpp
index a08a1d61c..d47b7dffc 100644
--- a/td/mtproto/PingConnection.cpp
+++ b/td/mtproto/PingConnection.cpp
@@ -153,7 +153,7 @@ class PingConnectionPingPong final
   }
   void on_message_failed(uint64 id, Status status) final {
   }
-  void on_message_info(uint64 id, int32 state, uint64 answer_id, int32 answer_size) final {
+  void on_message_info(uint64 id, int32 state, uint64 answer_id, int32 answer_size, int32 source) final {
   }
 
   Status on_destroy_auth_key() final {
diff --git a/td/mtproto/SessionConnection.cpp b/td/mtproto/SessionConnection.cpp
index 2c7764df5..54e4f2dea 100644
--- a/td/mtproto/SessionConnection.cpp
+++ b/td/mtproto/SessionConnection.cpp
@@ -446,7 +446,7 @@ Status SessionConnection::on_msgs_state_info(const vector<int64> &message_ids, S
   }
   size_t i = 0;
   for (auto message_id : message_ids) {
-    callback_->on_message_info(static_cast<uint64>(message_id), info[i], 0, 0);
+    callback_->on_message_info(static_cast<uint64>(message_id), info[i], 0, 0, 1);
     i++;
   }
   return Status::OK();
@@ -472,13 +472,13 @@ Status SessionConnection::on_packet(const MsgInfo &info, const mtproto_api::msgs
 
 Status SessionConnection::on_packet(const MsgInfo &info, const mtproto_api::msg_detailed_info &msg_detailed_info) {
   callback_->on_message_info(msg_detailed_info.msg_id_, msg_detailed_info.status_, msg_detailed_info.answer_msg_id_,
-                             msg_detailed_info.bytes_);
+                             msg_detailed_info.bytes_, 2);
   return Status::OK();
 }
 
 Status SessionConnection::on_packet(const MsgInfo &info,
                                     const mtproto_api::msg_new_detailed_info &msg_new_detailed_info) {
-  callback_->on_message_info(0, 0, msg_new_detailed_info.answer_msg_id_, msg_new_detailed_info.bytes_);
+  callback_->on_message_info(0, 0, msg_new_detailed_info.answer_msg_id_, msg_new_detailed_info.bytes_, 0);
   return Status::OK();
 }
 
diff --git a/td/mtproto/SessionConnection.h b/td/mtproto/SessionConnection.h
index 398dada72..67a4821e9 100644
--- a/td/mtproto/SessionConnection.h
+++ b/td/mtproto/SessionConnection.h
@@ -119,7 +119,7 @@ class SessionConnection final
     virtual Status on_message_result_ok(uint64 id, BufferSlice packet, size_t original_size) = 0;
     virtual void on_message_result_error(uint64 id, int code, string message) = 0;
     virtual void on_message_failed(uint64 id, Status status) = 0;
-    virtual void on_message_info(uint64 id, int32 state, uint64 answer_id, int32 answer_size) = 0;
+    virtual void on_message_info(uint64 id, int32 state, uint64 answer_id, int32 answer_size, int32 source) = 0;
 
     virtual Status on_destroy_auth_key() = 0;
   };
diff --git a/td/telegram/net/NetQueryStats.h b/td/telegram/net/NetQueryStats.h
index 84707e31b..9a28a4666 100644
--- a/td/telegram/net/NetQueryStats.h
+++ b/td/telegram/net/NetQueryStats.h
@@ -23,7 +23,7 @@ struct NetQueryDebug {
   double state_timestamp_ = 0;
   int32 state_change_count_ = 0;
   int32 send_failed_count_ = 0;
-  int ack_state_ = 0;
+  int32 ack_state_ = 0;
   bool unknown_state_ = false;
 };
 
diff --git a/td/telegram/net/Session.cpp b/td/telegram/net/Session.cpp
index 04910c30c..fe338d5bc 100644
--- a/td/telegram/net/Session.cpp
+++ b/td/telegram/net/Session.cpp
@@ -1046,7 +1046,8 @@ void Session::on_message_failed(uint64 message_id, Status status) {
   on_message_failed_inner(message_id, false);
 }
 
-void Session::on_message_info(uint64 message_id, int32 state, uint64 answer_message_id, int32 answer_size) {
+void Session::on_message_info(uint64 message_id, int32 state, uint64 answer_message_id, int32 answer_size,
+                              int32 source) {
   auto it = sent_queries_.find(message_id);
   if (it != sent_queries_.end()) {
     if (it->second.net_query_->update_is_ready()) {
@@ -1079,7 +1080,8 @@ void Session::on_message_info(uint64 message_id, int32 state, uint64 answer_mess
         }
       // fallthrough
       case 4:
-        on_message_ack_impl(message_id, answer_message_id == 0 ? 4 : 2);
+        CHECK(0 <= source && source <= 3);
+        on_message_ack_impl(message_id, (answer_message_id ? 2 : 0) | (((state | source) & ((1 << 28) - 1)) << 2));
         break;
       default:
         LOG(ERROR) << "Invalid message info " << tag("state", state);
diff --git a/td/telegram/net/Session.h b/td/telegram/net/Session.h
index 64159aa33..0b094c36a 100644
--- a/td/telegram/net/Session.h
+++ b/td/telegram/net/Session.h
@@ -226,7 +226,7 @@ class Session final
   void on_message_result_error(uint64 message_id, int error_code, string message) final;
   void on_message_failed(uint64 message_id, Status status) final;
 
-  void on_message_info(uint64 message_id, int32 state, uint64 answer_message_id, int32 answer_size) final;
+  void on_message_info(uint64 message_id, int32 state, uint64 answer_message_id, int32 answer_size, int32 source) final;
 
   Status on_destroy_auth_key() final;