Fix RI fast-path crash under nested C-level SPI

When a C-language function uses SPI_connect/SPI_execute/SPI_finish to
INSERT into a table with FK constraints, the FK AFTER triggers fire
and schedule ri_FastPathEndBatch via
RegisterAfterTriggerBatchCallback(), opening PK relations under
CurrentResourceOwner at the time of the SPI call.  The query_depth > 0
guard in FireAfterTriggerBatchCallbacks suppresses the callback at
that nesting level, deferring teardown to the outer query's
AfterTriggerEndQuery. By then the resource owner active during the SPI
call may have been released, decrementing the cached relations'
refcounts to zero. ri_FastPathTeardown, running under the outer
query's resource owner, then crashes in assert builds when it attempts
to close relations whose refcounts are already zero:

  TRAP: failed Assert("rel->rd_refcnt > 0")

Fix by storing batch callbacks at the level where they should fire:
in AfterTriggersQueryData.batch_callbacks for immediate constraints
(fired by AfterTriggerEndQuery) and in AfterTriggersData.batch_callbacks
for deferred constraints (fired by AfterTriggerFireDeferred and
AfterTriggerSetState).  RegisterAfterTriggerBatchCallback() routes the
callback to the current query-level list when query_depth >= 0, and to
the top-level list otherwise.  FireAfterTriggerBatchCallbacks() takes a
list parameter and simply iterates and invokes it; memory cleanup is
handled by the caller.  This replaces the query_depth > 0 guard with
list-level scoping.  Note that deferred constraints are unaffected by
this bug: their callbacks fire at commit via AfterTriggerFireDeferred,
under the outer transaction's resource owner, which remains valid
throughout.

Also add firing_batch_callbacks to AfterTriggersData to enforce that
callbacks do not register new callbacks during
FireAfterTriggerBatchCallbacks(), which would be unsafe as it could
modify the list being iterated.  An Assert in
RegisterAfterTriggerBatchCallback() enforces this discipline for
future callers.  The flag is reset at transaction and subtransaction
boundaries to handle cases where an error thrown by a callback is
caught and the subtransaction is rolled back.

While at it, ensure callbacks are properly accounted for at all
transaction boundaries, as cleanup of b7b27eb41a5c: discard any
remaining top-level callbacks on both commit and abort in
AfterTriggerEndXact(), and clean up query-level callbacks in
AfterTriggerFreeQuery().

Note that ri_PerformCheck() calls SPI with fire_triggers=false, which
skips AfterTriggerBeginQuery/EndQuery for that SPI command.  Any
triggers queued during that SPI command are not fired immediately but
deferred to the outer query level.  Since the fast-path check for
those triggers runs under the outer query's resource owner rather than
a nested SPI resource owner, and ri_PerformCheck() does not create
a dedicated child resource owner, the bug described above does not
apply.

Reported-by: Evan Montgomery-Recht <montge@mianetworks.net>
Reported-by: Sandro Santilli <strk@kbt.io>
Analyzed-by: Evan Montgomery-Recht <montge@mianetworks.net>
Author: Amit Langote <amitlangote09@gmail.com>
Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Discussion: https://postgr.es/m/CAEg7pwcKf01FmDqFAf-Hzu_pYnMYScY_Otid-pe9uw3BJ6gq9g@mail.gmail.com

diff --git a/src/backend/commands/trigger.c b/src/backend/commands/trigger.c
index c41005ba44e6294e7ebe182bb67e009b86cd796f..ba154178229a26a3070a4a4bb3142b88cefc26b5 100644 (file)
--- a/src/backend/commands/trigger.c
+++ b/src/backend/commands/trigger.c
@@ -3894,7 +3894,10 @@ typedef struct AfterTriggersData
        AfterTriggersTransData *trans_stack;    /* array of structs shown below */
        int                     maxtransdepth;  /* allocated len of above array */
 
-       List       *batch_callbacks;    /* List of AfterTriggerCallbackItem */
+       List       *batch_callbacks;    /* List of AfterTriggerCallbackItem; for
+                                                                        * deferred constraints */
+       bool            firing_batch_callbacks; /* true when in
+                                                                                * FireAfterTriggersBatchCallbacks() */
 } AfterTriggersData;
 
 struct AfterTriggersQueryData
@@ -3902,6 +3905,7 @@ struct AfterTriggersQueryData
        AfterTriggerEventList events;   /* events pending from this query */
        Tuplestorestate *fdw_tuplestore;        /* foreign tuples for said events */
        List       *tables;                     /* list of AfterTriggersTableData, see below */
+       List       *batch_callbacks;    /* List of AfterTriggerCallbackItem */
 };
 
 struct AfterTriggersTransData
@@ -3980,7 +3984,7 @@ static SetConstraintState SetConstraintStateAddItem(SetConstraintState state,
                                                                                                        Oid tgoid, bool tgisdeferred);
 static void cancel_prior_stmt_triggers(Oid relid, CmdType cmdType, int tgevent);
 
-static void FireAfterTriggerBatchCallbacks(void);
+static void FireAfterTriggerBatchCallbacks(List *callbacks);
 
 /*
  * Get the FDW tuplestore for the current trigger query level, creating it
@@ -5107,6 +5111,7 @@ AfterTriggerBeginXact(void)
        afterTriggers.firing_counter = (CommandId) 1;   /* mustn't be 0 */
        afterTriggers.query_depth = -1;
        afterTriggers.batch_callbacks = NIL;
