v1-0001-Fix-bug-plpython-function-causes-server-panic.patch

From 891ff26ba9b52f3eba6f0d112657c3b288524de2 Mon Sep 17 00:00:00 2001
From: Hao Zhang <zhrt1446384557@gmail.com>
Date: Thu, 30 Nov 2023 14:51:07 +0800
Subject: [PATCH v1] Fix bug: plpython function causes server panic.

---
 src/pl/plpython/plpy_cursorobject.c | 12 ++++++---
 src/pl/plpython/plpy_spi.c          | 41 +++++++++++++++++++++++------
 src/pl/plpython/plpy_spi.h          |  2 +-
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/pl/plpython/plpy_cursorobject.c b/src/pl/plpython/plpy_cursorobject.c
index 57e8f8ec21..8bfe05815b 100644
--- a/src/pl/plpython/plpy_cursorobject.c
+++ b/src/pl/plpython/plpy_cursorobject.c
@@ -98,7 +98,8 @@ PLy_cursor_query(const char *query)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	PLy_spi_subtransaction_begin(oldcontext, oldowner);
+	if (!PLy_spi_subtransaction_begin(oldcontext, oldowner))
+		return NULL;
 
 	PG_TRY();
 	{
@@ -196,7 +197,8 @@ PLy_cursor_plan(PyObject *ob, PyObject *args)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	PLy_spi_subtransaction_begin(oldcontext, oldowner);
+	if (!PLy_spi_subtransaction_begin(oldcontext, oldowner))
+		return NULL;
 
 	PG_TRY();
 	{
@@ -333,7 +335,8 @@ PLy_cursor_iternext(PyObject *self)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	PLy_spi_subtransaction_begin(oldcontext, oldowner);
+	if (!PLy_spi_subtransaction_begin(oldcontext, oldowner))
+		return NULL;
 
 	PG_TRY();
 	{
@@ -403,7 +406,8 @@ PLy_cursor_fetch(PyObject *self, PyObject *args)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	PLy_spi_subtransaction_begin(oldcontext, oldowner);
+	if (!PLy_spi_subtransaction_begin(oldcontext, oldowner))
+		return NULL;
 
 	PG_TRY();
 	{
diff --git a/src/pl/plpython/plpy_spi.c b/src/pl/plpython/plpy_spi.c
index ff87b27de0..58a274eda9 100644
--- a/src/pl/plpython/plpy_spi.c
+++ b/src/pl/plpython/plpy_spi.c
@@ -77,7 +77,8 @@ PLy_spi_prepare(PyObject *self, PyObject *args)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	PLy_spi_subtransaction_begin(oldcontext, oldowner);
+	if (!PLy_spi_subtransaction_begin(oldcontext, oldowner))
+		return NULL;
 
 	PG_TRY();
 	{
@@ -217,7 +218,8 @@ PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	PLy_spi_subtransaction_begin(oldcontext, oldowner);
+	if (!PLy_spi_subtransaction_begin(oldcontext, oldowner))
+		return NULL;
 
 	PG_TRY();
 	{
@@ -313,7 +315,8 @@ PLy_spi_execute_query(char *query, long limit)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	PLy_spi_subtransaction_begin(oldcontext, oldowner);
+	if (!PLy_spi_subtransaction_begin(oldcontext, oldowner))
+		return NULL;
 
 	PG_TRY();
 	{
@@ -556,7 +559,9 @@ PLy_rollback(PyObject *self, PyObject *args)
  *	MemoryContext oldcontext = CurrentMemoryContext;
  *	ResourceOwner oldowner = CurrentResourceOwner;
  *
- *	PLy_spi_subtransaction_begin(oldcontext, oldowner);
+ *	if (!PLy_spi_subtransaction_begin(oldcontext, oldowner))
+ *		return NULL;
+ *
  *	PG_TRY();
  *	{
  *		<call SPI functions>
@@ -573,12 +578,32 @@ PLy_rollback(PyObject *self, PyObject *args)
  * These utilities take care of restoring connection to the SPI manager and
  * setting a Python exception in case of an abort.
  */
-void
+bool
 PLy_spi_subtransaction_begin(MemoryContext oldcontext, ResourceOwner oldowner)
 {
-	BeginInternalSubTransaction(NULL);
-	/* Want to run inside function's memory context */
-	MemoryContextSwitchTo(oldcontext);
+	/*
+	 * Catch the error when begin a subtransaction.
+	 * Set the python error indicator if any error occurs to free
+	 * the traced pyobjects.
+	 */
+	PG_TRY();
+	{
+		BeginInternalSubTransaction(NULL);
+		/* Want to run inside function's memory context */
+		MemoryContextSwitchTo(oldcontext);
+	}
+	PG_CATCH();
+	{
+		MemoryContextSwitchTo(oldcontext);
+		CurrentResourceOwner = oldowner;
+
+		ErrorData *edata = CopyErrorData();
+		PLy_spi_exception_set(PyExc_Exception, edata);
+		FreeErrorData(edata);
+		return false;
+	}
+	PG_END_TRY();
+	return true;
 }
 
 void
diff --git a/src/pl/plpython/plpy_spi.h b/src/pl/plpython/plpy_spi.h
index 98ccd21093..a43d803928 100644
--- a/src/pl/plpython/plpy_spi.h
+++ b/src/pl/plpython/plpy_spi.h
@@ -22,7 +22,7 @@ typedef struct PLyExceptionEntry
 } PLyExceptionEntry;
 
 /* handling of SPI operations inside subtransactions */
-extern void PLy_spi_subtransaction_begin(MemoryContext oldcontext, ResourceOwner oldowner);
+extern bool PLy_spi_subtransaction_begin(MemoryContext oldcontext, ResourceOwner oldowner);
 extern void PLy_spi_subtransaction_commit(MemoryContext oldcontext, ResourceOwner oldowner);
 extern void PLy_spi_subtransaction_abort(MemoryContext oldcontext, ResourceOwner oldowner);
 
-- 
2.43.0