From d3ade3d64dfcd20c7bcd4faeeb6837d62af3c351 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Thu, 4 May 2023 15:54:21 +0300
Subject: [PATCH] Persist temporary keys for main sessions if multiple sessions
 enabled.

---
 td/telegram/net/SessionMultiProxy.cpp |  7 ++++---
 td/telegram/net/SessionProxy.cpp      | 25 ++++++++++++++++++++++++-
 td/telegram/net/SessionProxy.h        |  6 +++++-
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/td/telegram/net/SessionMultiProxy.cpp b/td/telegram/net/SessionMultiProxy.cpp
index be0c9a1b9..75a12ef2f 100644
--- a/td/telegram/net/SessionMultiProxy.cpp
+++ b/td/telegram/net/SessionMultiProxy.cpp
@@ -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));
   }
 }
diff --git a/td/telegram/net/SessionProxy.cpp b/td/telegram/net/SessionProxy.cpp
index 47971550b..defacc0d0 100644
--- a/td/telegram/net/SessionProxy.cpp
+++ b/td/telegram/net/SessionProxy.cpp
@@ -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) {
diff --git a/td/telegram/net/SessionProxy.h b/td/telegram/net/SessionProxy.h
index eb2f48e89..b01c44427 100644
--- a/td/telegram/net/SessionProxy.h
+++ b/td/telegram/net/SessionProxy.h
@@ -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;
 };