1 //
2 // Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2021
3 //
4 // Distributed under the Boost Software License, Version 1.0. (See accompanying
5 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
6 //
7 #include "td/telegram/net/DcAuthManager.h"
8
9 #include "td/telegram/Global.h"
10 #include "td/telegram/net/AuthDataShared.h"
11 #include "td/telegram/net/NetQuery.h"
12 #include "td/telegram/net/NetQueryDispatcher.h"
13 #include "td/telegram/TdDb.h"
14 #include "td/telegram/telegram_api.h"
15 #include "td/telegram/UniqueId.h"
16
17 #include "td/actor/actor.h"
18
19 #include "td/utils/logging.h"
20 #include "td/utils/misc.h"
21
22 #include <algorithm>
23 #include <limits>
24
25 namespace td {
26
27 int VERBOSITY_NAME(dc) = VERBOSITY_NAME(DEBUG) + 2;
28
DcAuthManager(ActorShared<> parent)29 DcAuthManager::DcAuthManager(ActorShared<> parent) {
30 parent_ = std::move(parent);
31 auto s_main_dc_id = G()->td_db()->get_binlog_pmc()->get("main_dc_id");
32 if (!s_main_dc_id.empty()) {
33 auto main_dc_id = to_integer<int32>(s_main_dc_id);
34 if (DcId::is_valid(main_dc_id)) {
35 main_dc_id_ = DcId::internal(main_dc_id);
36 VLOG(dc) << "Init main DcId to " << main_dc_id_;
37 } else {
38 LOG(ERROR) << "Receive invalid main DcId " << main_dc_id;
39 }
40 }
41 }
42
add_dc(std::shared_ptr<AuthDataShared> auth_data)43 void DcAuthManager::add_dc(std::shared_ptr<AuthDataShared> auth_data) {
44 VLOG(dc) << "Register " << auth_data->dc_id();
45 class Listener final : public AuthDataShared::Listener {
46 public:
47 explicit Listener(ActorShared<DcAuthManager> dc_manager) : dc_manager_(std::move(dc_manager)) {
48 }
49 bool notify() final {
50 if (!dc_manager_.is_alive()) {
51 return false;
52 }
53 send_closure(dc_manager_, &DcAuthManager::update_auth_key_state);
54 return true;
55 }
56
57 private:
58 ActorShared<DcAuthManager> dc_manager_;
59 };
60
61 DcInfo info;
62 info.dc_id = auth_data->dc_id();
63 CHECK(info.dc_id.is_exact());
64 info.shared_auth_data = std::move(auth_data);
65 info.auth_key_state = info.shared_auth_data->get_auth_key_state();
66 VLOG(dc) << "Add " << info.dc_id << " with auth key state " << info.auth_key_state;
67 if (!main_dc_id_.is_exact()) {
68 main_dc_id_ = info.dc_id;
69 VLOG(dc) << "Set main DcId to " << main_dc_id_;
70 }
71 info.shared_auth_data->add_auth_key_listener(make_unique<Listener>(actor_shared(this, info.dc_id.get_raw_id())));
72 dcs_.emplace_back(std::move(info));
73 loop();
74 }
75
update_main_dc(DcId new_main_dc_id)76 void DcAuthManager::update_main_dc(DcId new_main_dc_id) {
77 main_dc_id_ = new_main_dc_id;
78 VLOG(dc) << "Update main DcId to " << main_dc_id_;
79 loop();
80 }
81
get_dc(int32 dc_id)82 DcAuthManager::DcInfo &DcAuthManager::get_dc(int32 dc_id) {
83 auto *dc = find_dc(dc_id);
84 CHECK(dc);
85 return *dc;
86 }
find_dc(int32 dc_id)87 DcAuthManager::DcInfo *DcAuthManager::find_dc(int32 dc_id) {
88 auto it = std::find_if(dcs_.begin(), dcs_.end(), [&](auto &x) { return x.dc_id.get_raw_id() == dc_id; });
89 if (it == dcs_.end()) {
90 return nullptr;
91 }
92 return &*it;
93 }
94
update_auth_key_state()95 void DcAuthManager::update_auth_key_state() {
96 auto dc_id = narrow_cast<int32>(get_link_token());
97 auto &dc = get_dc(dc_id);
98 dc.auth_key_state = dc.shared_auth_data->get_auth_key_state();
99 VLOG(dc) << "Update " << dc_id << " auth key state from " << dc.auth_key_state << " to " << dc.auth_key_state;
100
101 loop();
102 }
103
on_result(NetQueryPtr result)104 void DcAuthManager::on_result(NetQueryPtr result) {
105 auto dc_id = narrow_cast<int32>(get_link_token());
106 auto &dc = get_dc(dc_id);
107 CHECK(dc.wait_id == result->id());
108 dc.wait_id = std::numeric_limits<decltype(dc.wait_id)>::max();
109 switch (dc.state) {
110 case DcInfo::State::Import: {
111 if (result->is_error()) {
112 LOG(WARNING) << "DC auth_exportAuthorization error: " << result->error();
113 dc.state = DcInfo::State::Export;
114 break;
115 }
116 auto r_result_auth_exported = fetch_result<telegram_api::auth_exportAuthorization>(result->ok());
117 if (r_result_auth_exported.is_error()) {
118 LOG(WARNING) << "Failed to parse result to auth_exportAuthorization: " << r_result_auth_exported.error();
119 dc.state = DcInfo::State::Export;
120 break;
121 }
122 auto result_auth_exported = r_result_auth_exported.move_as_ok();
123 dc.export_id = result_auth_exported->id_;
124 dc.export_bytes = std::move(result_auth_exported->bytes_);
125 break;
126 }
127 case DcInfo::State::BeforeOk: {
128 if (result->is_error()) {
129 LOG(WARNING) << "DC authImport error: " << result->error();
130 dc.state = DcInfo::State::Export;
131 break;
132 }
133 auto result_auth = fetch_result<telegram_api::auth_importAuthorization>(result->ok());
134 if (result_auth.is_error()) {
135 LOG(WARNING) << "Failed to parse result to auth_importAuthorization: " << result_auth.error();
136 dc.state = DcInfo::State::Export;
137 break;
138 }
139 dc.state = DcInfo::State::Ok;
140 break;
141 }
142 default:
143 UNREACHABLE();
144 }
145 result->clear();
146 loop();
147 }
148
dc_loop(DcInfo & dc)149 void DcAuthManager::dc_loop(DcInfo &dc) {
150 VLOG(dc) << "In dc_loop: " << dc.dc_id << " " << dc.auth_key_state;
151 if (dc.auth_key_state == AuthKeyState::OK) {
152 return;
153 }
154 if (dc.state == DcInfo::State::Ok) {
155 LOG(WARNING) << "Lost key in " << dc.dc_id << ", restart dc_loop";
156 dc.state = DcInfo::State::Waiting;
157 }
158 CHECK(dc.shared_auth_data);
159 switch (dc.state) {
160 case DcInfo::State::Waiting: {
161 // wait for timeout
162 // break;
163 }
164 case DcInfo::State::Export: {
165 // send auth.exportAuthorization to auth_dc
166 VLOG(dc) << "Send exportAuthorization to " << dc.dc_id;
167 auto id = UniqueId::next();
168 auto query = G()->net_query_creator().create(id, telegram_api::auth_exportAuthorization(dc.dc_id.get_raw_id()),
169 DcId::main(), NetQuery::Type::Common, NetQuery::AuthFlag::On);
170 query->total_timeout_limit_ = 60 * 60 * 24;
171 G()->net_query_dispatcher().dispatch_with_callback(std::move(query), actor_shared(this, dc.dc_id.get_raw_id()));
172 dc.wait_id = id;
173 dc.export_id = -1;
174 dc.state = DcInfo::State::Import;
175 break;
176 }
177 case DcInfo::State::Import: {
178 // send auth.importAuthorization to dc
179 if (dc.export_id == -1) {
180 break;
181 }
182 uint64 id = UniqueId::next();
183 VLOG(dc) << "Send importAuthorization to " << dc.dc_id;
184 auto query = G()->net_query_creator().create(
185 id, telegram_api::auth_importAuthorization(dc.export_id, std::move(dc.export_bytes)), dc.dc_id,
186 NetQuery::Type::Common, NetQuery::AuthFlag::Off);
187 query->total_timeout_limit_ = 60 * 60 * 24;
188 G()->net_query_dispatcher().dispatch_with_callback(std::move(query), actor_shared(this, dc.dc_id.get_raw_id()));
189 dc.wait_id = id;
190 dc.state = DcInfo::State::BeforeOk;
191 break;
192 }
193 case DcInfo::State::BeforeOk:
194 break;
195 case DcInfo::State::Ok:
196 break;
197 }
198 }
199
destroy(Promise<> promise)200 void DcAuthManager::destroy(Promise<> promise) {
201 destroy_promise_ = std::move(promise);
202 loop();
203 }
204
destroy_loop()205 void DcAuthManager::destroy_loop() {
206 if (!destroy_promise_) {
207 return;
208 }
209 bool is_ready{true};
210 for (auto &dc : dcs_) {
211 is_ready &= dc.auth_key_state == AuthKeyState::Empty;
212 }
213
214 if (is_ready) {
215 VLOG(dc) << "Destroy auth keys loop is ready, all keys are destroyed";
216 destroy_promise_.set_value(Unit());
217 } else {
218 VLOG(dc) << "DC is not ready for destroying auth key";
219 }
220 }
221
loop()222 void DcAuthManager::loop() {
223 if (close_flag_) {
224 VLOG(dc) << "Skip loop because close_flag";
225 return;
226 }
227 destroy_loop();
228 if (!main_dc_id_.is_exact()) {
229 VLOG(dc) << "Skip loop because main_dc_id is unknown";
230 return;
231 }
232 auto main_dc = find_dc(main_dc_id_.get_raw_id());
233 if (!main_dc || main_dc->auth_key_state != AuthKeyState::OK) {
234 VLOG(dc) << "Skip loop, because main DC is " << main_dc_id_ << ", main auth key state is "
235 << (main_dc != nullptr ? main_dc->auth_key_state : AuthKeyState::Empty);
236 return;
237 }
238 for (auto &dc : dcs_) {
239 dc_loop(dc);
240 }
241 }
242
243 } // namespace td
244