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_;
int32 session_id_;
};
info.proxy = create_actor<SessionProxy>(name, make_unique<Callback>(actor_id(this), sessions_generation_, i),
auth_data_, is_primary_, is_main_, allow_media_only_, is_media_,
get_pfs_flag(), is_cdn_, need_destroy_auth_key_ && i == 0);
info.proxy =
create_actor<SessionProxy>(name, make_unique<Callback>(actor_id(this), sessions_generation_, i), auth_data_,
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));
}
}

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,
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))
, auth_data_(std::move(shared_auth_data))
, is_primary_(is_primary)
@ -88,6 +88,7 @@ SessionProxy::SessionProxy(unique_ptr<Callback> callback, std::shared_ptr<AuthDa
, allow_media_only_(allow_media_only)
, is_media_(is_media)
, use_pfs_(use_pfs)
, persist_tmp_auth_key_(use_pfs && persist_tmp_auth_key)
, is_cdn_(is_cdn)
, 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_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();
}
@ -241,6 +257,13 @@ void SessionProxy::update_auth_key_state() {
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);
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) {

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,
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 update_main_flag(bool is_main);
@ -46,6 +47,7 @@ class SessionProxy final : public Actor {
bool allow_media_only_;
bool is_media_;
bool use_pfs_;
bool persist_tmp_auth_key_;
mtproto::AuthKey tmp_auth_key_;
std::vector<mtproto::ServerSalt> server_salts_;
bool is_cdn_;
@ -65,6 +67,8 @@ class SessionProxy final : public Actor {
void on_query_finished();
string tmp_auth_key_key() const;
void start_up() final;
void tear_down() final;
};