Session: cache salt for temp keys, better on_mtproto_header

GitOrigin-RevId: 02ba02aaf86e36402536d98d7bb80fd198c68378

td/telegram/ConfigManager.cpp

@@ -295,7 +295,8 @@ ActorOwn<> get_full_config(DcId dc_id, IPAddress ip_address, Promise<FullConfig>
         int_dc_id += 10000;
       }
       session_ = create_actor<Session>("ConfigSession", std::move(session_callback), std::move(auth_data), int_dc_id,
-                                       false /*is_main*/, true /*use_pfs*/, false /*is_cdn*/, mtproto::AuthKey());
+                                       false /*is_main*/, true /*use_pfs*/, false /*is_cdn*/, mtproto::AuthKey(),
+                                       std::vector<mtproto::ServerSalt>());
       auto query = G()->net_query_creator().create(create_storer(telegram_api::help_getConfig()), DcId::empty(),
                                                    NetQuery::Type::Common, NetQuery::AuthFlag::Off,
                                                    NetQuery::GzipFlag::On, 60 * 60 * 24);

td/telegram/net/Session.cpp

@@ -107,7 +107,8 @@ class GenAuthKeyActor : public Actor {
 }  // namespace detail
 
 Session::Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 dc_id,
-                 bool is_main, bool use_pfs, bool is_cdn, const mtproto::AuthKey &tmp_auth_key)
+                 bool is_main, bool use_pfs, bool is_cdn, const mtproto::AuthKey &tmp_auth_key,
+                 std::vector<mtproto::ServerSalt> server_salts)
     : dc_id_(dc_id), is_main_(is_main), is_cdn_(is_cdn) {
   VLOG(dc) << "Start connection";
 
@@ -118,6 +119,7 @@ Session::Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared>
   auth_data_.set_future_salts(shared_auth_data_->get_future_salts(), Time::now());
   if (use_pfs && !tmp_auth_key.empty()) {
     auth_data_.set_tmp_auth_key(tmp_auth_key);
+    auth_data_.set_future_salts(std::move(server_salts), Time::now());
   }
   uint64 session_id = 0;
   Random::secure_bytes(reinterpret_cast<uint8 *>(&session_id), sizeof(session_id));
@@ -376,6 +378,7 @@ void Session::on_tmp_auth_key_updated() {
 
 void Session::on_server_salt_updated() {
   if (auth_data_.use_pfs()) {
+    callback_->on_server_salt_updated(auth_data_.get_future_salts());
     return;
   }
   shared_auth_data_->set_future_salts(auth_data_.get_future_salts());

td/telegram/net/Session.h

@@ -56,11 +56,14 @@ class Session final
     virtual void on_closed() = 0;
     virtual void request_raw_connection(Promise<std::unique_ptr<mtproto::RawConnection>>) = 0;
     virtual void on_tmp_auth_key_updated(mtproto::AuthKey auth_key) = 0;
+    virtual void on_server_salt_updated(std::vector<mtproto::ServerSalt> server_salts) {
+    }
     // one still have to call close after on_closed
   };
 
   Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 dc_id, bool is_main,
-          bool use_pfs, bool is_cdn, const mtproto::AuthKey &tmp_auth_key);
+          bool use_pfs, bool is_cdn, const mtproto::AuthKey &tmp_auth_key,
+          std::vector<mtproto::ServerSalt> server_salts);
   void send(NetQueryPtr &&query);
   void on_network(bool network_flag, uint32 network_generation);
   void on_online(bool online_flag);

td/telegram/net/SessionMultiProxy.cpp

@@ -88,7 +88,9 @@ void SessionMultiProxy::update_options(int32 session_count, bool use_pfs) {
 }
 
 void SessionMultiProxy::update_mtproto_header() {
-  init();
+  for (auto &session : sessions_) {
+    send_closure_later(session, &SessionProxy::update_mtproto_header);
+  }
 }
 
 void SessionMultiProxy::start_up() {

td/telegram/net/SessionProxy.cpp

@@ -47,6 +47,10 @@ class SessionCallback : public Session::Callback {
     send_closure(parent_, &SessionProxy::on_tmp_auth_key_updated, std::move(auth_key));
   }
 
+  void on_server_salt_updated(std::vector<mtproto::ServerSalt> server_salts) override {
+    send_closure(parent_, &SessionProxy::on_server_salt_updated, std::move(server_salts));
+  }
+
  private:
   ActorShared<SessionProxy> parent_;
   DcId dc_id_;
@@ -128,6 +132,11 @@ void SessionProxy::on_failed() {
   open_session();
 }
 
+void SessionProxy::update_mtproto_header() {
+  close_session();
+  open_session();
+}
+
 void SessionProxy::on_closed() {
 }
 
@@ -154,7 +163,7 @@ void SessionProxy::open_session(bool force) {
   session_ = create_actor<Session>(
       name,
       make_unique<SessionCallback>(actor_shared(this, session_generation_), dc_id, allow_media_only_, is_media_, hash),
-      auth_data_, int_dc_id, is_main_, use_pfs_, is_cdn_, tmp_auth_key_);
+      auth_data_, int_dc_id, is_main_, use_pfs_, is_cdn_, tmp_auth_key_, server_salts_);
 }
 
 void SessionProxy::update_auth_state() {
@@ -187,5 +196,8 @@ void SessionProxy::on_tmp_auth_key_updated(mtproto::AuthKey auth_key) {
   LOG(WARNING) << "tmp_auth_key " << auth_key.id() << ": " << state;
   tmp_auth_key_ = std::move(auth_key);
 }
+void SessionProxy::on_server_salt_updated(std::vector<mtproto::ServerSalt> server_salts) {
+  server_salts_ = std::move(server_salts);
+}
 
 }  // namespace td

td/telegram/net/SessionProxy.h

@@ -25,6 +25,7 @@ class SessionProxy : public Actor {
 
   void send(NetQueryPtr query);
   void update_main_flag(bool is_main);
+  void update_mtproto_header();
 
  private:
   std::shared_ptr<AuthDataShared> auth_data_;
@@ -34,6 +35,7 @@ class SessionProxy : public Actor {
   bool is_media_;
   bool use_pfs_;
   mtproto::AuthKey tmp_auth_key_;
+  std::vector<mtproto::ServerSalt> server_salts_;
   bool need_wait_for_key_;
   bool is_cdn_;
   ActorOwn<Session> session_;
@@ -47,6 +49,7 @@ class SessionProxy : public Actor {
 
   void update_auth_state();
   void on_tmp_auth_key_updated(mtproto::AuthKey auth_key);
+  void on_server_salt_updated(std::vector<mtproto::ServerSalt> server_salts);
 
   void start_up() override;
   void tear_down() override;