51652: fix running of TRAPEXIT explicitly.

This is a special case where TRAPEXIT is unset within a TRAPEXIT
as it should never run in a nested context, so just save the
function structure temporarily on the heap.

3 files changed, 49 insertions, 1 deletions

diff --git a/ChangeLog b/ChangeLog
index 904d207de..981fdacd3 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,8 @@
+2023-04-17  Peter Stephenson  <p.stephenson@samsung.com>
+
+	* 51652 (plus typo correction): Src/exec.c, Test/C03traps.ztst:
+	fix running of TRAPEXIT explicitly.
+
 2023-04-11  Jun-ichi Takimoto  <takimoto-j@kba.biglobe.ne.jp>
 
 	* 51639: Doc/Zsh/params.yo, Src/init.c, configure.ac: add new
diff --git a/Src/exec.c b/Src/exec.c
index 3b3d1235e..8f9d5a885 100644
--- a/Src/exec.c
+++ b/Src/exec.c
@@ -5779,12 +5779,25 @@ doshfunc(Shfunc shfunc, LinkList doshargs, int noreturnval)
     char *name = shfunc->node.nam;
     int flags = shfunc->node.flags;
     char *fname = dupstring(name);
-    Eprog prog;
+    Eprog prog, marked_prog;
     static int oflags;
     static int funcdepth;
     Heap funcheap;
 
     queue_signals();	/* Lots of memory and global state changes coming */
+    /*
+     * In case this is a special function such as a trap, mark it
+     * as in use right now, so it doesn't get freed early.  The
+     * worst that can happen is this hangs around in memory a little
+     * longer than strictly needed.
+     *
+     * Classic example of this happening is running TRAPEXIT directly.
+     *
+     * Because the shell function's contents may change, we'll ensure
+     * we use a consistent structure for use / free.
+     */
+    marked_prog = shfunc->funcdef;
+    useeprog(marked_prog);
 
     NEWHEAPS(funcheap) {
 	/*
@@ -5818,6 +5831,22 @@ doshfunc(Shfunc shfunc, LinkList doshargs, int noreturnval)
 	    memcpy(funcsave->pipestats, pipestats, bytes);
 	}
 
+	if (!strcmp(fname, "TRAPEXIT")) {
+	    /*
+	     * If we are executing TRAPEXIT directly, starttrapscope()
+	     * will pull the rug out from under us to ensure the
+	     * exit trap isn't run inside the function.  We just need
+	     * the information locally here, so copy it on the heap.
+	     *
+	     * The funcdef is separately handled by reference counting.
+	     */
+	    Shfunc shcopy = (Shfunc)zhalloc(sizeof(struct shfunc));
+	    memcpy(shcopy, shfunc, sizeof(struct shfunc));
+	    shcopy->node.nam = dupstring(shfunc->node.nam);
+	    shfunc = shcopy;
+	    name = shfunc->node.nam;
+	}
+
 	starttrapscope();
 	startpatternscope();
 
@@ -5942,6 +5971,8 @@ doshfunc(Shfunc shfunc, LinkList doshargs, int noreturnval)
 	funcsave->fstack.filename = getshfuncfile(shfunc);
 
 	prog = shfunc->funcdef;
+	DPUTS1(!prog->nref, "function definition %s has zero reference count",
+	       (fname && *fname) ? fname : "<anon>");
 	if (prog->flags & EF_RUN) {
 	    Shfunc shf;
 
@@ -6046,6 +6077,7 @@ doshfunc(Shfunc shfunc, LinkList doshargs, int noreturnval)
 	}
     } OLDHEAPS;
 
+    freeeprog(marked_prog);
     unqueue_signals();
 
     /*
diff --git a/Test/C03traps.ztst b/Test/C03traps.ztst
index e0b6afb5f..de57765a0 100644
--- a/Test/C03traps.ztst
+++ b/Test/C03traps.ztst
@@ -1083,6 +1083,17 @@ F:Must be tested with a top-level script rather than source or function
 >trap1
 # As of 5.7.1-test-2, the output was "out1 fn1 trap1 fn2" (on separate lines).
 
+  TRAPEXIT() { echo This is TRAPEXIT; }
+  TRAPEXIT
+  TRAPEXIT
+  TRAPEXIT
+0:No memory problems with explicit call to TRAPEXIT.
+>This is TRAPEXIT
+>This is TRAPEXIT
+>This is TRAPEXIT
+>This is TRAPEXIT
+# Three explicit calls, one implicit call at function exit.
+
 %clean
 
   rm -f TRAPEXIT