290 lines
10 KiB
C++
290 lines
10 KiB
C++
//
|
|
// 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/telegram/net/AuthDataShared.h"
|
|
#include "td/telegram/net/NetQuery.h"
|
|
#include "td/telegram/net/TempAuthKeyWatchdog.h"
|
|
|
|
#include "td/mtproto/AuthData.h"
|
|
#include "td/mtproto/AuthKey.h"
|
|
#include "td/mtproto/ConnectionManager.h"
|
|
#include "td/mtproto/Handshake.h"
|
|
#include "td/mtproto/MessageId.h"
|
|
#include "td/mtproto/SessionConnection.h"
|
|
|
|
#include "td/actor/actor.h"
|
|
|
|
#include "td/utils/buffer.h"
|
|
#include "td/utils/CancellationToken.h"
|
|
#include "td/utils/common.h"
|
|
#include "td/utils/FlatHashMap.h"
|
|
#include "td/utils/FlatHashSet.h"
|
|
#include "td/utils/List.h"
|
|
#include "td/utils/Promise.h"
|
|
#include "td/utils/Status.h"
|
|
#include "td/utils/StringBuilder.h"
|
|
#include "td/utils/VectorQueue.h"
|
|
|
|
#include <array>
|
|
#include <deque>
|
|
#include <functional>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <utility>
|
|
|
|
namespace td {
|
|
|
|
namespace mtproto {
|
|
class RawConnection;
|
|
} // namespace mtproto
|
|
|
|
namespace detail {
|
|
class GenAuthKeyActor;
|
|
} // namespace detail
|
|
|
|
class Session final
|
|
: public NetQueryCallback
|
|
, private mtproto::SessionConnection::Callback {
|
|
public:
|
|
class Callback {
|
|
public:
|
|
Callback() = default;
|
|
Callback(const Callback &) = delete;
|
|
Callback &operator=(const Callback &) = delete;
|
|
virtual ~Callback() = default;
|
|
virtual void on_failed() = 0;
|
|
virtual void on_closed() = 0;
|
|
virtual void request_raw_connection(unique_ptr<mtproto::AuthData> auth_data,
|
|
Promise<unique_ptr<mtproto::RawConnection>>) = 0;
|
|
virtual void on_tmp_auth_key_updated(mtproto::AuthKey auth_key) = 0;
|
|
virtual void on_server_salt_updated(vector<mtproto::ServerSalt> server_salts) = 0;
|
|
virtual void on_update(BufferSlice &&update, uint64 auth_key_id) = 0;
|
|
virtual void on_result(NetQueryPtr net_query) = 0;
|
|
};
|
|
|
|
Session(unique_ptr<Callback> callback, std::shared_ptr<AuthDataShared> shared_auth_data, int32 raw_dc_id, int32 dc_id,
|
|
bool is_primary, bool is_main, bool use_pfs, bool persist_tmp_auth_key, bool is_cdn,
|
|
bool need_destroy_auth_key, const mtproto::AuthKey &tmp_auth_key,
|
|
const vector<mtproto::ServerSalt> &server_salts);
|
|
|
|
void send(NetQueryPtr &&query);
|
|
|
|
void close();
|
|
|
|
static bool is_high_loaded();
|
|
|
|
private:
|
|
struct Query final : private ListNode {
|
|
mtproto::MessageId container_message_id_;
|
|
NetQueryPtr net_query_;
|
|
|
|
bool is_acknowledged_ = false;
|
|
bool is_unknown_ = false;
|
|
|
|
const int8 connection_id_;
|
|
const double sent_at_;
|
|
|
|
Query(mtproto::MessageId message_id, NetQueryPtr &&net_query, int8 connection_id, double sent_at)
|
|
: container_message_id_(message_id)
|
|
, net_query_(std::move(net_query))
|
|
, connection_id_(connection_id)
|
|
, sent_at_(sent_at) {
|
|
}
|
|
|
|
ListNode *get_list_node() {
|
|
return static_cast<ListNode *>(this);
|
|
}
|
|
static Query *from_list_node(ListNode *list_node) {
|
|
return static_cast<Query *>(list_node);
|
|
}
|
|
};
|
|
|
|
// When connection is closed, mark all queries without ack as unknown.
|
|
// Ask state of all unknown queries when new connection is created.
|
|
//
|
|
// Just re-ask answer_message_id each time we get information about it.
|
|
// Though mtproto::Connection must ensure delivery of such query.
|
|
|
|
const int32 raw_dc_id_; // numerical datacenter ID, i.e. 2
|
|
const int32 dc_id_; // unique datacenter ID, i.e. -10002
|
|
const bool is_primary_; // true for primary Sessions to all DCs
|
|
const bool is_main_; // true only for the primary Session(s) to the main DC
|
|
const bool persist_tmp_auth_key_;
|
|
const bool is_cdn_;
|
|
const bool need_destroy_auth_key_;
|
|
bool was_on_network_ = false;
|
|
bool network_flag_ = false;
|
|
bool online_flag_ = false;
|
|
bool logging_out_flag_ = false;
|
|
bool connection_online_flag_ = false;
|
|
enum class Mode : int8 { Tcp, Http } mode_ = Mode::Tcp;
|
|
uint32 network_generation_ = 0;
|
|
uint64 being_binded_tmp_auth_key_id_ = 0;
|
|
uint64 being_checked_main_auth_key_id_ = 0;
|
|
uint64 last_bind_query_id_ = 0;
|
|
uint64 last_check_query_id_ = 0;
|
|
double last_activity_timestamp_ = 0;
|
|
double last_success_timestamp_ = 0; // time when auth_key and Session definitely was valid
|
|
double last_bind_success_timestamp_ = 0; // time when auth_key and Session definitely was valid and authorized
|
|
size_t dropped_size_ = 0;
|
|
|
|
FlatHashSet<mtproto::MessageId, mtproto::MessageIdHash> unknown_queries_;
|
|
vector<mtproto::MessageId> to_cancel_message_ids_;
|
|
|
|
// Do not invalidate iterators of these two containers!
|
|
// TODO: better data structures
|
|
struct PriorityQueue {
|
|
void push(NetQueryPtr query);
|
|
NetQueryPtr pop();
|
|
bool empty() const;
|
|
|
|
private:
|
|
std::map<int8, VectorQueue<NetQueryPtr>, std::greater<>> queries_;
|
|
};
|
|
PriorityQueue pending_queries_;
|
|
std::map<mtproto::MessageId, Query> sent_queries_;
|
|
std::deque<NetQueryPtr> pending_invoke_after_queries_;
|
|
ListNode sent_queries_list_;
|
|
|
|
struct ConnectionInfo {
|
|
int8 connection_id_ = 0;
|
|
Mode mode_ = Mode::Tcp;
|
|
enum class State : int8 { Empty, Connecting, Ready } state_ = State::Empty;
|
|
CancellationTokenSource cancellation_token_source_;
|
|
unique_ptr<mtproto::SessionConnection> connection_;
|
|
bool ask_info_ = false;
|
|
double wakeup_at_ = 0;
|
|
double created_at_ = 0;
|
|
};
|
|
|
|
ConnectionInfo *current_info_;
|
|
ConnectionInfo main_connection_;
|
|
ConnectionInfo long_poll_connection_;
|
|
mtproto::ConnectionManager::ConnectionToken connection_token_;
|
|
|
|
double cached_connection_timestamp_ = 0;
|
|
unique_ptr<mtproto::RawConnection> cached_connection_;
|
|
|
|
std::shared_ptr<Callback> callback_;
|
|
bool use_pfs_{false};
|
|
bool need_check_main_key_{false};
|
|
TempAuthKeyWatchdog::RegisteredAuthKey registered_temp_auth_key_;
|
|
std::shared_ptr<AuthDataShared> shared_auth_data_;
|
|
bool close_flag_ = false;
|
|
|
|
static constexpr double ACTIVITY_TIMEOUT = 60 * 5;
|
|
static constexpr size_t MAX_INFLIGHT_QUERIES = 1024;
|
|
|
|
struct ContainerInfo {
|
|
size_t ref_cnt;
|
|
vector<mtproto::MessageId> message_ids;
|
|
};
|
|
FlatHashMap<mtproto::MessageId, ContainerInfo, mtproto::MessageIdHash> sent_containers_;
|
|
|
|
friend class GenAuthKeyActor;
|
|
struct HandshakeInfo {
|
|
bool flag_ = false;
|
|
ActorOwn<detail::GenAuthKeyActor> actor_;
|
|
unique_ptr<mtproto::AuthKeyHandshake> handshake_;
|
|
};
|
|
enum HandshakeId : int32 { MainAuthKeyHandshake = 0, TmpAuthKeyHandshake = 1 };
|
|
std::array<HandshakeInfo, 2> handshake_info_;
|
|
|
|
double wakeup_at_;
|
|
|
|
// mtproto::AuthData should be the last field, because it's size is about 32 KB
|
|
mtproto::AuthData auth_data_;
|
|
|
|
void on_handshake_ready(Result<unique_ptr<mtproto::AuthKeyHandshake>> r_handshake);
|
|
void create_gen_auth_key_actor(HandshakeId handshake_id);
|
|
void auth_loop(double now);
|
|
|
|
// mtproto::Connection::Callback
|
|
void on_connected() final;
|
|
void on_closed(Status status) final;
|
|
|
|
Status on_pong(double ping_time, double pong_time) final;
|
|
|
|
void on_network(bool network_flag, uint32 network_generation);
|
|
void on_online(bool online_flag);
|
|
void on_logging_out(bool logging_out_flag);
|
|
|
|
void on_auth_key_updated();
|
|
void on_tmp_auth_key_updated();
|
|
|
|
void on_server_salt_updated() final;
|
|
void on_server_time_difference_updated(bool force) final;
|
|
|
|
void on_new_session_created(uint64 unique_id, mtproto::MessageId first_message_id) final;
|
|
void on_session_failed(Status status) final;
|
|
|
|
void on_container_sent(mtproto::MessageId container_message_id, vector<mtproto::MessageId> message_ids) final;
|
|
|
|
Status on_update(BufferSlice packet) final;
|
|
|
|
void on_message_ack(mtproto::MessageId message_id) final;
|
|
Status on_message_result_ok(mtproto::MessageId message_id, BufferSlice packet, size_t original_size) final;
|
|
void on_message_result_error(mtproto::MessageId message_id, int error_code, string message) final;
|
|
void on_message_failed(mtproto::MessageId message_id, Status status) final;
|
|
|
|
void on_message_info(mtproto::MessageId message_id, int32 state, mtproto::MessageId answer_message_id,
|
|
int32 answer_size, int32 source) final;
|
|
|
|
Status on_destroy_auth_key() final;
|
|
|
|
void flush_pending_invoke_after_queries();
|
|
bool has_queries() const;
|
|
|
|
void dec_container(mtproto::MessageId container_message_id, Query *query);
|
|
void cleanup_container(mtproto::MessageId container_message_id, Query *query);
|
|
void mark_as_known(mtproto::MessageId message_id, Query *query);
|
|
void mark_as_unknown(mtproto::MessageId message_id, Query *query);
|
|
|
|
void on_message_ack_impl(mtproto::MessageId container_message_id, int32 type);
|
|
void on_message_ack_impl_inner(mtproto::MessageId message_id, int32 type, bool in_container);
|
|
void on_message_failed_inner(mtproto::MessageId message_id, bool in_container);
|
|
|
|
// send NetQueryPtr to parent
|
|
void return_query(NetQueryPtr &&query);
|
|
void add_query(NetQueryPtr &&net_query);
|
|
void resend_query(NetQueryPtr query);
|
|
|
|
void connection_open(ConnectionInfo *info, double now, bool ask_info = false);
|
|
void connection_add(unique_ptr<mtproto::RawConnection> raw_connection);
|
|
void connection_check_mode(ConnectionInfo *info);
|
|
void connection_open_finish(ConnectionInfo *info, Result<unique_ptr<mtproto::RawConnection>> r_raw_connection);
|
|
|
|
void connection_online_update(double now, bool force);
|
|
void connection_close(ConnectionInfo *info);
|
|
void connection_flush(ConnectionInfo *info);
|
|
void connection_send_query(ConnectionInfo *info, NetQueryPtr &&net_query, mtproto::MessageId message_id = {});
|
|
bool need_send_bind_key() const;
|
|
bool need_send_query() const;
|
|
bool can_destroy_auth_key() const;
|
|
bool connection_send_bind_key(ConnectionInfo *info);
|
|
bool need_send_check_main_key() const;
|
|
bool connection_send_check_main_key(ConnectionInfo *info);
|
|
|
|
void on_result(NetQueryPtr query) final;
|
|
|
|
void on_bind_result(NetQueryPtr query);
|
|
void on_check_key_result(NetQueryPtr query);
|
|
|
|
void start_up() final;
|
|
void timeout_expired() final;
|
|
void loop() final;
|
|
void hangup() final;
|
|
void raw_event(const Event::Raw &event) final;
|
|
|
|
friend StringBuilder &operator<<(StringBuilder &sb, Mode mode) {
|
|
return sb << (mode == Mode::Http ? "HTTP" : "TCP");
|
|
}
|
|
};
|
|
|
|
} // namespace td
|