diff --git a/td/telegram/net/AuthDataShared.cpp b/td/telegram/net/AuthDataShared.cpp index 5417e499c..318d47206 100644 --- a/td/telegram/net/AuthDataShared.cpp +++ b/td/telegram/net/AuthDataShared.cpp @@ -19,8 +19,8 @@ namespace td { class AuthDataSharedImpl : public AuthDataShared { public: - AuthDataSharedImpl(DcId dc_id, std::shared_ptr public_rsa_key) - : dc_id_(dc_id), public_rsa_key_(std::move(public_rsa_key)) { + AuthDataSharedImpl(DcId dc_id, std::shared_ptr public_rsa_key, std::shared_ptr guard) + : dc_id_(dc_id), public_rsa_key_(std::move(public_rsa_key)), guard_(std::move(guard)) { log_auth_key(get_auth_key()); } @@ -89,6 +89,7 @@ class AuthDataSharedImpl : public AuthDataShared { DcId dc_id_; std::vector> auth_key_listeners_; std::shared_ptr public_rsa_key_; + std::shared_ptr guard_; RwMutex rw_mutex_; string auth_key_key() { @@ -111,7 +112,8 @@ class AuthDataSharedImpl : public AuthDataShared { } }; -std::shared_ptr AuthDataShared::create(DcId dc_id, std::shared_ptr public_rsa_key) { - return std::make_shared(dc_id, std::move(public_rsa_key)); +std::shared_ptr AuthDataShared::create(DcId dc_id, std::shared_ptr public_rsa_key, + std::shared_ptr guard) { + return std::make_shared(dc_id, std::move(public_rsa_key), std::move(guard)); } } // namespace td diff --git a/td/telegram/net/AuthDataShared.h b/td/telegram/net/AuthDataShared.h index ed55b8db0..027aa80ed 100644 --- a/td/telegram/net/AuthDataShared.h +++ b/td/telegram/net/AuthDataShared.h @@ -13,6 +13,7 @@ #include "td/telegram/net/PublicRsaKeyShared.h" #include "td/utils/common.h" +#include "td/utils/ScopeGuard.h" #include "td/utils/StringBuilder.h" #include @@ -70,7 +71,8 @@ class AuthDataShared { return state; } - static std::shared_ptr create(DcId dc_id, std::shared_ptr public_rsa_key); + static std::shared_ptr create(DcId dc_id, std::shared_ptr public_rsa_key, + std::shared_ptr guard); }; }; // namespace td diff --git a/td/telegram/net/NetQueryDispatcher.cpp b/td/telegram/net/NetQueryDispatcher.cpp index 1ea925d3c..cea313146 100644 --- a/td/telegram/net/NetQueryDispatcher.cpp +++ b/td/telegram/net/NetQueryDispatcher.cpp @@ -140,7 +140,7 @@ Status NetQueryDispatcher::wait_dc_init(DcId dc_id, bool force) { send_closure_later(public_rsa_key_watchdog_, &PublicRsaKeyWatchdog::add_public_rsa_key, public_rsa_key); is_cdn = true; } - auto auth_data = AuthDataShared::create(dc_id, std::move(public_rsa_key)); + auto auth_data = AuthDataShared::create(dc_id, std::move(public_rsa_key), td_guard_); int32 session_count = get_session_count(); bool use_pfs = get_use_pfs(); @@ -184,6 +184,7 @@ void NetQueryDispatcher::dispatch_with_callback(NetQueryPtr net_query, ActorShar void NetQueryDispatcher::stop() { std::lock_guard guard(main_dc_id_mutex_); + td_guard_.reset(); stop_flag_ = true; delayer_.hangup(); for (const auto &dc : dcs_) { @@ -246,6 +247,8 @@ NetQueryDispatcher::NetQueryDispatcher(std::function()> create_ref dc_auth_manager_ = create_actor("DcAuthManager", create_reference()); common_public_rsa_key_ = std::make_shared(DcId::empty()); public_rsa_key_watchdog_ = create_actor("PublicRsaKeyWatchdog", create_reference()); + + td_guard_ = create_shared_lamda_guard([actor = create_reference] {}); } NetQueryDispatcher::NetQueryDispatcher() = default; diff --git a/td/telegram/net/NetQueryDispatcher.h b/td/telegram/net/NetQueryDispatcher.h index 5020e8cac..5ccc5201c 100644 --- a/td/telegram/net/NetQueryDispatcher.h +++ b/td/telegram/net/NetQueryDispatcher.h @@ -74,6 +74,7 @@ class NetQueryDispatcher { std::shared_ptr common_public_rsa_key_; ActorOwn public_rsa_key_watchdog_; std::mutex main_dc_id_mutex_; + std::shared_ptr td_guard_; Status wait_dc_init(DcId dc_id, bool force); bool is_dc_inited(int32 raw_dc_id);