handle opcode 9 invalid session and improve handling socket closure

3 files changed, 86 insertions, 3 deletions

diff --git a/discord/discord.cpp b/discord/discord.cpp
index 9d1dd04..382e3ac 100644
--- a/discord/discord.cpp
+++ b/discord/discord.cpp
@@ -7,6 +7,14 @@ DiscordClient::DiscordClient(bool mem_store)
     , m_decompress_buf(InflateChunkSize)
     , m_store(mem_store) {
     m_msg_dispatch.connect(sigc::mem_fun(*this, &DiscordClient::MessageDispatch));
+    auto dispatch_cb = [this]() {
+        m_generic_mutex.lock();
+        auto func = m_generic_queue.front();
+        m_generic_queue.pop();
+        m_generic_mutex.unlock();
+        func();
+    };
+    m_generic_dispatch.connect(dispatch_cb);
 
     m_websocket.signal_message().connect(sigc::mem_fun(*this, &DiscordClient::HandleGatewayMessageRaw));
     m_websocket.signal_open().connect(sigc::mem_fun(*this, &DiscordClient::HandleSocketOpen));
@@ -41,7 +49,7 @@ void DiscordClient::Stop() {
 
     m_websocket.Stop();
 
-    m_signal_disconnected.emit(false);
+    m_signal_disconnected.emit(false, GatewayCloseCode::UserDisconnect);
 }
 
 bool DiscordClient::IsStarted() const {
@@ -652,6 +660,9 @@ void DiscordClient::HandleGatewayMessage(std::string str) {
             case GatewayOp::Reconnect: {
                 HandleGatewayReconnect(m);
             } break;
+            case GatewayOp::InvalidSession: {
+                HandleGatewayInvalidSession(m);
+            } break;
             case GatewayOp::Event: {
                 auto iter = m_event_map.find(m.Type);
                 if (iter == m_event_map.end()) {
@@ -1100,7 +1111,8 @@ void DiscordClient::HandleGatewayInviteDelete(const GatewayMessage &msg) {
 }
 
 void DiscordClient::HandleGatewayReconnect(const GatewayMessage &msg) {
-    m_signal_disconnected.emit(true);
+    printf("received reconnect\n");
+    m_signal_disconnected.emit(true, GatewayCloseCode::Reconnecting);
     inflateEnd(&m_zstream);
     m_compressed_buf.clear();
 
@@ -1117,6 +1129,27 @@ void DiscordClient::HandleGatewayReconnect(const GatewayMessage &msg) {
     m_websocket.StartConnection(DiscordGateway);
 }
 
+void DiscordClient::HandleGatewayInvalidSession(const GatewayMessage &msg) {
+    printf("invalid session! re-identifying\n");
+
+    m_signal_disconnected.emit(true, GatewayCloseCode::Reconnecting);
+    inflateEnd(&m_zstream);
+    m_compressed_buf.clear();
+
+    std::memset(&m_zstream, 0, sizeof(m_zstream));
+    inflateInit2(&m_zstream, MAX_WBITS + 32);
+
+    m_heartbeat_acked = true;
+    m_wants_resume = false;
+
+    m_heartbeat_waiter.kill();
+    if (m_heartbeat_thread.joinable()) m_heartbeat_thread.join();
+
+    m_websocket.Stop(1000);
+
+    m_websocket.StartConnection(DiscordGateway);
+}
+
 void DiscordClient::HandleGatewayMessageUpdate(const GatewayMessage &msg) {
     Snowflake id = msg.Data.at("id");
 
@@ -1263,6 +1296,28 @@ void DiscordClient::HandleSocketOpen() {
 }
 
 void DiscordClient::HandleSocketClose(uint16_t code) {
+    printf("got socket close code: %d\n", code);
+    auto close_code = static_cast<GatewayCloseCode>(code);
+    auto cb = [this, close_code]() {
+        inflateEnd(&m_zstream);
+        m_compressed_buf.clear();
+
+        m_heartbeat_waiter.kill();
+        if (m_heartbeat_thread.joinable()) m_heartbeat_thread.join();
+        m_client_connected = false;
+
+        m_store.ClearAll();
+        m_chan_to_message_map.clear();
+        m_guild_to_users.clear();
+
+        m_websocket.Stop();
+
+        m_signal_disconnected.emit(false, close_code);
+    };
+    m_generic_mutex.lock();
+    m_generic_queue.push(cb);
+    m_generic_dispatch.emit();
+    m_generic_mutex.unlock();
 }
 
 bool DiscordClient::CheckCode(const http::response_type &r) {
diff --git a/discord/discord.hpp b/discord/discord.hpp
index 025684a..cef811e 100644
--- a/discord/discord.hpp
+++ b/discord/discord.hpp
@@ -170,6 +170,7 @@ private:
     void HandleGatewayInviteCreate(const GatewayMessage &msg);
     void HandleGatewayInviteDelete(const GatewayMessage &msg);
     void HandleGatewayReconnect(const GatewayMessage &msg);
+    void HandleGatewayInvalidSession(const GatewayMessage &msg);
     void HeartbeatThread();
     void SendIdentify();
     void SendResume();
@@ -217,6 +218,10 @@ private:
     std::queue<std::string> m_msg_queue;
     void MessageDispatch();
 
+    mutable std::mutex m_generic_mutex;
+    Glib::Dispatcher m_generic_dispatch;
+    std::queue<std::function<void()>> m_generic_queue;
+
     // signals
 public:
     typedef sigc::signal<void> type_signal_gateway_ready;
@@ -241,7 +246,7 @@ public:
     typedef sigc::signal<void, Snowflake, Snowflake> type_signal_guild_ban_add;       // guild id, user id
     typedef sigc::signal<void, InviteData> type_signal_invite_create;
     typedef sigc::signal<void, InviteDeleteObject> type_signal_invite_delete;
-    typedef sigc::signal<void, bool> type_signal_disconnected; // bool true if reconnecting
+    typedef sigc::signal<void, bool, GatewayCloseCode> type_signal_disconnected; // bool true if reconnecting
     typedef sigc::signal<void> type_signal_connected;
 
     type_signal_gateway_ready signal_gateway_ready();
diff --git a/discord/objects.hpp b/discord/objects.hpp
index 575ee85..e7e574c 100644
--- a/discord/objects.hpp
+++ b/discord/objects.hpp
@@ -28,6 +28,7 @@ enum class GatewayOp : int {
     UpdateStatus = 3,
     Resume = 6,
     Reconnect = 7,
+    InvalidSession = 9,
     Hello = 10,
     HeartbeatAck = 11,
     LazyLoadRequest = 14,
@@ -62,6 +63,28 @@ enum class GatewayEvent : int {
     INVITE_DELETE,
 };
 
+enum class GatewayCloseCode : uint16_t {
+    // discord
+    UnknownError = 4000,
+    UnknownOpcode = 4001,
+    DecodeError = 4002,
+    NotAuthenticated = 4003,
+    AuthenticationFailed = 4004,
+    AlreadyAuthenticated = 4005,
+    InvalidSequence = 4007,
+    RateLimited = 4008,
+    SessionTimedOut = 4009,
+    InvalidShard = 4010,
+    ShardingRequired = 4011,
+    InvalidAPIVersion = 4012,
+    InvalidIntents = 4013,
+    DisallowedIntents = 4014,
+
+    // internal
+    UserDisconnect = 4091,
+    Reconnecting = 4092,
+};
+
 struct GatewayMessage {
     GatewayOp Opcode;
     nlohmann::json Data;