fix webhook messages not using right username + avatar (fixes #209)

8 files changed, 118 insertions, 6 deletions

diff --git a/src/discord/discord.cpp b/src/discord/discord.cpp
index 084e19f..ccc61b4 100644
--- a/src/discord/discord.cpp
+++ b/src/discord/discord.cpp
@@ -333,6 +333,10 @@ std::vector<Snowflake> DiscordClient::GetChildChannelIDs(Snowflake parent_id) co
     return m_store.GetChannelIDsWithParentID(parent_id);
 }
 
+std::optional<WebhookMessageData> DiscordClient::GetWebhookMessageData(Snowflake message_id) const {
+    return m_store.GetWebhookMessage(message_id);
+}
+
 bool DiscordClient::IsThreadJoined(Snowflake thread_id) const {
     return std::find(m_joined_threads.begin(), m_joined_threads.end(), thread_id) != m_joined_threads.end();
 }
@@ -2735,15 +2739,18 @@ void DiscordClient::StoreMessageData(Message &msg) {
         for (const auto &r : *msg.Reactions) {
             if (!r.Emoji.ID.IsValid()) continue;
             const auto cur = m_store.GetEmoji(r.Emoji.ID);
-            if (!cur.has_value())
+            if (!cur.has_value()) {
                 m_store.SetEmoji(r.Emoji.ID, r.Emoji);
+            }
         }
 
-    for (const auto &user : msg.Mentions)
+    for (const auto &user : msg.Mentions) {
         m_store.SetUser(user.ID, user);
+    }
 
-    if (msg.Member.has_value())
+    if (msg.Member.has_value()) {
         m_store.SetGuildMember(*msg.GuildID, msg.Author.ID, *msg.Member);
+    }
 
     if (msg.Interaction.has_value()) {
         m_store.SetUser(msg.Interaction->User.ID, msg.Interaction->User);
@@ -2752,11 +2759,17 @@ void DiscordClient::StoreMessageData(Message &msg) {
         }
     }
 
+    if (msg.IsWebhook()) {
+        m_store.SetWebhookMessage(msg);
+    }
+
     m_store.EndTransaction();
 
-    if (msg.ReferencedMessage.has_value() && msg.MessageReference.has_value() && msg.MessageReference->ChannelID.has_value())
-        if (msg.ReferencedMessage.value() != nullptr)
+    if (msg.ReferencedMessage.has_value() && msg.MessageReference.has_value() && msg.MessageReference->ChannelID.has_value()) {
+        if (msg.ReferencedMessage.value() != nullptr) {
             StoreMessageData(**msg.ReferencedMessage);
+        }
+    }
 }
 
 // some notes for myself
diff --git a/src/discord/discord.hpp b/src/discord/discord.hpp
index eba190a..ebbf5f9 100644
--- a/src/discord/discord.hpp
+++ b/src/discord/discord.hpp
@@ -63,6 +63,7 @@ public:
     void GetArchivedPublicThreads(Snowflake channel_id, const sigc::slot<void(DiscordError, const ArchivedThreadsResponseData &)> &callback);
     void GetArchivedPrivateThreads(Snowflake channel_id, const sigc::slot<void(DiscordError, const ArchivedThreadsResponseData &)> &callback);
     std::vector<Snowflake> GetChildChannelIDs(Snowflake parent_id) const;
+    std::optional<WebhookMessageData> GetWebhookMessageData(Snowflake message_id) const;
 
     // get ids of given list of members for who we do not have the member data
     template<typename Iter>
diff --git a/src/discord/message.cpp b/src/discord/message.cpp
index 9adc7e6..bc4c6c8 100644
--- a/src/discord/message.cpp
+++ b/src/discord/message.cpp
@@ -283,3 +283,11 @@ bool Message::DoesMention(Snowflake id) const noexcept {
     if (!member.has_value()) return false;
     return std::find_first_of(MentionRoles.begin(), MentionRoles.end(), member->Roles.begin(), member->Roles.end()) != MentionRoles.end();
 }
