1 files changed, 87 insertions, 104 deletions

diff --git a/protocols/Telegram/tdlib/td/td/mtproto/Transport.cpp b/protocols/Telegram/tdlib/td/td/mtproto/Transport.cpp
index ef0440ef6f..1a98182dd8 100644
--- a/protocols/Telegram/tdlib/td/td/mtproto/Transport.cpp
+++ b/protocols/Telegram/tdlib/td/td/mtproto/Transport.cpp
@@ -1,5 +1,5 @@
 //
-// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2023
+// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2024
 //
 // Distributed under the Boost Software License, Version 1.0. (See accompanying
 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
@@ -8,12 +8,12 @@
 
 #include "td/mtproto/AuthKey.h"
 #include "td/mtproto/KDF.h"
+#include "td/mtproto/MessageId.h"
 
 #include "td/utils/as.h"
 #include "td/utils/crypto.h"
 #include "td/utils/format.h"
 #include "td/utils/logging.h"
-#include "td/utils/misc.h"
 #include "td/utils/Random.h"
 #include "td/utils/SliceBuilder.h"
 #include "td/utils/Status.h"
@@ -43,7 +43,7 @@ struct CryptoHeader {
 
   // It is weird to generate message_id and seq_no while writing a packet.
   //
-  // uint64 message_id;
+  // uint64 msg_id;
   // uint32 seq_no;
   // uint32 message_data_length;
   uint8 data[0];  // use compiler extension
@@ -69,7 +69,7 @@ struct CryptoHeader {
 };
 
 struct CryptoPrefix {
-  uint64 message_id;
+  uint64 msg_id;
   uint32 seq_no;
   uint32 message_data_length;
 };
@@ -109,9 +109,9 @@ struct EndToEndPrefix {
 struct NoCryptoHeader {
   uint64 auth_key_id;
 
-  // message_id is removed from CryptoHeader. Should be removed from here too.
+  // msg_id is removed from CryptoHeader. Should be removed from here too.
   //
-  // uint64 message_id;
+  // uint64 msg_id;
   // uint32 message_data_length;
   uint8 data[0];  // use compiler extension
 
@@ -186,26 +186,21 @@ size_t do_calc_crypto_size2_rand(size_t data_size, size_t enc_size, size_t raw_s
 }  // namespace
 
 template <class HeaderT>
-size_t Transport::calc_crypto_size2(size_t data_size, PacketInfo *info) {
-  if (info->size != 0) {
-    return info->size;
-  }
-
+size_t Transport::calc_crypto_size2(size_t data_size, PacketInfo *packet_info) {
   size_t enc_size = HeaderT::encrypted_header_size();
   size_t raw_size = sizeof(HeaderT) - enc_size;
-  if (info->use_random_padding) {
-    info->size = narrow_cast<uint32>(do_calc_crypto_size2_rand(data_size, enc_size, raw_size));
+  if (packet_info->use_random_padding) {
+    return do_calc_crypto_size2_rand(data_size, enc_size, raw_size);
   } else {
-    info->size = narrow_cast<uint32>(do_calc_crypto_size2_basic(data_size, enc_size, raw_size));
+    return do_calc_crypto_size2_basic(data_size, enc_size, raw_size);
   }
-  return info->size;
 }
 
 size_t Transport::calc_no_crypto_size(size_t data_size) {
   return sizeof(NoCryptoHeader) + data_size;
 }
 
-Status Transport::read_no_crypto(MutableSlice message, PacketInfo *info, MutableSlice *data) {
+Status Transport::read_no_crypto(MutableSlice message, PacketInfo *packet_info, MutableSlice *data) {
   if (message.size() < sizeof(NoCryptoHeader)) {
     return Status::Error(PSLICE() << "Invalid MTProto message: too small [message.size() = " << message.size()
                                   << "] < [sizeof(NoCryptoHeader) = " << sizeof(NoCryptoHeader) << "]");
@@ -218,7 +213,7 @@ Status Transport::read_no_crypto(MutableSlice message, PacketInfo *info, Mutable
 
 template <class HeaderT, class PrefixT>
 Status Transport::read_crypto_impl(int X, MutableSlice message, const AuthKey &auth_key, HeaderT **header_ptr,
-                                   PrefixT **prefix_ptr, MutableSlice *data, PacketInfo *info) {
+                                   PrefixT **prefix_ptr, MutableSlice *data, PacketInfo *packet_info) {
   if (message.size() < sizeof(HeaderT)) {
     return Status::Error(PSLICE() << "Invalid MTProto message: too small [message.size() = " << message.size()
                                   << "] < [sizeof(HeaderT) = " << sizeof(HeaderT) << "]");
@@ -237,7 +232,7 @@ Status Transport::read_crypto_impl(int X, MutableSlice message, const AuthKey &a
 
   UInt256 aes_key;
   UInt256 aes_iv;
-  if (info->version == 1) {
+  if (packet_info->version == 1) {
     KDF(auth_key.key(), header->message_key, X, &aes_key, &aes_iv);
   } else {
     KDF2(auth_key.key(), header->message_key, X, &aes_key, &aes_iv);
@@ -256,14 +251,14 @@ Status Transport::read_crypto_impl(int X, MutableSlice message, const AuthKey &a
   bool is_length_bad = false;
   UInt128 real_message_key;
 
-  if (info->version == 1) {
-    is_length_bad |= info->check_mod4 && prefix->message_data_length % 4 != 0;
+  if (packet_info->version == 1) {
+    is_length_bad |= packet_info->check_mod4 && prefix->message_data_length % 4 != 0;
     auto expected_size = calc_crypto_size<HeaderT>(data_size);
     is_length_bad |= expected_size != message.size();
     auto check_size = data_size * (1 - is_length_bad) + tail_size * is_length_bad;
-    std::tie(info->message_ack, real_message_key) = calc_message_ack_and_key(*header, check_size);
+    std::tie(packet_info->message_ack, real_message_key) = calc_message_ack_and_key(*header, check_size);
   } else {
-    std::tie(info->message_ack, real_message_key) = calc_message_key2(auth_key, X, to_decrypt);
+    std::tie(packet_info->message_ack, real_message_key) = calc_message_key2(auth_key, X, to_decrypt);
   }
 
   int is_key_bad = false;
@@ -276,8 +271,8 @@ Status Transport::read_crypto_impl(int X, MutableSlice message, const AuthKey &a
                                   << "] [expected = " << format::as_hex_dump(real_message_key) << "]");
   }
 
-  if (info->version == 2) {
-    if (info->check_mod4 && prefix->message_data_length % 4 != 0) {
+  if (packet_info->version == 2) {
+    if (packet_info->check_mod4 && prefix->message_data_length % 4 != 0) {
       return Status::Error(PSLICE() << "Invalid MTProto message: invalid length (not divisible by four)"
                                     << tag("total_size", message.size())
                                     << tag("message_data_length", prefix->message_data_length));
@@ -304,125 +299,113 @@ Status Transport::read_crypto_impl(int X, MutableSlice message, const AuthKey &a
   return Status::OK();
 }
 
-Status Transport::read_crypto(MutableSlice message, const AuthKey &auth_key, PacketInfo *info, MutableSlice *data) {
+Status Transport::read_crypto(MutableSlice message, const AuthKey &auth_key, PacketInfo *packet_info,
+                              MutableSlice *data) {
   CryptoHeader *header = nullptr;
   CryptoPrefix *prefix = nullptr;
-  TRY_STATUS(read_crypto_impl(8, message, auth_key, &header, &prefix, data, info));
+  TRY_STATUS(read_crypto_impl(8, message, auth_key, &header, &prefix, data, packet_info));
   CHECK(header != nullptr);
   CHECK(prefix != nullptr);
-  CHECK(info != nullptr);
-  info->type = PacketInfo::Common;
-  info->salt = header->salt;
-  info->session_id = header->session_id;
-  info->message_id = prefix->message_id;
-  info->seq_no = prefix->seq_no;
+  CHECK(packet_info != nullptr);
+  packet_info->type = PacketInfo::Common;
+  packet_info->salt = header->salt;
+  packet_info->session_id = header->session_id;
+  packet_info->message_id = MessageId(prefix->msg_id);
+  packet_info->seq_no = prefix->seq_no;
   return Status::OK();
 }
-Status Transport::read_e2e_crypto(MutableSlice message, const AuthKey &auth_key, PacketInfo *info, MutableSlice *data) {
+Status Transport::read_e2e_crypto(MutableSlice message, const AuthKey &auth_key, PacketInfo *packet_info,
+                                  MutableSlice *data) {
   EndToEndHeader *header = nullptr;
   EndToEndPrefix *prefix = nullptr;
-  TRY_STATUS(read_crypto_impl(info->is_creator && info->version != 1 ? 8 : 0, message, auth_key, &header, &prefix, data,
-                              info));
+  TRY_STATUS(read_crypto_impl(packet_info->is_creator && packet_info->version != 1 ? 8 : 0, message, auth_key, &header,
+                              &prefix, data, packet_info));
   CHECK(header != nullptr);
   CHECK(prefix != nullptr);
-  CHECK(info != nullptr);
-  info->type = PacketInfo::EndToEnd;
+  CHECK(packet_info != nullptr);
+  packet_info->type = PacketInfo::EndToEnd;
   return Status::OK();
 }
 
-size_t Transport::write_no_crypto(const Storer &storer, PacketInfo *info, MutableSlice dest) {
+BufferWriter Transport::write_no_crypto(const Storer &storer, PacketInfo *packet_info, size_t prepend_size,
+                                        size_t append_size) {
   size_t size = calc_no_crypto_size(storer.size());
-  if (size > dest.size()) {
-    return size;
-  }
+  auto packet = BufferWriter{size, prepend_size, append_size};
+
   // NoCryptoHeader
-  as<uint64>(dest.begin()) = 0;
-  auto real_size = storer.store(dest.ubegin() + sizeof(uint64));
+  auto *begin = packet.as_mutable_slice().ubegin();
+  as<uint64>(begin) = 0;
+  auto real_size = storer.store(begin + sizeof(uint64));
   CHECK(real_size == storer.size());
-  return size;
+  return packet;
 }
 
 template <class HeaderT>
-void Transport::write_crypto_impl(int X, const Storer &storer, const AuthKey &auth_key, PacketInfo *info,
-                                  HeaderT *header, size_t data_size) {
+void Transport::write_crypto_impl(int X, const Storer &storer, const AuthKey &auth_key, PacketInfo *packet_info,
+                                  HeaderT *header, size_t data_size, size_t padded_size) {
   auto real_data_size = storer.store(header->data);
   CHECK(real_data_size == data_size);
-  VLOG(raw_mtproto) << "Send packet of size " << data_size << " to session " << format::as_hex(info->session_id) << ":"
+  VLOG(raw_mtproto) << "Send packet of size " << data_size << ':'
                     << format::as_hex_dump<4>(Slice(header->data, data_size));
-
-  size_t size = 0;
-  if (info->version == 1) {
-    size = calc_crypto_size<HeaderT>(data_size);
-  } else {
-    size = calc_crypto_size2<HeaderT>(data_size, info);
-  }
-
-  size_t pad_size = size - (sizeof(HeaderT) + data_size);
+  size_t pad_size = padded_size - (sizeof(HeaderT) + data_size);
   MutableSlice pad(header->data + data_size, pad_size);
   Random::secure_bytes(pad.ubegin(), pad.size());
   MutableSlice to_encrypt = MutableSlice(header->encrypt_begin(), pad.uend());
 
-  if (info->version == 1) {
-    std::tie(info->message_ack, info->message_key) = calc_message_ack_and_key(*header, data_size);
-  } else {
-    std::tie(info->message_ack, info->message_key) = calc_message_key2(auth_key, X, to_encrypt);
-  }
-
-  header->message_key = info->message_key;
-
   UInt256 aes_key;
   UInt256 aes_iv;
-  if (info->version == 1) {
+  if (packet_info->version == 1) {
+    std::tie(packet_info->message_ack, header->message_key) = calc_message_ack_and_key(*header, data_size);
     KDF(auth_key.key(), header->message_key, X, &aes_key, &aes_iv);
   } else {
+    std::tie(packet_info->message_ack, header->message_key) = calc_message_key2(auth_key, X, to_encrypt);
     KDF2(auth_key.key(), header->message_key, X, &aes_key, &aes_iv);
   }
 
   aes_ige_encrypt(as_slice(aes_key), as_mutable_slice(aes_iv), to_encrypt, to_encrypt);
 }
 
-size_t Transport::write_crypto(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, MutableSlice dest) {
+BufferWriter Transport::write_crypto(const Storer &storer, const AuthKey &auth_key, PacketInfo *packet_info,
+                                     size_t prepend_size, size_t append_size) {
   size_t data_size = storer.size();
-  size_t size;
-  if (info->version == 1) {
-    size = calc_crypto_size<CryptoHeader>(data_size);
+  size_t padded_size;
+  if (packet_info->version == 1) {
+    padded_size = calc_crypto_size<CryptoHeader>(data_size);
   } else {
-    size = calc_crypto_size2<CryptoHeader>(data_size, info);
-  }
-  if (size > dest.size()) {
-    return size;
+    padded_size = calc_crypto_size2<CryptoHeader>(data_size, packet_info);
   }
+  auto packet = BufferWriter{padded_size, prepend_size, append_size};
 
   //FIXME: rewrite without reinterpret cast
-  auto &header = *reinterpret_cast<CryptoHeader *>(dest.begin());
+  auto &header = *reinterpret_cast<CryptoHeader *>(packet.as_mutable_slice().begin());
   header.auth_key_id = auth_key.id();
-  header.salt = info->salt;
-  header.session_id = info->session_id;
+  header.salt = packet_info->salt;
+  header.session_id = packet_info->session_id;
 
-  write_crypto_impl(0, storer, auth_key, info, &header, data_size);
+  write_crypto_impl(0, storer, auth_key, packet_info, &header, data_size, padded_size);
 
-  return size;
+  return packet;
 }
 
-size_t Transport::write_e2e_crypto(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, MutableSlice dest) {
+BufferWriter Transport::write_e2e_crypto(const Storer &storer, const AuthKey &auth_key, PacketInfo *packet_info,
+                                         size_t prepend_size, size_t append_size) {
   size_t data_size = storer.size();
-  size_t size;
-  if (info->version == 1) {
-    size = calc_crypto_size<EndToEndHeader>(data_size);
+  size_t padded_size;
+  if (packet_info->version == 1) {
+    padded_size = calc_crypto_size<EndToEndHeader>(data_size);
   } else {
-    size = calc_crypto_size2<EndToEndHeader>(data_size, info);
-  }
-  if (size > dest.size()) {
-    return size;
+    padded_size = calc_crypto_size2<EndToEndHeader>(data_size, packet_info);
   }
+  auto packet = BufferWriter{padded_size, prepend_size, append_size};
 
   //FIXME: rewrite without reinterpret cast
-  auto &header = *reinterpret_cast<EndToEndHeader *>(dest.begin());
+  auto &header = *reinterpret_cast<EndToEndHeader *>(packet.as_mutable_slice().begin());
   header.auth_key_id = auth_key.id();
 
-  write_crypto_impl(info->is_creator || info->version == 1 ? 0 : 8, storer, auth_key, info, &header, data_size);
+  write_crypto_impl(packet_info->is_creator || packet_info->version == 1 ? 0 : 8, storer, auth_key, packet_info,
+                    &header, data_size, padded_size);
 
-  return size;
+  return packet;
 }
 
 Result<uint64> Transport::read_auth_key_id(Slice message) {
@@ -432,8 +415,8 @@ Result<uint64> Transport::read_auth_key_id(Slice message) {
   return as<uint64>(message.begin());
 }
 
-Result<Transport::ReadResult> Transport::read(MutableSlice message, const AuthKey &auth_key, PacketInfo *info) {
-  if (message.size() < 12) {
+Result<Transport::ReadResult> Transport::read(MutableSlice message, const AuthKey &auth_key, PacketInfo *packet_info) {
+  if (message.size() < 16) {
     if (message.size() < 4) {
       return Status::Error(PSLICE() << "Invalid MTProto message: smaller than 4 bytes [size = " << message.size()
                                     << "]");
@@ -449,31 +432,31 @@ Result<Transport::ReadResult> Transport::read(MutableSlice message, const AuthKe
     }
   }
 
-  info->auth_key_id = as<int64>(message.begin());
-  info->no_crypto_flag = info->auth_key_id == 0;
+  packet_info->no_crypto_flag = as<int64>(message.begin()) == 0;
   MutableSlice data;
-  if (info->type == PacketInfo::EndToEnd) {
-    TRY_STATUS(read_e2e_crypto(message, auth_key, info, &data));
-  } else if (info->no_crypto_flag) {
-    TRY_STATUS(read_no_crypto(message, info, &data));
+  if (packet_info->type == PacketInfo::EndToEnd) {
+    TRY_STATUS(read_e2e_crypto(message, auth_key, packet_info, &data));
+  } else if (packet_info->no_crypto_flag) {
+    TRY_STATUS(read_no_crypto(message, packet_info, &data));
   } else {
     if (auth_key.empty()) {
       return Status::Error("Failed to decrypt MTProto message: auth key is empty");
     }
-    TRY_STATUS(read_crypto(message, auth_key, info, &data));
+    TRY_STATUS(read_crypto(message, auth_key, packet_info, &data));
   }
   return ReadResult::make_packet(data);
 }
 
-size_t Transport::write(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, MutableSlice dest) {
-  if (info->type == PacketInfo::EndToEnd) {
-    return write_e2e_crypto(storer, auth_key, info, dest);
+BufferWriter Transport::write(const Storer &storer, const AuthKey &auth_key, PacketInfo *packet_info,
+                              size_t prepend_size, size_t append_size) {
+  if (packet_info->type == PacketInfo::EndToEnd) {
+    return write_e2e_crypto(storer, auth_key, packet_info, prepend_size, append_size);
   }
-  if (info->no_crypto_flag) {
-    return write_no_crypto(storer, info, dest);
+  if (packet_info->no_crypto_flag) {
+    return write_no_crypto(storer, packet_info, prepend_size, append_size);
   } else {
     CHECK(!auth_key.empty());
-    return write_crypto(storer, auth_key, info, dest);
+    return write_crypto(storer, auth_key, packet_info, prepend_size, append_size);
   }
 }