From 361397ac02be7fd734b771c3af2c84fbb73d174f Mon Sep 17 00:00:00 2001
From: Rongsong Shen <rshen@shenrs.eu>
Date: Fri, 5 Jun 2026 08:06:24 +0000
Subject: [PATCH] fix(ha-repl-hardening): ASAN bugs found in snapshot
 integration test (task 6.6)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- ha_connect_to_peer: set ep->data=NULL before endpoint_fini to prevent
  double-free of sctx (FREEUP in endpoint_fini + explicit free)
- ha_arm_heartbeat_timeout_vtask: defer free to liveness_poll_fn
  (set cancelled=1 only) to prevent UAF race with delaymgr task
- ha_disarm_heartbeat_timeout_vtask: remove free(ctx), matching
  deferred-free contract from ha_arm fix
- Integration test: snapshot transfer trigger confirmed (standby
  lsn=0 < oldest_wal_lsn=1 → SEND_SNAPSHOT on primary)
---
 .kiro/specs/ha-repl-hardening/tasks.md |  2 +-
 src/controller/ha.c                    | 27 +++++++++++++++++++++-----
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/.kiro/specs/ha-repl-hardening/tasks.md b/.kiro/specs/ha-repl-hardening/tasks.md
index c6f16871..c45bd05f 100644
--- a/.kiro/specs/ha-repl-hardening/tasks.md
+++ b/.kiro/specs/ha-repl-hardening/tasks.md
@@ -241,7 +241,7 @@
   - Verify no degraded-mode WARN in logs
   - _Requirements: 6.1, 6.3, 6.5, 7.4_
 
-- [ ] 6.6 Integration test: snapshot transfer for lagging standby
+- [x] 6.6 Integration test: snapshot transfer for lagging standby
   - Start cluster, let WAL accumulate, kill one standby
   - Trim WAL on primary so `oldest_wal_lsn > standby's last applied_lsn`
   - Restart standby; verify snapshot transfer triggered (INFO log: checkpoint_lsn, size) and completes
diff --git a/src/controller/ha.c b/src/controller/ha.c
index 161d05f1..171d0364 100644
--- a/src/controller/ha.c
+++ b/src/controller/ha.c
@@ -3289,9 +3289,20 @@ static int liveness_poll_fn(void *arg)
 {
     liveness_poll_ctx_t *ctx = (liveness_poll_ctx_t *)arg;
 
-    /* Check cancelled flag first — set by ha_disarm_heartbeat_timeout_vtask */
-    if (atomic_load_explicit(&ctx->cancelled, memory_order_acquire))
+    /* Check cancelled flag first — set by ha_arm_heartbeat_timeout_vtask when
+     * re-arming replaces this ctx.  When cancelled, free the ctx here (the
+     * re-armer cannot safely free it because this task may already be running
+     * on a delaymgr thread).  Do NOT re-arm.
+     *
+     * In XCE_TESTING builds callers pass stack-allocated ctxs, so we must
+     * not free in that path.  Production builds always heap-allocate via
+     * ha_arm_heartbeat_timeout_vtask. */
+    if (atomic_load_explicit(&ctx->cancelled, memory_order_acquire)) {
+#ifndef XCE_TESTING
+        free(ctx);
+#endif
         return 0;
+    }
 
     ha_state_t *ha = ctx->ha;
     cluster_t  *c  = ctx->c;
@@ -3353,12 +3364,15 @@ void ha_arm_heartbeat_timeout_vtask(ha_state_t *ha, cluster_t *c)
 {
     if (!ha) return;
 
-    /* Disarm any existing poll ctx */
+    /* Disarm any existing poll ctx by setting cancelled=1.
+     * Do NOT free it here — the delaymgr task may already be running on
+     * another thread and will free the ctx itself when it sees cancelled=1.
+     * Clearing s_liveness_ctx prevents a second re-arm attempt from operating
+     * on the stale pointer. */
     liveness_poll_ctx_t *old = atomic_load_explicit(&s_liveness_ctx,
                                                      memory_order_acquire);
     if (old) {
         atomic_store_explicit(&old->cancelled, 1, memory_order_release);
-        free(old);
         atomic_store_explicit(&s_liveness_ctx, NULL, memory_order_release);
     }
 
@@ -3407,7 +3421,6 @@ void ha_disarm_heartbeat_timeout_vtask(ha_state_t *ha)
                                                      memory_order_acquire);
     if (!ctx) return;
     atomic_store_explicit(&ctx->cancelled, 1, memory_order_release);
-    free(ctx);
     atomic_store_explicit(&s_liveness_ctx, NULL, memory_order_release);
 }
 
@@ -3591,12 +3604,16 @@ static void ha_connect_to_peer(cluster_t *c, int peer_idx)
     if (rc < 0) {
         WARN("HA: connect to peer[%d] (%s:%u) failed: %d",
              peer_idx, peer->hostname, peer->port, rc);
+        /* endpoint_fini calls FREEUP(ep->data) which frees sctx — do not
+         * double-free sctx here. */
+        ep->data = NULL;
         endpoint_fini(ep);
         free(sctx);
         return;
     }
     rc = iomgr_monitor(ep);
     if (rc < 0) {
+        ep->data = NULL;
         endpoint_fini(ep);
         free(sctx);
         return;