From 908dddc3cb44b7864c1a89388e9aa68fa8f6cd20 Mon Sep 17 00:00:00 2001
From: Amolith <amolith@secluded.site>
Date: Sun, 5 Jul 2026 02:55:32 +0000
Subject: [PATCH] fix: enforce receive timeout during MAM requests

Amp-Thread-ID: https://ampcode.com/threads/T-019f2fe4-f392-7688-8b2f-8d40eb9fa4b8
---
 src/mode/mam.c | 48 ++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/mode/mam.c b/src/mode/mam.c
index 4b022fe6c3268c704ef5ebccfe3bc8ddc2fb4b79..a914623be658edead15261a7b1acee1eb106ac59 100644
--- a/src/mode/mam.c
+++ b/src/mode/mam.c
@@ -76,6 +76,7 @@ typedef struct {
   int messages_total;
   gint64 deadline_us;
   char *queryid;
+  bool done;
 } mam_run_t;
 
 static void _mam_options_init(mam_options_t *options);
@@ -85,9 +86,10 @@ static void _mam_start_request(xmppc_t *xmppc, mam_options_t *options, bool pret
 static void _mam_send_request(mam_run_t *run);
 static int _mam_finish(xmpp_conn_t *const conn, xmpp_stanza_t *const stanza, void *const userdata);
 static int _mam_receive_poll(xmpp_conn_t *const conn, void *const userdata);
+static int _mam_receive_timeout(xmpp_conn_t *const conn, void *const userdata);
 static int _mam_display_message(xmpp_conn_t *const conn, xmpp_stanza_t *const stanza, void *const userdata);
 static int _mam_display_pretty_message(xmpp_conn_t *const conn, xmpp_stanza_t *const stanza, void *const userdata);
-static void _mam_disconnect_and_free(xmpp_conn_t *conn, mam_run_t *run);
+static void _mam_disconnect_and_free(xmpp_conn_t *conn, mam_run_t *run, bool delete_id_handler);
 static void _mam_add_field(xmpp_ctx_t *ctx, xmpp_stanza_t *x, const char *var, const char *value, const char *type);
 static void _mam_add_ids_field(xmpp_ctx_t *ctx, xmpp_stanza_t *x, const char *ids);
 static void _mam_add_text_child(xmpp_ctx_t *ctx, xmpp_stanza_t *parent, const char *name, const char *text);
@@ -270,7 +272,9 @@ static void _mam_start_request(xmppc_t *xmppc, mam_options_t *options, bool pret
   run->options = *options;
   run->receive = receive;
   if (receive) {
+    unsigned long timeout_ms = (unsigned long)run->options.timeout * 1000;
     run->deadline_us = g_get_monotonic_time() + ((gint64)run->options.timeout * G_USEC_PER_SEC);
+    xmpp_timed_handler_add(xmppc->conn, _mam_receive_timeout, timeout_ms, run);
   }
 
   if (pretty) {
@@ -347,9 +351,13 @@ static int _mam_finish(xmpp_conn_t *const conn, xmpp_stanza_t *const stanza, voi
   mam_run_t *run = (mam_run_t *)userdata;
   const char *type = xmpp_stanza_get_type(stanza);
 
+  if (run->done) {
+    return 0;
+  }
+
   if (type && strcmp(type, "error") == 0) {
     logError(run->xmppc, "MAM request failed\n");
-    _mam_disconnect_and_free(conn, run);
+    _mam_disconnect_and_free(conn, run, false);
     return 0;
   }
 
@@ -358,7 +366,7 @@ static int _mam_finish(xmpp_conn_t *const conn, xmpp_stanza_t *const stanza, voi
     if (run->receive && run->messages_total == 0 && g_get_monotonic_time() < run->deadline_us) {
       xmpp_timed_handler_add(conn, _mam_receive_poll, MAM_POLL_PERIOD_MS, run);
     } else {
-      _mam_disconnect_and_free(conn, run);
+      _mam_disconnect_and_free(conn, run, false);
     }
   }
   return 0;
@@ -367,8 +375,12 @@ static int _mam_finish(xmpp_conn_t *const conn, xmpp_stanza_t *const stanza, voi
 static int _mam_receive_poll(xmpp_conn_t *const conn, void *const userdata) {
   mam_run_t *run = (mam_run_t *)userdata;
 
+  if (run->done) {
+    return 0;
+  }
+
   if (g_get_monotonic_time() >= run->deadline_us) {
-    _mam_disconnect_and_free(conn, run);
+    _mam_disconnect_and_free(conn, run, false);
     return 0;
   }
 
@@ -376,9 +388,29 @@ static int _mam_receive_poll(xmpp_conn_t *const conn, void *const userdata) {
   return 0;
 }
 
-static void _mam_disconnect_and_free(xmpp_conn_t *conn, mam_run_t *run) {
+static int _mam_receive_timeout(xmpp_conn_t *const conn, void *const userdata) {
+  mam_run_t *run = (mam_run_t *)userdata;
+
+  if (!run->done) {
+    _mam_disconnect_and_free(conn, run, true);
+  }
+
+  return 0;
+}
+
+static void _mam_disconnect_and_free(xmpp_conn_t *conn, mam_run_t *run, bool delete_id_handler) {
+  if (run->done) {
+    return;
+  }
+
+  run->done = true;
   xmpp_handler_delete(conn, _mam_display_message);
   xmpp_handler_delete(conn, _mam_display_pretty_message);
+  xmpp_timed_handler_delete(conn, _mam_receive_poll);
+  xmpp_timed_handler_delete(conn, _mam_receive_timeout);
+  if (delete_id_handler && run->queryid) {
+    xmpp_id_handler_delete(conn, _mam_finish, run->queryid);
+  }
   _mam_options_free(&run->options);
   g_free(run->queryid);
   g_free(run);
@@ -461,7 +493,7 @@ static bool _mam_parse_max(xmppc_t *xmppc, mam_options_t *options, const char *v
 static bool _mam_parse_timeout(xmppc_t *xmppc, mam_options_t *options, const char *value) {
   gint64 seconds = 0;
 
-  if (!_mam_parse_duration_seconds(value, &seconds) || seconds < 0 || seconds > G_MAXINT) {
+  if (!_mam_parse_duration_seconds(value, &seconds) || seconds < 0 || seconds > G_MAXINT || (guint64)seconds > ((guint64)G_MAXULONG / 1000)) {
     logError(xmppc, "timeout must be a duration like 30s, 5m, or 1h\n");
     return false;
   }
@@ -502,6 +534,10 @@ static bool _mam_parse_duration_seconds(const char *value, gint64 *seconds) {
     return false;
   }
 
+  if (parsed > G_MAXINT64 / multiplier) {
+    return false;
+  }
+
   *seconds = parsed * multiplier;
   return true;
 }