Simplify mtproto::Transport::write usage.

This commit is contained in:
levlam 2023-09-09 22:49:45 +03:00
parent d47c862560
commit d87cc2b143
5 changed files with 25 additions and 23 deletions

View File

@ -67,10 +67,8 @@ class RawConnectionDefault final : public RawConnection {
info.salt = salt; info.salt = salt;
info.session_id = session_id; info.session_id = session_id;
info.use_random_padding = transport_->use_random_padding(); info.use_random_padding = transport_->use_random_padding();
auto packet =
auto packet = BufferWriter{Transport::write(storer, auth_key, &info), transport_->max_prepend_size(), Transport::write(storer, auth_key, &info, transport_->max_prepend_size(), transport_->max_append_size());
transport_->max_append_size()};
Transport::write(storer, auth_key, &info, packet.as_mutable_slice());
bool use_quick_ack = false; bool use_quick_ack = false;
if (quick_ack_token != 0 && transport_->support_quick_ack()) { if (quick_ack_token != 0 && transport_->support_quick_ack()) {
@ -90,11 +88,10 @@ class RawConnectionDefault final : public RawConnection {
uint64 send_no_crypto(const Storer &storer) final { uint64 send_no_crypto(const Storer &storer) final {
PacketInfo info; PacketInfo info;
info.no_crypto_flag = true; info.no_crypto_flag = true;
auto packet = BufferWriter{Transport::write(storer, AuthKey(), &info), transport_->max_prepend_size(), auto packet =
transport_->max_append_size()}; Transport::write(storer, AuthKey(), &info, transport_->max_prepend_size(), transport_->max_append_size());
Transport::write(storer, AuthKey(), &info, packet.as_mutable_slice());
LOG(INFO) << "Send handshake packet: " << format::as_hex_dump<4>(packet.as_slice()); LOG(INFO) << "Send handshake packet: " << format::as_hex_dump<4>(packet.as_slice());
transport_->write(std::move(packet), false); transport_->write(std::move(packet), false);
return info.message_id; return info.message_id;
@ -311,9 +308,7 @@ class RawConnectionHttp final : public RawConnection {
info.salt = salt; info.salt = salt;
info.session_id = session_id; info.session_id = session_id;
info.use_random_padding = false; info.use_random_padding = false;
auto packet = Transport::write(storer, auth_key, &info);
auto packet = BufferWriter{Transport::write(storer, auth_key, &info), 0, 0};
Transport::write(storer, auth_key, &info, packet.as_mutable_slice());
auto packet_size = packet.size(); auto packet_size = packet.size();
send_packet(packet.as_buffer_slice()); send_packet(packet.as_buffer_slice());
@ -322,10 +317,9 @@ class RawConnectionHttp final : public RawConnection {
uint64 send_no_crypto(const Storer &storer) final { uint64 send_no_crypto(const Storer &storer) final {
PacketInfo info; PacketInfo info;
info.no_crypto_flag = true; info.no_crypto_flag = true;
auto packet = BufferWriter{Transport::write(storer, AuthKey(), &info), 0, 0}; auto packet = Transport::write(storer, AuthKey(), &info);
Transport::write(storer, AuthKey(), &info, packet.as_mutable_slice());
LOG(INFO) << "Send handshake packet: " << format::as_hex_dump<4>(packet.as_slice()); LOG(INFO) << "Send handshake packet: " << format::as_hex_dump<4>(packet.as_slice());
send_packet(packet.as_buffer_slice()); send_packet(packet.as_buffer_slice());
return info.message_id; return info.message_id;

View File

@ -843,15 +843,13 @@ std::pair<uint64, BufferSlice> SessionConnection::encrypted_bind(int64 perm_key,
auth_data_->next_message_id(Time::now_cached()), 0, object_packet.as_buffer_slice(), false, {}, false}; auth_data_->next_message_id(Time::now_cached()), 0, object_packet.as_buffer_slice(), false, {}, false};
PacketStorer<QueryImpl> query_storer(query, Slice()); PacketStorer<QueryImpl> query_storer(query, Slice());
const AuthKey &main_auth_key = auth_data_->get_main_auth_key();
PacketInfo info; PacketInfo info;
info.version = 1; info.version = 1;
info.no_crypto_flag = false; info.no_crypto_flag = false;
info.salt = Random::secure_int64(); info.salt = Random::secure_int64();
info.session_id = Random::secure_int64(); info.session_id = Random::secure_int64();
auto packet = Transport::write(query_storer, main_auth_key, &info);
const AuthKey &main_auth_key = auth_data_->get_main_auth_key();
auto packet = BufferWriter{Transport::write(query_storer, main_auth_key, &info), 0, 0};
Transport::write(query_storer, main_auth_key, &info, packet.as_mutable_slice());
return std::make_pair(query.message_id, packet.as_buffer_slice()); return std::make_pair(query.message_id, packet.as_buffer_slice());
} }

View File

@ -451,7 +451,7 @@ Result<Transport::ReadResult> Transport::read(MutableSlice message, const AuthKe
return ReadResult::make_packet(data); return ReadResult::make_packet(data);
} }
size_t Transport::write(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, MutableSlice dest) { size_t Transport::do_write(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, MutableSlice dest) {
if (info->type == PacketInfo::EndToEnd) { if (info->type == PacketInfo::EndToEnd) {
return write_e2e_crypto(storer, auth_key, info, dest); return write_e2e_crypto(storer, auth_key, info, dest);
} }
@ -463,5 +463,12 @@ size_t Transport::write(const Storer &storer, const AuthKey &auth_key, PacketInf
} }
} }
BufferWriter Transport::write(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, size_t prepend_size,
size_t append_size) {
auto packet = BufferWriter{do_write(storer, auth_key, info, MutableSlice()), prepend_size, append_size};
do_write(storer, auth_key, info, packet.as_mutable_slice());
return packet;
}
} // namespace mtproto } // namespace mtproto
} // namespace td } // namespace td

View File

@ -8,6 +8,7 @@
#include "td/mtproto/PacketInfo.h" #include "td/mtproto/PacketInfo.h"
#include "td/utils/buffer.h"
#include "td/utils/common.h" #include "td/utils/common.h"
#include "td/utils/logging.h" #include "td/utils/logging.h"
#include "td/utils/Slice.h" #include "td/utils/Slice.h"
@ -88,9 +89,10 @@ class Transport {
// If auth_key is nonempty, encryption will be used. // If auth_key is nonempty, encryption will be used.
static Result<ReadResult> read(MutableSlice message, const AuthKey &auth_key, PacketInfo *info) TD_WARN_UNUSED_RESULT; static Result<ReadResult> read(MutableSlice message, const AuthKey &auth_key, PacketInfo *info) TD_WARN_UNUSED_RESULT;
static size_t write(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, static BufferWriter write(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, size_t prepend_size = 0,
MutableSlice dest = MutableSlice()); size_t append_size = 0);
// public for testing purposes
static std::pair<uint32, UInt128> calc_message_key2(const AuthKey &auth_key, int X, Slice to_encrypt); static std::pair<uint32, UInt128> calc_message_key2(const AuthKey &auth_key, int X, Slice to_encrypt);
private: private:
@ -126,6 +128,8 @@ class Transport {
template <class HeaderT> template <class HeaderT>
static void write_crypto_impl(int X, const Storer &storer, const AuthKey &auth_key, PacketInfo *info, HeaderT *header, static void write_crypto_impl(int X, const Storer &storer, const AuthKey &auth_key, PacketInfo *info, HeaderT *header,
size_t data_size, size_t padded_size); size_t data_size, size_t padded_size);
static size_t do_write(const Storer &storer, const AuthKey &auth_key, PacketInfo *info, MutableSlice dest);
}; };
} // namespace mtproto } // namespace mtproto

View File

@ -227,8 +227,7 @@ Result<BufferSlice> SecretChatActor::create_encrypted_message(int32 my_in_seq_no
info.type = mtproto::PacketInfo::EndToEnd; info.type = mtproto::PacketInfo::EndToEnd;
info.version = 2; info.version = 2;
info.is_creator = auth_state_.x == 0; info.is_creator = auth_state_.x == 0;
auto packet_writer = BufferWriter{mtproto::Transport::write(new_storer, *auth_key, &info), 0, 0}; auto packet_writer = mtproto::Transport::write(new_storer, *auth_key, &info);
mtproto::Transport::write(new_storer, *auth_key, &info, packet_writer.as_mutable_slice());
message = std::move(message_with_layer->message_); message = std::move(message_with_layer->message_);
return packet_writer.as_buffer_slice(); return packet_writer.as_buffer_slice();
} }