tdlight-team/tdlight
tdlight-team
/
tdlight
1
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Use separate implementations for main and CDN RSA keys.

This commit is contained in:
levlam 2024-01-22 02:15:26 +03:00
parent e52d08da74
commit 0dd1934090
12 changed files with 155 additions and 74 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
CMakeLists.txt
View File

@ -429,7 +429,8 @@ set(TDLIB_SOURCE
td/telegram/net/NetQueryStats.cpp
td/telegram/net/NetStatsManager.cpp
td/telegram/net/Proxy.cpp
td/telegram/net/PublicRsaKeyShared.cpp
td/telegram/net/PublicRsaKeySharedCdn.cpp
td/telegram/net/PublicRsaKeySharedMain.cpp
td/telegram/net/PublicRsaKeyWatchdog.cpp
td/telegram/net/Session.cpp
td/telegram/net/SessionMultiProxy.cpp
@ -735,7 +736,8 @@ set(TDLIB_SOURCE
td/telegram/net/NetStatsManager.h
td/telegram/net/NetType.h
td/telegram/net/Proxy.h
td/telegram/net/PublicRsaKeyShared.h
td/telegram/net/PublicRsaKeySharedCdn.h
td/telegram/net/PublicRsaKeySharedMain.h
td/telegram/net/PublicRsaKeyWatchdog.h
td/telegram/net/Session.h
td/telegram/net/SessionProxy.h

7
td/telegram/ConfigManager.cpp
View File

@ -21,7 +21,7 @@
#include "td/telegram/net/NetQuery.h"
#include "td/telegram/net/NetQueryDispatcher.h"
#include "td/telegram/net/NetType.h"
#include "td/telegram/net/PublicRsaKeyShared.h"
#include "td/telegram/net/PublicRsaKeySharedMain.h"
#include "td/telegram/net/Session.h"
#include "td/telegram/OptionManager.h"
#include "td/telegram/Premium.h"
@ -493,7 +493,7 @@ static ActorOwn<> get_full_config(DcOption option, Promise<tl_object_ptr<telegra
private:
DcId dc_id_;
std::shared_ptr<mtproto::PublicRsaKeyInterface> public_rsa_key_ =
std::make_shared<PublicRsaKeyShared>(DcId::empty(), G()->is_test_dc());
std::make_shared<PublicRsaKeySharedMain>(G()->is_test_dc());
std::vector<unique_ptr<Listener>> auth_key_listeners_;
void notify() {
@ -1497,7 +1497,8 @@ void ConfigManager::process_app_config(tl_object_ptr<telegram_api::JSONValue> &c
key == "getfile_experimental_params" || key == "message_animated_emoji_max" ||
key == "stickers_emoji_cache_time" || key == "stories_export_nopublic_link" || key == "test" ||
key == "upload_max_fileparts_default" || key == "upload_max_fileparts_premium" ||
key == "wallet_blockchain_name" || key == "wallet_config" || key == "wallet_enabled" || key == "channel_color_level_min") {
key == "wallet_blockchain_name" || key == "wallet_config" || key == "wallet_enabled" ||
key == "channel_color_level_min") {
continue;
}
if (key == "ignore_restriction_reasons") {

6
td/telegram/Td.cpp
View File

@ -100,7 +100,7 @@
#include "td/telegram/net/NetStatsManager.h"
#include "td/telegram/net/NetType.h"
#include "td/telegram/net/Proxy.h"
#include "td/telegram/net/PublicRsaKeyShared.h"
#include "td/telegram/net/PublicRsaKeySharedMain.h"
#include "td/telegram/net/TempAuthKeyWatchdog.h"
#include "td/telegram/NotificationGroupId.h"
#include "td/telegram/NotificationId.h"
@ -529,11 +529,11 @@ class TestProxyRequest final : public RequestOnceActor {
return nullptr;
}
mtproto::PublicRsaKeyInterface *get_public_rsa_key_interface() final {
return &public_rsa_key;
return &public_rsa_key_;
}
private:
PublicRsaKeyShared public_rsa_key{DcId::empty(), false};
PublicRsaKeySharedMain public_rsa_key_{false};
};
auto handshake = make_unique<mtproto::AuthKeyHandshake>(dc_id_, 3600);
auto data = r_data.move_as_ok();

12
td/telegram/net/NetQueryDispatcher.cpp
View File

@ -11,7 +11,8 @@
#include "td/telegram/net/DcAuthManager.h"
#include "td/telegram/net/NetQuery.h"
#include "td/telegram/net/NetQueryDelayer.h"
#include "td/telegram/net/PublicRsaKeyShared.h"
#include "td/telegram/net/PublicRsaKeySharedCdn.h"
#include "td/telegram/net/PublicRsaKeySharedMain.h"
#include "td/telegram/net/PublicRsaKeyWatchdog.h"
#include "td/telegram/net/SessionMultiProxy.h"
#include "td/telegram/SequenceDispatcher.h"
@ -150,13 +151,14 @@ Status NetQueryDispatcher::wait_dc_init(DcId dc_id, bool force) {
}
// init dc
dc.id_ = dc_id;
decltype(common_public_rsa_key_) public_rsa_key;
std::shared_ptr<mtproto::PublicRsaKeyInterface> public_rsa_key;
bool is_cdn = false;
if (dc_id.is_internal()) {
public_rsa_key = common_public_rsa_key_;
} else {
public_rsa_key = std::make_shared<PublicRsaKeyShared>(dc_id, G()->is_test_dc());
send_closure_later(public_rsa_key_watchdog_, &PublicRsaKeyWatchdog::add_public_rsa_key, public_rsa_key);
auto public_rsa_key_cdn = std::make_shared<PublicRsaKeySharedCdn>(dc_id);
send_closure_later(public_rsa_key_watchdog_, &PublicRsaKeyWatchdog::add_public_rsa_key, public_rsa_key_cdn);
public_rsa_key = public_rsa_key_cdn;
is_cdn = true;
}
auto auth_data = AuthDataShared::create(dc_id, std::move(public_rsa_key), td_guard_);
@ -299,7 +301,7 @@ NetQueryDispatcher::NetQueryDispatcher(const std::function<ActorShared<>()> &cre
LOG(INFO) << tag("main_dc_id", main_dc_id_.load(std::memory_order_relaxed));
delayer_ = create_actor<NetQueryDelayer>("NetQueryDelayer", create_reference());
dc_auth_manager_ = create_actor<DcAuthManager>("DcAuthManager", create_reference());
common_public_rsa_key_ = std::make_shared<PublicRsaKeyShared>(DcId::empty(), G()->is_test_dc());
common_public_rsa_key_ = std::make_shared<PublicRsaKeySharedMain>(G()->is_test_dc());
public_rsa_key_watchdog_ = create_actor<PublicRsaKeyWatchdog>("PublicRsaKeyWatchdog", create_reference());
sequence_dispatcher_ = MultiSequenceDispatcher::create("MultiSequenceDispatcher");

4
td/telegram/net/NetQueryDispatcher.h
View File

@ -27,7 +27,7 @@ namespace td {
class DcAuthManager;
class MultiSequenceDispatcher;
class NetQueryDelayer;
class PublicRsaKeyShared;
class PublicRsaKeySharedMain;
class PublicRsaKeyWatchdog;
class SessionMultiProxy;
@ -79,7 +79,7 @@ class NetQueryDispatcher {
#else
std::atomic<int32> main_dc_id_{1};
#endif
std::shared_ptr<PublicRsaKeyShared> common_public_rsa_key_;
std::shared_ptr<PublicRsaKeySharedMain> common_public_rsa_key_;
ActorOwn<PublicRsaKeyWatchdog> public_rsa_key_watchdog_;
std::mutex main_dc_id_mutex_;
std::shared_ptr<Guard> td_guard_;

76
td/telegram/net/PublicRsaKeySharedCdn.cpp Normal file
View File

@ -0,0 +1,76 @@
//
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2024
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
#include "td/telegram/net/PublicRsaKeySharedCdn.h"
#include "td/utils/algorithm.h"
#include "td/utils/format.h"
#include "td/utils/logging.h"
#include "td/utils/Slice.h"
#include "td/utils/SliceBuilder.h"
#include <algorithm>
namespace td {
PublicRsaKeySharedCdn::PublicRsaKeySharedCdn(DcId dc_id) : dc_id_(dc_id) {
CHECK(!dc_id_.is_empty());
CHECK(!dc_id_.is_internal());
}
void PublicRsaKeySharedCdn::add_rsa(mtproto::RSA rsa) {
auto lock = rw_mutex_.lock_write();
auto fingerprint = rsa.get_fingerprint();
if (get_rsa_key_unsafe(fingerprint) != nullptr) {
return;
}
keys_.push_back(RsaKey{std::move(rsa), fingerprint});
}
Result<mtproto::PublicRsaKeyInterface::RsaKey> PublicRsaKeySharedCdn::get_rsa_key(const vector<int64> &fingerprints) {
auto lock = rw_mutex_.lock_read();
for (auto fingerprint : fingerprints) {
auto *rsa_key = get_rsa_key_unsafe(fingerprint);
if (rsa_key != nullptr) {
return RsaKey{rsa_key->rsa.clone(), fingerprint};
}
}
return Status::Error(PSLICE() << "Unknown fingerprints " << format::as_array(fingerprints));
}
void PublicRsaKeySharedCdn::drop_keys() {
LOG(INFO) << "Drop " << keys_.size() << " keys for " << dc_id_;
auto lock = rw_mutex_.lock_write();
keys_.clear();
notify();
}
bool PublicRsaKeySharedCdn::has_keys() {
auto lock = rw_mutex_.lock_read();
return !keys_.empty();
}
void PublicRsaKeySharedCdn::add_listener(unique_ptr<Listener> listener) {
if (listener->notify()) {
auto lock = rw_mutex_.lock_write();
listeners_.push_back(std::move(listener));
}
}
mtproto::PublicRsaKeyInterface::RsaKey *PublicRsaKeySharedCdn::get_rsa_key_unsafe(int64 fingerprint) {
auto it = std::find_if(keys_.begin(), keys_.end(),
[fingerprint](const auto &value) { return value.fingerprint == fingerprint; });
if (it == keys_.end()) {
return nullptr;
}
return &*it;
}
void PublicRsaKeySharedCdn::notify() {
td::remove_if(listeners_, [&](auto &listener) { return !listener->notify(); });
}
} // namespace td

11
td/telegram/net/PublicRsaKeyShared.h → td/telegram/net/PublicRsaKeySharedCdn.h
View File

@ -16,9 +16,9 @@
namespace td {
class PublicRsaKeyShared final : public mtproto::PublicRsaKeyInterface {
class PublicRsaKeySharedCdn final : public mtproto::PublicRsaKeyInterface {
public:
PublicRsaKeyShared(DcId dc_id, bool is_test);
explicit PublicRsaKeySharedCdn(DcId dc_id);
class Listener {
public:
@ -32,8 +32,11 @@ class PublicRsaKeyShared final : public mtproto::PublicRsaKeyInterface {
};
void add_rsa(mtproto::RSA rsa);
Result<RsaKey> get_rsa_key(const vector<int64> &fingerprints) final;
void drop_keys() final;
bool has_keys();
void add_listener(unique_ptr<Listener> listener);
@ -44,8 +47,8 @@ class PublicRsaKeyShared final : public mtproto::PublicRsaKeyInterface {
private:
DcId dc_id_;
std::vector<RsaKey> keys_;
std::vector<unique_ptr<Listener>> listeners_;
vector<RsaKey> keys_;
vector<unique_ptr<Listener>> listeners_;
RwMutex rw_mutex_;
RsaKey *get_rsa_key_unsafe(int64 fingerprint);

59
td/telegram/net/PublicRsaKeyShared.cpp → td/telegram/net/PublicRsaKeySharedMain.cpp
View File

@ -4,29 +4,29 @@
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
#include "td/telegram/net/PublicRsaKeyShared.h"
#include "td/telegram/net/PublicRsaKeySharedMain.h"
#include "td/utils/algorithm.h"
#include "td/utils/format.h"
#include "td/utils/logging.h"
#include "td/utils/Slice.h"
#include "td/utils/SliceBuilder.h"
#include "td/utils/Status.h"
#include <algorithm>
namespace td {
PublicRsaKeyShared::PublicRsaKeyShared(DcId dc_id, bool is_test) : dc_id_(dc_id) {
if (!dc_id_.is_empty()) {
return;
}
PublicRsaKeySharedMain::PublicRsaKeySharedMain(bool is_test) {
auto add_pem = [this](CSlice pem) {
auto r_rsa = mtproto::RSA::from_pem_public_key(pem);
LOG_CHECK(r_rsa.is_ok()) << r_rsa.error() << " " << pem;
if (r_rsa.is_ok()) {
add_rsa(r_rsa.move_as_ok());
auto rsa = r_rsa.move_as_ok();
auto fingerprint = rsa.get_fingerprint();
if (get_rsa_key_unsafe(fingerprint) != nullptr) {
return;
}
keys_.push_back(RsaKey{std::move(rsa), fingerprint});
}
};
@ -54,19 +54,9 @@ PublicRsaKeyShared::PublicRsaKeyShared(DcId dc_id, bool is_test) : dc_id_(dc_id)
"-----END RSA PUBLIC KEY-----");
}
void PublicRsaKeyShared::add_rsa(mtproto::RSA rsa) {
auto lock = rw_mutex_.lock_write();
auto fingerprint = rsa.get_fingerprint();
if (get_rsa_key_unsafe(fingerprint) != nullptr) {
return;
}
keys_.push_back(RsaKey{std::move(rsa), fingerprint});
}
Result<mtproto::PublicRsaKeyInterface::RsaKey> PublicRsaKeyShared::get_rsa_key(const vector<int64> &fingerprints) {
auto lock = rw_mutex_.lock_read();
Result<mtproto::PublicRsaKeyInterface::RsaKey> PublicRsaKeySharedMain::get_rsa_key(const vector<int64> &fingerprints) {
for (auto fingerprint : fingerprints) {
auto *rsa_key = get_rsa_key_unsafe(fingerprint);
const auto *rsa_key = get_rsa_key_unsafe(fingerprint);
if (rsa_key != nullptr) {
return RsaKey{rsa_key->rsa.clone(), fingerprint};
}
@ -74,30 +64,11 @@ Result<mtproto::PublicRsaKeyInterface::RsaKey> PublicRsaKeyShared::get_rsa_key(c
return Status::Error(PSLICE() << "Unknown fingerprints " << format::as_array(fingerprints));
}
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();
void PublicRsaKeySharedMain::drop_keys() {
// nothing to do
}
bool PublicRsaKeyShared::has_keys() {
auto lock = rw_mutex_.lock_read();
return !keys_.empty();
}
void PublicRsaKeyShared::add_listener(unique_ptr<Listener> listener) {
if (listener->notify()) {
auto lock = rw_mutex_.lock_write();
listeners_.push_back(std::move(listener));
}
}
mtproto::PublicRsaKeyInterface::RsaKey *PublicRsaKeyShared::get_rsa_key_unsafe(int64 fingerprint) {
const mtproto::PublicRsaKeyInterface::RsaKey *PublicRsaKeySharedMain::get_rsa_key_unsafe(int64 fingerprint) const {
auto it = std::find_if(keys_.begin(), keys_.end(),
[fingerprint](const auto &value) { return value.fingerprint == fingerprint; });
if (it == keys_.end()) {
@ -106,8 +77,4 @@ mtproto::PublicRsaKeyInterface::RsaKey *PublicRsaKeyShared::get_rsa_key_unsafe(i
return &*it;
}
void PublicRsaKeyShared::notify() {
td::remove_if(listeners_, [&](auto &listener) { return !listener->notify(); });
}
} // namespace td

30
td/telegram/net/PublicRsaKeySharedMain.h Normal file
View File

@ -0,0 +1,30 @@
//
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2024
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
#pragma once
#include "td/mtproto/RSA.h"
#include "td/utils/common.h"
#include "td/utils/Status.h"
namespace td {
class PublicRsaKeySharedMain final : public mtproto::PublicRsaKeyInterface {
public:
explicit PublicRsaKeySharedMain(bool is_test);
Result<RsaKey> get_rsa_key(const vector<int64> &fingerprints) final;
void drop_keys() final;
private:
vector<RsaKey> keys_;
const RsaKey *get_rsa_key_unsafe(int64 fingerprint) const;
};
} // namespace td

6
td/telegram/net/PublicRsaKeyWatchdog.cpp
View File

@ -22,8 +22,8 @@ namespace td {
PublicRsaKeyWatchdog::PublicRsaKeyWatchdog(ActorShared<> parent) : parent_(std::move(parent)) {
}
void PublicRsaKeyWatchdog::add_public_rsa_key(std::shared_ptr<PublicRsaKeyShared> key) {
class Listener final : public PublicRsaKeyShared::Listener {
void PublicRsaKeyWatchdog::add_public_rsa_key(std::shared_ptr<PublicRsaKeySharedCdn> key) {
class Listener final : public PublicRsaKeySharedCdn::Listener {
public:
explicit Listener(ActorId<PublicRsaKeyWatchdog> parent) : parent_(std::move(parent)) {
}
@ -119,7 +119,7 @@ void PublicRsaKeyWatchdog::sync(BufferSlice cdn_config_serialized) {
}
}
void PublicRsaKeyWatchdog::sync_key(std::shared_ptr<PublicRsaKeyShared> &key) {
void PublicRsaKeyWatchdog::sync_key(std::shared_ptr<PublicRsaKeySharedCdn> &key) {
if (!cdn_config_) {
return;
}

8
td/telegram/net/PublicRsaKeyWatchdog.h
View File

@ -9,7 +9,7 @@
#include "td/telegram/net/NetActor.h"
#include "td/telegram/net/NetQueryCreator.h"
#include "td/telegram/net/NetQueryDispatcher.h"
#include "td/telegram/net/PublicRsaKeyShared.h"
#include "td/telegram/net/PublicRsaKeySharedCdn.h"
#include "td/telegram/telegram_api.h"
#include "td/actor/actor.h"
@ -26,11 +26,11 @@ class PublicRsaKeyWatchdog final : public NetActor {
public:
explicit PublicRsaKeyWatchdog(ActorShared<> parent);
void add_public_rsa_key(std::shared_ptr<PublicRsaKeyShared> key);
void add_public_rsa_key(std::shared_ptr<PublicRsaKeySharedCdn> key);
private:
ActorShared<> parent_;
vector<std::shared_ptr<PublicRsaKeyShared>> keys_;
vector<std::shared_ptr<PublicRsaKeySharedCdn>> keys_;
tl_object_ptr<telegram_api::cdnConfig> cdn_config_;
FloodControlStrict flood_control_;
bool has_query_{false};
@ -41,7 +41,7 @@ class PublicRsaKeyWatchdog final : public NetActor {
void on_result(NetQueryPtr net_query) final;
void sync(BufferSlice cdn_config_serialized);
void sync_key(std::shared_ptr<PublicRsaKeyShared> &key);
void sync_key(std::shared_ptr<PublicRsaKeySharedCdn> &key);
};
} // namespace td

4
test/mtproto.cpp
View File

@ -6,7 +6,7 @@
//
#include "td/telegram/ConfigManager.h"
#include "td/telegram/net/DcId.h"
#include "td/telegram/net/PublicRsaKeyShared.h"
#include "td/telegram/net/PublicRsaKeySharedMain.h"
#include "td/telegram/net/Session.h"
#include "td/telegram/NotificationManager.h"
@ -304,7 +304,7 @@ class HandshakeContext final : public td::mtproto::AuthKeyHandshakeContext {
}
private:
td::PublicRsaKeyShared public_rsa_key{td::DcId::empty(), true};
td::PublicRsaKeySharedMain public_rsa_key{true};
};
class HandshakeTestActor final : public td::Actor {