+       afterTriggers.firing_batch_callbacks = false;
 
        /*
         * Verify that there is no leftover state remaining.  If these assertions
@@ -5233,11 +5238,9 @@ AfterTriggerEndQuery(EState *estate)
        /*
         * Fire batch callbacks before releasing query-level storage and before
         * decrementing query_depth.  Callbacks may do real work (index probes,
-        * error reporting) and rely on query_depth still reflecting the current
-        * batch level so that nested calls from SPI inside AFTER triggers are
-        * correctly suppressed by FireAfterTriggerBatchCallbacks's depth guard.
+        * error reporting).
         */
-       FireAfterTriggerBatchCallbacks();
+       FireAfterTriggerBatchCallbacks(qs->batch_callbacks);
 
        /* Release query-level-local storage, including tuplestores if any */
        AfterTriggerFreeQuery(&afterTriggers.query_stack[afterTriggers.query_depth]);
@@ -5300,6 +5303,9 @@ AfterTriggerFreeQuery(AfterTriggersQueryData *qs)
         */
        qs->tables = NIL;
        list_free_deep(tables);
+
+       list_free_deep(qs->batch_callbacks);
+       qs->batch_callbacks = NIL;
 }
 
 
@@ -5349,13 +5355,14 @@ AfterTriggerFireDeferred(void)
        }
 
        /* Flush any fast-path batches accumulated by the triggers just fired. */
-       FireAfterTriggerBatchCallbacks();
+       FireAfterTriggerBatchCallbacks(afterTriggers.batch_callbacks);
 
        afterTriggerFiringDepth--;
 
        /*
-        * We don't bother freeing the event list, since it will go away anyway
-        * (and more efficiently than via pfree) in AfterTriggerEndXact.
+        * We don't bother freeing the event list or batch_callbacks, since they
+        * will go away anyway (and more efficiently than via pfree) in
+        * AfterTriggerEndXact.
         */
 
        if (snap_pushed)
@@ -5419,6 +5426,10 @@ AfterTriggerEndXact(bool isCommit)
        afterTriggers.query_depth = -1;
 
        afterTriggerFiringDepth = 0;
+
+       list_free_deep(afterTriggers.batch_callbacks);
+       afterTriggers.batch_callbacks = NIL;
+       afterTriggers.firing_batch_callbacks = false;
 }
 
 /*
@@ -5565,6 +5576,9 @@ AfterTriggerEndSubXact(bool isCommit)
                        }
                }
        }
+
+       /* Reset in case a callback threw an error while firing. */
+       afterTriggers.firing_batch_callbacks = false;
 }
 
 /*
@@ -5719,6 +5733,7 @@ AfterTriggerEnlargeQueryState(void)
                qs->events.tailfree = NULL;
                qs->fdw_tuplestore = NULL;
                qs->tables = NIL;
+               qs->batch_callbacks = NIL;
 
                ++init_depth;
        }
@@ -6101,8 +6116,10 @@ AfterTriggerSetState(ConstraintsSetStmt *stmt)
                /*
                 * Flush any fast-path batches accumulated by the triggers just fired.
                 */
-               FireAfterTriggerBatchCallbacks();
+               FireAfterTriggerBatchCallbacks(afterTriggers.batch_callbacks);
                afterTriggerFiringDepth--;
+               list_free_deep(afterTriggers.batch_callbacks);
+               afterTriggers.batch_callbacks = NIL;
 
                if (snapshot_set)
                        PopActiveSnapshot();
@@ -6827,42 +6844,46 @@ RegisterAfterTriggerBatchCallback(AfterTriggerBatchCallback callback,
         * outside a trigger-firing context would never fire.
         */
        Assert(afterTriggerFiringDepth > 0);
+       Assert(!afterTriggers.firing_batch_callbacks);
        oldcxt = MemoryContextSwitchTo(TopTransactionContext);
        item = palloc(sizeof(AfterTriggerCallbackItem));
        item->callback = callback;
        item->arg = arg;
-       afterTriggers.batch_callbacks =
-               lappend(afterTriggers.batch_callbacks, item);
+       if (afterTriggers.query_depth >= 0)
+       {
+               AfterTriggersQueryData *qs =
+                       &afterTriggers.query_stack[afterTriggers.query_depth];
+
+               qs->batch_callbacks = lappend(qs->batch_callbacks, item);
+       }
+       else
+               afterTriggers.batch_callbacks =
+                       lappend(afterTriggers.batch_callbacks, item);
        MemoryContextSwitchTo(oldcxt);
 }
 
 /*
  * FireAfterTriggerBatchCallbacks
- *             Invoke and clear all registered batch callbacks.
+ *             Invoke all callbacks in the given list.
  *
- * Only fires at the outermost query level (query_depth == 0) or from
- * top-level operations (query_depth == -1, e.g. AfterTriggerFireDeferred
- * at COMMIT).  Nested queries from SPI inside AFTER triggers run at
- * depth > 0 and must not tear down resources the outer batch still needs.
+ * Memory cleanup of the list and its items is handled by the caller
+ * (AfterTriggerFreeQuery for query-level callbacks, AfterTriggerEndXact
+ * for top-level deferred callbacks).
  */
 static void
-FireAfterTriggerBatchCallbacks(void)
+FireAfterTriggerBatchCallbacks(List *callbacks)
 {
        ListCell   *lc;
 
-       if (afterTriggers.query_depth > 0)
-               return;
-
        Assert(afterTriggerFiringDepth > 0);
-       foreach(lc, afterTriggers.batch_callbacks)
+       afterTriggers.firing_batch_callbacks = true;
+       foreach(lc, callbacks)
        {
                AfterTriggerCallbackItem *item = lfirst(lc);
 
                item->callback(item->arg);
        }
-
-       list_free_deep(afterTriggers.batch_callbacks);
-       afterTriggers.batch_callbacks = NIL;
+       afterTriggers.firing_batch_callbacks = false;
 }
 
 /*