Persist temporary keys for main sessions if multiple sessions enabled.

This commit is contained in:
levlam 2023-05-04 15:54:21 +03:00
parent 796a1f83c3
commit d3ade3d64d
3 changed files with 33 additions and 5 deletions

View File

@ -141,9 +141,10 @@ void SessionMultiProxy::init() {
uint32 generation_; uint32 generation_;
int32 session_id_; int32 session_id_;
}; };
info.proxy = create_actor<SessionProxy>(name, make_unique<Callback>(actor_id(this), sessions_generation_, i), info.proxy =
auth_data_, is_primary_, is_main_, allow_media_only_, is_media_, create_actor<SessionProxy>(name, make_unique<Callback>(actor_id(this), sessions_generation_, i), auth_data_,
get_pfs_flag(), is_cdn_, need_destroy_auth_key_ && i == 0); is_primary_, is_main_, allow_media_only_, is_media_, get_pfs_flag(),
session_count_ > 1 && is_primary_, is_cdn_, need_destroy_auth_key_ && i == 0);
sessions_.push_back(std::move(info)); sessions_.push_back(std::move(info));
} }
} }

View File

@ -80,7 +80,7 @@ class SessionCallback final : public Session::Callback {
SessionProxy::SessionProxy(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, SessionProxy::SessionProxy(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data,
bool is_primary, bool is_main, bool allow_media_only, bool is_media, bool use_pfs, bool is_primary, bool is_main, bool allow_media_only, bool is_media, bool use_pfs,
bool is_cdn, bool need_destroy) bool persist_tmp_auth_key, bool is_cdn, bool need_destroy)
: callback_(std::move(callback)) : callback_(std::move(callback))
, auth_data_(std::move(shared_auth_data)) , auth_data_(std::move(shared_auth_data))
, is_primary_(is_primary) , is_primary_(is_primary)
@ -88,6 +88,7 @@ SessionProxy::SessionProxy(unique_ptr<Callback> callback, std::shared_ptr<AuthDa
, allow_media_only_(allow_media_only) , allow_media_only_(allow_media_only)
, is_media_(is_media) , is_media_(is_media)
, use_pfs_(use_pfs) , use_pfs_(use_pfs)
, persist_tmp_auth_key_(use_pfs && persist_tmp_auth_key)
, is_cdn_(is_cdn) , is_cdn_(is_cdn)
, need_destroy_(need_destroy) { , need_destroy_(need_destroy) {
} }
@ -110,6 +111,21 @@ void SessionProxy::start_up() {
}; };
auth_key_state_ = get_auth_key_state(auth_data_->get_auth_key()); auth_key_state_ = get_auth_key_state(auth_data_->get_auth_key());
auth_data_->add_auth_key_listener(make_unique<Listener>(actor_shared(this))); auth_data_->add_auth_key_listener(make_unique<Listener>(actor_shared(this)));
string saved_auth_key = G()->td_db()->get_binlog_pmc()->get(tmp_auth_key_key());
if (!saved_auth_key.empty()) {
if (persist_tmp_auth_key_) {
unserialize(tmp_auth_key_, saved_auth_key).ensure();
if (tmp_auth_key_.expires_at() < Time::now()) {
tmp_auth_key_ = {};
} else {
LOG(WARNING) << "Loaded tmp_auth_key " << tmp_auth_key_.id() << ": " << get_auth_key_state(tmp_auth_key_);
}
} else {
LOG(WARNING) << "Drop saved tmp_auth_key";
G()->td_db()->get_binlog_pmc()->erase(tmp_auth_key_key());
}
}
open_session(); open_session();
} }
@ -241,6 +257,13 @@ void SessionProxy::update_auth_key_state() {
void SessionProxy::on_tmp_auth_key_updated(mtproto::AuthKey auth_key) { void SessionProxy::on_tmp_auth_key_updated(mtproto::AuthKey auth_key) {
LOG(WARNING) << "Have tmp_auth_key " << auth_key.id() << ": " << get_auth_key_state(auth_key); LOG(WARNING) << "Have tmp_auth_key " << auth_key.id() << ": " << get_auth_key_state(auth_key);
tmp_auth_key_ = std::move(auth_key); tmp_auth_key_ = std::move(auth_key);
if (persist_tmp_auth_key_) {
G()->td_db()->get_binlog_pmc()->set(tmp_auth_key_key(), serialize(tmp_auth_key_));
}
}
string SessionProxy::tmp_auth_key_key() const {
return PSTRING() << "tmp_auth" << get_name();
} }
void SessionProxy::on_server_salt_updated(std::vector<mtproto::ServerSalt> server_salts) { void SessionProxy::on_server_salt_updated(std::vector<mtproto::ServerSalt> server_salts) {

View File

@ -30,7 +30,8 @@ class SessionProxy final : public Actor {
}; };
SessionProxy(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, bool is_primary, SessionProxy(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, bool is_primary,
bool is_main, bool allow_media_only, bool is_media, bool use_pfs, bool is_cdn, bool need_destroy); bool is_main, bool allow_media_only, bool is_media, bool use_pfs, bool persist_tmp_auth_key, bool is_cdn,
bool need_destroy);
void send(NetQueryPtr query); void send(NetQueryPtr query);
void update_main_flag(bool is_main); void update_main_flag(bool is_main);
@ -46,6 +47,7 @@ class SessionProxy final : public Actor {
bool allow_media_only_; bool allow_media_only_;
bool is_media_; bool is_media_;
bool use_pfs_; bool use_pfs_;
bool persist_tmp_auth_key_;
mtproto::AuthKey tmp_auth_key_; mtproto::AuthKey tmp_auth_key_;
std::vector<mtproto::ServerSalt> server_salts_; std::vector<mtproto::ServerSalt> server_salts_;
bool is_cdn_; bool is_cdn_;
@ -65,6 +67,8 @@ class SessionProxy final : public Actor {
void on_query_finished(); void on_query_finished();
string tmp_auth_key_key() const;
void start_up() final; void start_up() final;
void tear_down() final; void tear_down() final;
}; };