From 23b0a91e88493028380394ff9f5cc6e15708148c Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Fri, 9 Jul 2021 02:17:26 +0300
Subject: [PATCH] Fix PublicRsaKeyWatchdog.

---
 td/telegram/net/AuthDataShared.cpp       | 4 ++--
 td/telegram/net/PublicRsaKeyShared.cpp   | 4 +++-
 td/telegram/net/PublicRsaKeyWatchdog.cpp | 4 ++++
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/td/telegram/net/AuthDataShared.cpp b/td/telegram/net/AuthDataShared.cpp
index 2dc28eaca..164dac567 100644
--- a/td/telegram/net/AuthDataShared.cpp
+++ b/td/telegram/net/AuthDataShared.cpp
@@ -90,10 +90,10 @@ class AuthDataSharedImpl final : public AuthDataShared {
   std::shared_ptr<Guard> guard_;
   RwMutex rw_mutex_;
 
-  string auth_key_key() {
+  string auth_key_key() const {
     return PSTRING() << "auth" << dc_id_.get_raw_id();
   }
-  string future_salts_key() {
+  string future_salts_key() const {
     return PSTRING() << "salt" << dc_id_.get_raw_id();
   }
 
diff --git a/td/telegram/net/PublicRsaKeyShared.cpp b/td/telegram/net/PublicRsaKeyShared.cpp
index 294658ade..54b893012 100644
--- a/td/telegram/net/PublicRsaKeyShared.cpp
+++ b/td/telegram/net/PublicRsaKeyShared.cpp
@@ -76,10 +76,13 @@ Result<mtproto::PublicRsaKeyInterface::RsaKey> PublicRsaKeyShared::get_rsa_key(c
 
 void PublicRsaKeyShared::drop_keys() {
   if (dc_id_.is_empty()) {
+    // not CDN
     return;
   }
   auto lock = rw_mutex_.lock_write();
+  LOG(INFO) << "Drop " << keys_.size() << " keys for " << dc_id_;
   keys_.clear();
+  notify();
 }
 
 bool PublicRsaKeyShared::has_keys() {
@@ -104,7 +107,6 @@ mtproto::PublicRsaKeyInterface::RsaKey *PublicRsaKeyShared::get_rsa_key_unsafe(i
 }
 
 void PublicRsaKeyShared::notify() {
-  auto lock = rw_mutex_.lock_read();
   td::remove_if(listeners_, [&](auto &listener) { return !listener->notify(); });
 }
 
diff --git a/td/telegram/net/PublicRsaKeyWatchdog.cpp b/td/telegram/net/PublicRsaKeyWatchdog.cpp
index 045e287fc..131997206 100644
--- a/td/telegram/net/PublicRsaKeyWatchdog.cpp
+++ b/td/telegram/net/PublicRsaKeyWatchdog.cpp
@@ -55,6 +55,7 @@ void PublicRsaKeyWatchdog::loop() {
     return;
   }
   if (Time::now_cached() < flood_control_.get_wakeup_at()) {
+    set_timeout_in(flood_control_.get_wakeup_at() - Time::now_cached() + 0.01);
     return;
   }
   bool ok = true;
@@ -78,6 +79,7 @@ void PublicRsaKeyWatchdog::on_result(NetQueryPtr net_query) {
   yield();
   if (net_query->is_error()) {
     LOG(ERROR) << "Receive error for GetCdnConfig: " << net_query->move_as_error();
+    loop();
     return;
   }
 
@@ -88,11 +90,13 @@ void PublicRsaKeyWatchdog::on_result(NetQueryPtr net_query) {
 
 void PublicRsaKeyWatchdog::sync(BufferSlice cdn_config_serialized) {
   if (cdn_config_serialized.empty()) {
+    loop();
     return;
   }
   auto r_keys = fetch_result<telegram_api::help_getCdnConfig>(cdn_config_serialized);
   if (r_keys.is_error()) {
     LOG(WARNING) << "Failed to deserialize help_getCdnConfig (probably not a problem) " << r_keys.error();
+    loop();
     return;
   }
   cdn_config_ = r_keys.move_as_ok();