1 files changed, 38 insertions, 8 deletions

diff --git a/xmpp-vala/src/module/xep/0313_message_archive_management.vala b/xmpp-vala/src/module/xep/0313_message_archive_management.vala
index 1caa1bc3..2235e118 100644
--- a/xmpp-vala/src/module/xep/0313_message_archive_management.vala
+++ b/xmpp-vala/src/module/xep/0313_message_archive_management.vala
@@ -11,8 +11,8 @@ public class QueryResult {
     public bool error { get; set; default=false; }
     public bool malformed { get; set; default=false; }
     public bool complete { get; set; default=false; }
-    public string first { get; set; }
-    public string last { get; set; }
+    public string? first { get; set; }
+    public string? last { get; set; }
 }
 
 public class Module : XmppStreamModule {
@@ -65,16 +65,17 @@ public class Module : XmppStreamModule {
         }
 
         StanzaNode query_node = new StanzaNode.build("query", NS_VER(stream)).add_self_xmlns().put_node(data_form.get_submit_node());
-        if (queryid != null) {
-            query_node.put_attribute("queryid", queryid);
-        }
+        query_node.put_attribute("queryid", queryid);
         return query_node;
     }
 
     internal async QueryResult query_archive(XmppStream stream, string ns, Jid? mam_server, StanzaNode query_node, Cancellable? cancellable = null) {
-        var res = new QueryResult();
 
-        if (stream.get_flag(Flag.IDENTITY) == null) { res.error = true; return res; }
+        var res = new QueryResult();
+        Flag? flag = stream.get_flag(Flag.IDENTITY);
+        string? query_id = query_node.get_attribute("queryid");
+        if (flag == null || query_id == null) { res.error = true; return res; }
+        flag.active_query_ids.add(query_id);
 
         // Build and send query
         Iq.Stanza iq = new Iq.Stanza.set(query_node) { to=mam_server };
@@ -93,6 +94,11 @@ public class Module : XmppStreamModule {
         if ((res.first == null) != (res.last == null)) { res.malformed = true; return res; }
         res.complete = fin_node.get_attribute_bool("complete", false, ns);
 
+        Idle.add(() => {
+            flag.active_query_ids.remove(query_id);
+            return Source.REMOVE;
+        }, Priority.LOW);
+
         return res;
     }
 
@@ -104,7 +110,8 @@ public class ReceivedPipelineListener : StanzaListener<MessageStanza> {
     public override string[] after_actions { get { return after_actions_const; } }
 
     public override async bool run(XmppStream stream, MessageStanza message) {
-        if (stream.get_flag(Flag.IDENTITY) == null) return false;
+        Flag? flag = stream.get_flag(Flag.IDENTITY);
+        if (flag == null) return false;
 
         StanzaNode? message_node = message.stanza.get_deep_subnode(NS_VER(stream) + ":result", StanzaForwarding.NS_URI + ":forwarded", Xmpp.NS_URI + ":message");
         if (message_node != null) {
@@ -112,6 +119,28 @@ public class ReceivedPipelineListener : StanzaListener<MessageStanza> {
             DateTime? datetime = DelayedDelivery.get_time_for_node(forward_node);
             string? mam_id = message.stanza.get_deep_attribute(NS_VER(stream) + ":result", NS_VER(stream) + ":id");
             string? query_id = message.stanza.get_deep_attribute(NS_VER(stream) + ":result", NS_VER(stream) + ":queryid");
+
+            if (query_id == null) {
+                warning("Received MAM message without queryid from %s, ignoring", message.from.to_string());
+                return true;
+            }
+
+            if (!flag.active_query_ids.contains(query_id)) {
+                warning("Received MAM message from %s with unknown query id %s, ignoring", message.from.to_string(), query_id ?? "<none>");
+                return true;
+            }
+            Jid? inner_from = null;
+            try {
+                inner_from = new Jid(message_node.get_attribute("from"));
+            } catch (InvalidJidError e) {
+                warning("Received MAM message with invalid from attribute in forwarded message from %s, ignoring", message.from.to_string());
+                return true;
+            }
+            if (!message.from.equals(stream.get_flag(Bind.Flag.IDENTITY).my_jid.bare_jid) && !message.from.equals_bare(inner_from)) {
+                warning("Received MAM message from %s illegally impersonating %s, ignoring", message.from.to_string(), inner_from.to_string());
+                return true;
+            }
+
             message.add_flag(new MessageFlag(message.from, datetime, mam_id, query_id));
 
             message.stanza = message_node;
@@ -124,6 +153,7 @@ public class ReceivedPipelineListener : StanzaListener<MessageStanza> {
 public class Flag : XmppStreamFlag {
     public static FlagIdentity<Flag> IDENTITY = new FlagIdentity<Flag>(NS_URI, "message_archive_management");
     public bool cought_up { get; set; default=false; }
+    public Gee.Set<string> active_query_ids { get; set; default = new HashSet<string>(); }
     public string ns_ver;
 
     public Flag(string ns_ver) {