+
+bool Message::IsWebhook() const noexcept {
+    return WebhookID.has_value();
+}
+
+std::optional<WebhookMessageData> Message::GetWebhookData() const {
+    return Abaddon::Get().GetDiscordClient().GetWebhookMessageData(ID);
+}
diff --git a/src/discord/message.hpp b/src/discord/message.hpp
index fb11604..0f53021 100644
--- a/src/discord/message.hpp
+++ b/src/discord/message.hpp
@@ -8,6 +8,7 @@
 #include "emoji.hpp"
 #include "member.hpp"
 #include "interactions.hpp"
+#include "webhook.hpp"
 #include "misc/bitwise.hpp"
 
 enum class MessageType {
@@ -226,6 +227,9 @@ struct Message {
 
     bool DoesMentionEveryoneOrUser(Snowflake id) const noexcept;
     bool DoesMention(Snowflake id) const noexcept;
+    bool IsWebhook() const noexcept;
+
+    std::optional<WebhookMessageData> GetWebhookData() const;
 
 private:
     bool m_deleted = false;
diff --git a/src/discord/store.cpp b/src/discord/store.cpp
index 7ee4d87..d8994c4 100644
--- a/src/discord/store.cpp
+++ b/src/discord/store.cpp
@@ -57,6 +57,20 @@ void Store::SetBan(Snowflake guild_id, Snowflake user_id, const BanData &ban) {
     s->Reset();
 }
 
+void Store::SetWebhookMessage(const Message &message) {
+    auto &s = m_stmt_set_webhook_msg;
+
+    s->Bind(1, message.ID);
+    s->Bind(2, message.WebhookID);
+    s->Bind(3, message.Author.Username);
+    s->Bind(4, message.Author.Avatar);
+
+    if (!s->Insert())
+        fprintf(stderr, "webhook message insert failed for %" PRIu64 ": %s\n", static_cast<uint64_t>(message.ID), m_db.ErrStr());
+
+    s->Reset();
+}
+
 void Store::SetChannel(Snowflake id, const ChannelData &chan) {
     auto &s = m_stmt_set_chan;
 
@@ -483,6 +497,29 @@ std::vector<BanData> Store::GetBans(Snowflake guild_id) const {
     return ret;
 }
 
+std::optional<WebhookMessageData> Store::GetWebhookMessage(Snowflake message_id) const {
+    auto &s = m_stmt_get_webhook_msg;
+
+    s->Bind(1, message_id);
+    if (!s->FetchOne()) {
+        if (m_db.Error() != SQLITE_DONE)
+            fprintf(stderr, "error while fetching webhook message %" PRIu64 ": %s\n", static_cast<uint64_t>(message_id), m_db.ErrStr());
+        s->Reset();
+        return {};
+    }
+
+    WebhookMessageData data;
+    s->Get(0, data.MessageID);
+    s->Get(1, data.WebhookID);
+    s->Get(2, data.Username);
+    s->Get(3, data.Avatar);
+
+    s->Reset();
+
+    return data;
+}
+
+
 Snowflake Store::GetGuildOwner(Snowflake guild_id) const {
     auto &s = m_stmt_get_guild_owner;
 
@@ -1503,6 +1540,15 @@ bool Store::CreateTables() {
         )
     )";
 
+    const char *create_webhook_messages = R"(
+        CREATE TABLE IF NOT EXISTS webhook_messages (
+            message_id INTEGER NOT NULL,
+            webhook_id INTEGER NOT NULL,
+            username TEXT,
+            avatar TEXT
+        )
+    )";
+
     if (m_db.Execute(create_users) != SQLITE_OK) {
         fprintf(stderr, "failed to create user table: %s\n", m_db.ErrStr());
         return false;
@@ -1608,6 +1654,11 @@ bool Store::CreateTables() {
         return false;
     }
 
+    if (m_db.Execute(create_webhook_messages) != SQLITE_OK) {
+        fprintf(stderr, "failed to create webhook messages table: %s\n", m_db.ErrStr());
+        return false;
+    }
+
     if (m_db.Execute(R"(
         CREATE TRIGGER remove_zero_reactions AFTER UPDATE ON reactions WHEN new.count = 0
         BEGIN
@@ -2288,6 +2339,24 @@ bool Store::CreateStatements() {
         return false;
     }
 
+    m_stmt_set_webhook_msg = std::make_unique<Statement>(m_db, R"(
+        REPLACE INTO webhook_messages VALUES (
+            ?, ?, ?, ?
+        )
+    )");
+    if (!m_stmt_set_webhook_msg->OK()) {
+        fprintf(stderr, "failed to prepare set webhook message statement: %s\n", m_db.ErrStr());
+        return false;
+    }
+
+    m_stmt_get_webhook_msg = std::make_unique<Statement>(m_db, R"(
+        SELECT * FROM webhook_messages WHERE message_id = ?
+    )");
+    if (!m_stmt_get_webhook_msg->OK()) {
+        fprintf(stderr, "failed to prepare get webhook message statement: %s\n", m_db.ErrStr());
+        return false;
+    }
+
     return true;
 }
 
diff --git a/src/discord/store.hpp b/src/discord/store.hpp
index 875a5af..b6979d0 100644
--- a/src/discord/store.hpp
+++ b/src/discord/store.hpp
@@ -7,7 +7,7 @@
 #include <sqlite3.h>
 
 #ifdef GetMessage // fuck you windows.h
-    #undef GetMessage
+#undef GetMessage
 #endif
 
 class Store {
@@ -26,6 +26,7 @@ public:
     void SetPermissionOverwrite(Snowflake channel_id, Snowflake id, const PermissionOverwrite &perm);
     void SetEmoji(Snowflake id, const EmojiData &emoji);
     void SetBan(Snowflake guild_id, Snowflake user_id, const BanData &ban);
+    void SetWebhookMessage(const Message &message);
 
     std::optional<ChannelData> GetChannel(Snowflake id) const;
     std::optional<EmojiData> GetEmoji(Snowflake id) const;
@@ -37,6 +38,7 @@ public:
     std::optional<UserData> GetUser(Snowflake id) const;
     std::optional<BanData> GetBan(Snowflake guild_id, Snowflake user_id) const;
     std::vector<BanData> GetBans(Snowflake guild_id) const;
+    std::optional<WebhookMessageData> GetWebhookMessage(Snowflake message_id) const;
 
     Snowflake GetGuildOwner(Snowflake guild_id) const;
     std::vector<Snowflake> GetMemberRoles(Snowflake guild_id, Snowflake user_id) const;
@@ -313,5 +315,7 @@ private:
     STMT(get_guild_member_ids);
     STMT(clr_role);
     STMT(get_guild_owner);
+    STMT(set_webhook_msg);
+    STMT(get_webhook_msg);
 #undef STMT
 };
diff --git a/src/discord/webhook.cpp b/src/discord/webhook.cpp
index 4e8b422..318f8c8 100644
--- a/src/discord/webhook.cpp
+++ b/src/discord/webhook.cpp
@@ -11,3 +11,7 @@ void from_json(const nlohmann::json &j, WebhookData &m) {
     JS_O("token", m.Token);
     JS_N("application_id", m.ApplicationID);
 }
+
+std::string WebhookMessageData::GetAvatarURL() const {
+    return Avatar.empty() ? "" : "https://cdn.discordapp.com/avatars/" + std::to_string(WebhookID) + "/" + Avatar + ".png?size=32";
+}
diff --git a/src/discord/webhook.hpp b/src/discord/webhook.hpp
index f0214df..3b9eeb9 100644
--- a/src/discord/webhook.hpp
+++ b/src/discord/webhook.hpp
@@ -22,3 +22,12 @@ struct WebhookData {
 
     friend void from_json(const nlohmann::json &j, WebhookData &m);
 };
+
+struct WebhookMessageData {
+    Snowflake MessageID;
+    Snowflake WebhookID;
+    std::string Username;
+    std::string Avatar;
+
+    std::string GetAvatarURL() const;
+};