diff --git a/td/mtproto/Handshake.cpp b/td/mtproto/Handshake.cpp index b917d24bd..d8f47bba6 100644 --- a/td/mtproto/Handshake.cpp +++ b/td/mtproto/Handshake.cpp @@ -50,7 +50,8 @@ Result AuthKeyHandshake::fill_data_with_hash(uint8 *data_with_hash, cons return Status::Error("Too big data"); } as(data_ptr) = data.get_id(); - tl_store_unsafe(data, data_ptr + 4); + auto real_size = tl_store_unsafe(data, data_ptr + 4); + CHECK(real_size == data_size); sha1(Slice(data_ptr, data_size + 4), data_with_hash); return data_size + 20 + 4; } @@ -193,7 +194,8 @@ Status AuthKeyHandshake::on_server_dh_params(Slice message, Callback *connection string encrypted_data_str(encrypted_data_size_with_pad, 0); MutableSlice encrypted_data = encrypted_data_str; as(encrypted_data.begin() + 20) = data.get_id(); - tl_store_unsafe(data, encrypted_data.ubegin() + 20 + 4); + auto real_size = tl_store_unsafe(data, encrypted_data.ubegin() + 20 + 4); + CHECK(real_size + 4 == data_size); sha1(Slice(encrypted_data.ubegin() + 20, data_size), encrypted_data.ubegin()); Random::secure_bytes(encrypted_data.ubegin() + encrypted_data_size, encrypted_data_size_with_pad - encrypted_data_size); @@ -230,8 +232,10 @@ Status AuthKeyHandshake::on_dh_gen_response(Slice message, Callback *connection) return Status::OK(); } void AuthKeyHandshake::send(Callback *connection, const Storer &storer) { - auto writer = BufferWriter{storer.size(), 0, 0}; - storer.store(writer.as_slice().ubegin()); + auto size = storer.size(); + auto writer = BufferWriter{size, 0, 0}; + auto real_size = storer.store(writer.as_slice().ubegin()); + CHECK(real_size == size); last_query_ = writer.as_buffer_slice(); return do_send(connection, create_storer(last_query_.as_slice())); } diff --git a/td/mtproto/SessionConnection.cpp b/td/mtproto/SessionConnection.cpp index 0f7c466cc..96284fbfe 100644 --- a/td/mtproto/SessionConnection.cpp +++ b/td/mtproto/SessionConnection.cpp @@ -743,8 +743,10 @@ std::pair SessionConnection::encrypted_bind(int64 perm_key, mtproto_api::bind_auth_key_inner object(nonce, temp_key, perm_key, auth_data_->session_id_, expire_at); auto object_storer = create_storer(object); - auto object_packet = BufferWriter{object_storer.size(), 0, 0}; - object_storer.store(object_packet.as_slice().ubegin()); + auto size = object_storer.size(); + auto object_packet = BufferWriter{size, 0, 0}; + auto real_size = object_storer.store(object_packet.as_slice().ubegin()); + CHECK(size == real_size); Query query{auth_data_->next_message_id(Time::now_cached()), 0, object_packet.as_buffer_slice(), false, 0, false}; PacketStorer query_storer(query, Slice()); diff --git a/td/mtproto/Transport.cpp b/td/mtproto/Transport.cpp index 62c5cbfd7..4b7723085 100644 --- a/td/mtproto/Transport.cpp +++ b/td/mtproto/Transport.cpp @@ -194,14 +194,16 @@ size_t Transport::write_no_crypto(const Storer &storer, PacketInfo *info, Mutabl } auto &header = as(dest.begin()); header.auth_key_id = 0; - storer.store(header.data); + auto real_size = storer.store(header.data); + CHECK(real_size == storer.size()); return size; } template void Transport::write_crypto_impl(int X, const Storer &storer, const AuthKey &auth_key, PacketInfo *info, HeaderT *header, size_t data_size) { - storer.store(header->data); + auto real_data_size = storer.store(header->data); + CHECK(real_data_size == data_size); VLOG(raw_mtproto) << "SEND" << format::as_hex_dump<4>(Slice(header->data, data_size)); // LOG(ERROR) << "SEND" << format::as_hex_dump<4>(Slice(header->data, data_size)) << info->version; diff --git a/td/telegram/files/FileLocation.h b/td/telegram/files/FileLocation.h index 158933342..66118c25f 100644 --- a/td/telegram/files/FileLocation.h +++ b/td/telegram/files/FileLocation.h @@ -1300,6 +1300,7 @@ string as_key(const T &object) { TlStorerUnsafe storer(key.ubegin()); storer.store_int(T::KEY_MAGIC); object.as_key().store(storer); + CHECK(storer.get_buf() == key.uend()); return key.str(); } diff --git a/td/telegram/net/NetQueryCreator.cpp b/td/telegram/net/NetQueryCreator.cpp index 541bbfaae..c201b961d 100644 --- a/td/telegram/net/NetQueryCreator.cpp +++ b/td/telegram/net/NetQueryCreator.cpp @@ -14,7 +14,8 @@ NetQueryCreator::Ptr NetQueryCreator::create(uint64 id, const Storer &storer, Dc NetQuery::AuthFlag auth_flag, NetQuery::GzipFlag gzip_flag, double total_timeout_limit) { BufferSlice slice(storer.size()); - storer.store(slice.as_slice().ubegin()); + auto real_size = storer.store(slice.as_slice().ubegin()); + CHECK(real_size == slice.size()); // TODO: magic constant if (slice.size() < (1 << 8)) { diff --git a/td/telegram/net/Session.cpp b/td/telegram/net/Session.cpp index 8549ae173..a88b02636 100644 --- a/td/telegram/net/Session.cpp +++ b/td/telegram/net/Session.cpp @@ -461,7 +461,8 @@ void Session::on_session_created(uint64 unique_id, uint64 first_id) { telegram_api::updatesTooLong too_long_; auto storer = create_storer(too_long_); BufferSlice packet(storer.size()); - storer.store(packet.as_slice().ubegin()); + auto real_size = storer.store(packet.as_slice().ubegin()); + CHECK(real_size == packet.size()); return_query(G()->net_query_creator().create_result(0, std::move(packet))); } diff --git a/tdutils/td/utils/StorerBase.h b/tdutils/td/utils/StorerBase.h index e6fea28e1..c9db6879d 100644 --- a/tdutils/td/utils/StorerBase.h +++ b/tdutils/td/utils/StorerBase.h @@ -19,7 +19,7 @@ class Storer { Storer &operator=(Storer &&) = default; virtual ~Storer() = default; virtual size_t size() const = 0; - virtual size_t store(uint8 *ptr) const = 0; + virtual size_t store(uint8 *ptr) const TD_WARN_UNUSED_RESULT = 0; }; } // namespace td diff --git a/tdutils/td/utils/tl_helpers.h b/tdutils/td/utils/tl_helpers.h index fe6343458..7cb0b6711 100644 --- a/tdutils/td/utils/tl_helpers.h +++ b/tdutils/td/utils/tl_helpers.h @@ -189,11 +189,13 @@ string serialize(const T &object) { MutableSlice data = ptr.as_slice(); TlStorerUnsafe storer(data.ubegin()); store(object, storer); + CHECK(storer.get_buf() == data.uend()); key.assign(data.begin(), data.size()); } else { MutableSlice data = key; TlStorerUnsafe storer(data.ubegin()); store(object, storer); + CHECK(storer.get_buf() == data.uend()); } return key; } diff --git a/tdutils/td/utils/tl_storers.h b/tdutils/td/utils/tl_storers.h index da2eb07cb..95e0fe2e1 100644 --- a/tdutils/td/utils/tl_storers.h +++ b/tdutils/td/utils/tl_storers.h @@ -272,6 +272,9 @@ size_t tl_calc_length(const T &data) { return storer_calc_length.get_length(); } +template +size_t tl_store_unsafe(const T &data, unsigned char *dst) TD_WARN_UNUSED_RESULT; + template size_t tl_store_unsafe(const T &data, unsigned char *dst) { TlStorerUnsafe storer_unsafe(dst); diff --git a/test/secret.cpp b/test/secret.cpp index 618e37607..85a3d3fba 100644 --- a/test/secret.cpp +++ b/test/secret.cpp @@ -867,7 +867,8 @@ class Master : public Actor { config.version_ = 12; auto storer = TLObjectStorer(config); BufferSlice answer(storer.size()); - storer.store(answer.as_slice().ubegin()); + auto real_size = storer.store(answer.as_slice().ubegin()); + CHECK(real_size == answer.size()); net_query->set_ok(std::move(answer)); send_closure(std::move(callback), &NetQueryCallback::on_result, std::move(net_query)); } @@ -891,7 +892,8 @@ class Master : public Actor { my_api::encryptedChat encrypted_chat(123, 321, 0, 1, 2, BufferSlice(), request_encryption.key_fingerprint_); auto storer = TLObjectStorer(encrypted_chat); BufferSlice answer(storer.size()); - storer.store(answer.as_slice().ubegin()); + auto real_size = storer.store(answer.as_slice().ubegin()); + CHECK(real_size == answer.size()); net_query->set_ok(std::move(answer)); send_closure(std::move(callback), &NetQueryCallback::on_result, std::move(net_query)); send_closure(alice_, &SecretChatProxy::start_test); @@ -935,7 +937,8 @@ class Master : public Actor { sent_message.date_ = 0; auto storer = TLObjectStorer(sent_message); BufferSlice answer(storer.size()); - storer.store(answer.as_slice().ubegin()); + auto real_size = storer.store(answer.as_slice().ubegin()); + CHECK(real_size == answer.size()); net_query->set_ok(std::move(answer)); send_closure(std::move(callback), &NetQueryCallback::on_result, std::move(net_query));