Skip to content

Commit 1f116cd

Browse files
levlamlevlam
committed
Reuse the same PublicRsaKeySharedMain object.
1 parent 0dd1934 commit 1f116cd

7 files changed

+54
-63
lines changed

‎td/telegram/ConfigManager.cpp

+5-4
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,8 @@ static ActorOwn<> get_full_config(DcOption option, Promise<tl_object_ptr<telegra
442442

443443
class SimpleAuthData final : public AuthDataShared {
444444
public:
445-
explicit SimpleAuthData(DcId dc_id) : dc_id_(dc_id) {
445+
explicit SimpleAuthData(DcId dc_id)
446+
: dc_id_(dc_id), public_rsa_key_(PublicRsaKeySharedMain::create(G()->is_test_dc())) {
446447
}
447448
DcId dc_id() const final {
448449
return dc_id_;
@@ -492,10 +493,9 @@ static ActorOwn<> get_full_config(DcOption option, Promise<tl_object_ptr<telegra
492493

493494
private:
494495
DcId dc_id_;
495-
std::shared_ptr<mtproto::PublicRsaKeyInterface> public_rsa_key_ =
496-
std::make_shared<PublicRsaKeySharedMain>(G()->is_test_dc());
496+
std::shared_ptr<mtproto::PublicRsaKeyInterface> public_rsa_key_;
497+
vector<unique_ptr<Listener>> auth_key_listeners_;
497498

498-
std::vector<unique_ptr<Listener>> auth_key_listeners_;
499499
void notify() {
500500
td::remove_if(auth_key_listeners_, [&](auto &listener) {
501501
CHECK(listener != nullptr);
@@ -506,6 +506,7 @@ static ActorOwn<> get_full_config(DcOption option, Promise<tl_object_ptr<telegra
506506
string auth_key_key() const {
507507
return PSTRING() << "config_recovery_auth" << dc_id().get_raw_id();
508508
}
509+
509510
string future_salts_key() const {
510511
return PSTRING() << "config_recovery_salt" << dc_id().get_raw_id();
511512
}

‎td/telegram/Td.cpp

+2-2
Original file line numberDiff line numberDiff line change
@@ -529,11 +529,11 @@ class TestProxyRequest final : public RequestOnceActor {
529529
return nullptr;
530530
}
531531
mtproto::PublicRsaKeyInterface *get_public_rsa_key_interface() final {
532-
return &public_rsa_key_;
532+
return public_rsa_key_.get();
533533
}
534534

535535
private:
536-
PublicRsaKeySharedMain public_rsa_key_{false};
536+
std::shared_ptr<mtproto::PublicRsaKeyInterface> public_rsa_key_ = PublicRsaKeySharedMain::create(false);
537537
};
538538
auto handshake = make_unique<mtproto::AuthKeyHandshake>(dc_id_, 3600);
539539
auto data = r_data.move_as_ok();

‎td/telegram/net/NetQueryDispatcher.cpp

+1-2
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ Status NetQueryDispatcher::wait_dc_init(DcId dc_id, bool force) {
154154
std::shared_ptr<mtproto::PublicRsaKeyInterface> public_rsa_key;
155155
bool is_cdn = false;
156156
if (dc_id.is_internal()) {
157-
public_rsa_key = common_public_rsa_key_;
157+
public_rsa_key = PublicRsaKeySharedMain::create(G()->is_test_dc());
158158
} else {
159159
auto public_rsa_key_cdn = std::make_shared<PublicRsaKeySharedCdn>(dc_id);
160160
send_closure_later(public_rsa_key_watchdog_, &PublicRsaKeyWatchdog::add_public_rsa_key, public_rsa_key_cdn);
@@ -301,7 +301,6 @@ NetQueryDispatcher::NetQueryDispatcher(const std::function<ActorShared<>()> &cre
301301
LOG(INFO) << tag("main_dc_id", main_dc_id_.load(std::memory_order_relaxed));
302302
delayer_ = create_actor<NetQueryDelayer>("NetQueryDelayer", create_reference());
303303
dc_auth_manager_ = create_actor<DcAuthManager>("DcAuthManager", create_reference());
304-
common_public_rsa_key_ = std::make_shared<PublicRsaKeySharedMain>(G()->is_test_dc());
305304
public_rsa_key_watchdog_ = create_actor<PublicRsaKeyWatchdog>("PublicRsaKeyWatchdog", create_reference());
306305
sequence_dispatcher_ = MultiSequenceDispatcher::create("MultiSequenceDispatcher");
307306

‎td/telegram/net/NetQueryDispatcher.h

-2
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@ namespace td {
2727
class DcAuthManager;
2828
class MultiSequenceDispatcher;
2929
class NetQueryDelayer;
30-
class PublicRsaKeySharedMain;
3130
class PublicRsaKeyWatchdog;
3231
class SessionMultiProxy;
3332

@@ -79,7 +78,6 @@ class NetQueryDispatcher {
7978
#else
8079
std::atomic<int32> main_dc_id_{1};
8180
#endif
82-
std::shared_ptr<PublicRsaKeySharedMain> common_public_rsa_key_;
8381
ActorOwn<PublicRsaKeyWatchdog> public_rsa_key_watchdog_;
8482
std::mutex main_dc_id_mutex_;
8583
std::shared_ptr<Guard> td_guard_;

‎td/telegram/net/PublicRsaKeySharedMain.cpp

+38-48
Original file line numberDiff line numberDiff line change
@@ -7,58 +7,57 @@
77
#include "td/telegram/net/PublicRsaKeySharedMain.h"
88

99
#include "td/utils/format.h"
10-
#include "td/utils/logging.h"
1110
#include "td/utils/Slice.h"
1211
#include "td/utils/SliceBuilder.h"
1312

14-
#include <algorithm>
15-
1613
namespace td {
1714

18-
PublicRsaKeySharedMain::PublicRsaKeySharedMain(bool is_test) {
19-
auto add_pem = [this](CSlice pem) {
20-
auto r_rsa = mtproto::RSA::from_pem_public_key(pem);
21-
LOG_CHECK(r_rsa.is_ok()) << r_rsa.error() << " " << pem;
22-
23-
if (r_rsa.is_ok()) {
24-
auto rsa = r_rsa.move_as_ok();
25-
auto fingerprint = rsa.get_fingerprint();
26-
if (get_rsa_key_unsafe(fingerprint) != nullptr) {
27-
return;
28-
}
29-
keys_.push_back(RsaKey{std::move(rsa), fingerprint});
30-
}
15+
std::shared_ptr<PublicRsaKeySharedMain> PublicRsaKeySharedMain::create(bool is_test) {
16+
auto add_pem = [](vector<RsaKey> &keys, CSlice pem) {
17+
auto rsa = mtproto::RSA::from_pem_public_key(pem).move_as_ok();
18+
auto fingerprint = rsa.get_fingerprint();
19+
keys.push_back(RsaKey{std::move(rsa), fingerprint});
3120
};
3221

3322
if (is_test) {
34-
add_pem(
35-
"-----BEGIN RSA PUBLIC KEY-----\n"
36-
"MIIBCgKCAQEAyMEdY1aR+sCR3ZSJrtztKTKqigvO/vBfqACJLZtS7QMgCGXJ6XIR\n"
37-
"yy7mx66W0/sOFa7/1mAZtEoIokDP3ShoqF4fVNb6XeqgQfaUHd8wJpDWHcR2OFwv\n"
38-
"plUUI1PLTktZ9uW2WE23b+ixNwJjJGwBDJPQEQFBE+vfmH0JP503wr5INS1poWg/\n"
39-
"j25sIWeYPHYeOrFp/eXaqhISP6G+q2IeTaWTXpwZj4LzXq5YOpk4bYEQ6mvRq7D1\n"
40-
"aHWfYmlEGepfaYR8Q0YqvvhYtMte3ITnuSJs171+GDqpdKcSwHnd6FudwGO4pcCO\n"
41-
"j4WcDuXc2CTHgH8gFTNhp/Y8/SpDOhvn9QIDAQAB\n"
42-
"-----END RSA PUBLIC KEY-----");
43-
return;
23+
static auto test_public_rsa_key = [&] {
24+
vector<RsaKey> keys;
25+
add_pem(keys,
26+
"-----BEGIN RSA PUBLIC KEY-----\n"
27+
"MIIBCgKCAQEAyMEdY1aR+sCR3ZSJrtztKTKqigvO/vBfqACJLZtS7QMgCGXJ6XIR\n"
28+
"yy7mx66W0/sOFa7/1mAZtEoIokDP3ShoqF4fVNb6XeqgQfaUHd8wJpDWHcR2OFwv\n"
29+
"plUUI1PLTktZ9uW2WE23b+ixNwJjJGwBDJPQEQFBE+vfmH0JP503wr5INS1poWg/\n"
30+
"j25sIWeYPHYeOrFp/eXaqhISP6G+q2IeTaWTXpwZj4LzXq5YOpk4bYEQ6mvRq7D1\n"
31+
"aHWfYmlEGepfaYR8Q0YqvvhYtMte3ITnuSJs171+GDqpdKcSwHnd6FudwGO4pcCO\n"
32+
"j4WcDuXc2CTHgH8gFTNhp/Y8/SpDOhvn9QIDAQAB\n"
33+
"-----END RSA PUBLIC KEY-----");
34+
return std::make_shared<PublicRsaKeySharedMain>(std::move(keys));
35+
}();
36+
return test_public_rsa_key;
37+
} else {
38+
static auto main_public_rsa_key = [&] {
39+
vector<RsaKey> keys;
40+
add_pem(keys,
41+
"-----BEGIN RSA PUBLIC KEY-----\n"
42+
"MIIBCgKCAQEA6LszBcC1LGzyr992NzE0ieY+BSaOW622Aa9Bd4ZHLl+TuFQ4lo4g\n"
43+
"5nKaMBwK/BIb9xUfg0Q29/2mgIR6Zr9krM7HjuIcCzFvDtr+L0GQjae9H0pRB2OO\n"
44+
"62cECs5HKhT5DZ98K33vmWiLowc621dQuwKWSQKjWf50XYFw42h21P2KXUGyp2y/\n"
45+
"+aEyZ+uVgLLQbRA1dEjSDZ2iGRy12Mk5gpYc397aYp438fsJoHIgJ2lgMv5h7WY9\n"
46+
"t6N/byY9Nw9p21Og3AoXSL2q/2IJ1WRUhebgAdGVMlV1fkuOQoEzR7EdpqtQD9Cs\n"
47+
"5+bfo3Nhmcyvk5ftB0WkJ9z6bNZ7yxrP8wIDAQAB\n"
48+
"-----END RSA PUBLIC KEY-----");
49+
return std::make_shared<PublicRsaKeySharedMain>(std::move(keys));
50+
}();
51+
return main_public_rsa_key;
4452
}
45-
46-
add_pem(
47-
"-----BEGIN RSA PUBLIC KEY-----\n"
48-
"MIIBCgKCAQEA6LszBcC1LGzyr992NzE0ieY+BSaOW622Aa9Bd4ZHLl+TuFQ4lo4g\n"
49-
"5nKaMBwK/BIb9xUfg0Q29/2mgIR6Zr9krM7HjuIcCzFvDtr+L0GQjae9H0pRB2OO\n"
50-
"62cECs5HKhT5DZ98K33vmWiLowc621dQuwKWSQKjWf50XYFw42h21P2KXUGyp2y/\n"
51-
"+aEyZ+uVgLLQbRA1dEjSDZ2iGRy12Mk5gpYc397aYp438fsJoHIgJ2lgMv5h7WY9\n"
52-
"t6N/byY9Nw9p21Og3AoXSL2q/2IJ1WRUhebgAdGVMlV1fkuOQoEzR7EdpqtQD9Cs\n"
53-
"5+bfo3Nhmcyvk5ftB0WkJ9z6bNZ7yxrP8wIDAQAB\n"
54-
"-----END RSA PUBLIC KEY-----");
5553
}
5654

5755
Result<mtproto::PublicRsaKeyInterface::RsaKey> PublicRsaKeySharedMain::get_rsa_key(const vector<int64> &fingerprints) {
5856
for (auto fingerprint : fingerprints) {
59-
const auto *rsa_key = get_rsa_key_unsafe(fingerprint);
60-
if (rsa_key != nullptr) {
61-
return RsaKey{rsa_key->rsa.clone(), fingerprint};
57+
for (const auto &key : keys_) {
58+
if (key.fingerprint == fingerprint) {
59+
return RsaKey{key.rsa.clone(), fingerprint};
60+
}
6261
}
6362
}
6463
return Status::Error(PSLICE() << "Unknown fingerprints " << format::as_array(fingerprints));
@@ -68,13 +67,4 @@ void PublicRsaKeySharedMain::drop_keys() {
6867
// nothing to do
6968
}
7069

71-
const mtproto::PublicRsaKeyInterface::RsaKey *PublicRsaKeySharedMain::get_rsa_key_unsafe(int64 fingerprint) const {
72-
auto it = std::find_if(keys_.begin(), keys_.end(),
73-
[fingerprint](const auto &value) { return value.fingerprint == fingerprint; });
74-
if (it == keys_.end()) {
75-
return nullptr;
76-
}
77-
return &*it;
78-
}
79-
8070
} // namespace td

‎td/telegram/net/PublicRsaKeySharedMain.h

+4-3
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,17 @@ namespace td {
1515

1616
class PublicRsaKeySharedMain final : public mtproto::PublicRsaKeyInterface {
1717
public:
18-
explicit PublicRsaKeySharedMain(bool is_test);
18+
explicit PublicRsaKeySharedMain(vector<RsaKey> &&keys) : keys_(std::move(keys)) {
19+
}
20+
21+
static std::shared_ptr<PublicRsaKeySharedMain> create(bool is_test);
1922

2023
Result<RsaKey> get_rsa_key(const vector<int64> &fingerprints) final;
2124

2225
void drop_keys() final;
2326

2427
private:
2528
vector<RsaKey> keys_;
26-
27-
const RsaKey *get_rsa_key_unsafe(int64 fingerprint) const;
2829
};
2930

3031
} // namespace td

‎test/mtproto.cpp

+4-2
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@
4646
#include "td/utils/tests.h"
4747
#include "td/utils/Time.h"
4848

49+
#include <memory>
50+
4951
TEST(Mtproto, GetHostByNameActor) {
5052
int threads_n = 1;
5153
td::ConcurrentScheduler sched(threads_n, 0);
@@ -300,11 +302,11 @@ class HandshakeContext final : public td::mtproto::AuthKeyHandshakeContext {
300302
return nullptr;
301303
}
302304
td::mtproto::PublicRsaKeyInterface *get_public_rsa_key_interface() final {
303-
return &public_rsa_key;
305+
return public_rsa_key_.get();
304306
}
305307

306308
private:
307-
td::PublicRsaKeySharedMain public_rsa_key{true};
309+
std::shared_ptr<td::mtproto::PublicRsaKeyInterface> public_rsa_key_ = td::PublicRsaKeySharedMain::create(true);
308310
};
309311

310312
class HandshakeTestActor final : public td::Actor {

0 commit comments

Comments
 (0)