From 05ab316008845925c4723194693ccf01baacae1d Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Mon, 27 Sep 2021 13:53:31 +0300
Subject: [PATCH 1/1] Fix snapshot reference leak if lo_export fails.

If lo_export() fails to open the target file or to write to it, it leaks
the created LargeObjectDesc and its snapshot in the top-transaction
context and resource owner. That's pretty harmless, it's a small leak
after all, but it gives the user a "Snapshot reference leak" warning.

Fix by using a short-lived memory context and no resource owner for
transient LargeObjectDescs that are opened and closed within one function
call. The leak is easiest to reproduce with lo_export() on a directory
that doesn't exist, but in principle the other lo_* functions could also
fail.

Backpatch to all supported versions.

Reported-by: Andrew B
Discussion: TODO
---
 src/backend/libpq/be-fsstubs.c             | 88 +++++++++-------------
 src/backend/storage/large_object/inv_api.c | 33 ++++----
 src/include/storage/large_object.h         |  5 +-
 src/test/regress/input/largeobject.source  | 11 +++
 src/test/regress/output/largeobject.source | 11 +++
 5 files changed, 81 insertions(+), 67 deletions(-)

diff --git a/src/backend/libpq/be-fsstubs.c b/src/backend/libpq/be-fsstubs.c
index 70b4111c14..680c1498e9 100644
--- a/src/backend/libpq/be-fsstubs.c
+++ b/src/backend/libpq/be-fsstubs.c
@@ -67,18 +67,10 @@
 static LargeObjectDesc **cookies = NULL;
 static int	cookies_size = 0;
 
+static bool lo_cleanup_needed = false;
 static MemoryContext fscxt = NULL;
 
-#define CreateFSContext() \
-	do { \
-		if (fscxt == NULL) \
-			fscxt = AllocSetContextCreate(TopMemoryContext, \
-										  "Filesystem", \
-										  ALLOCSET_DEFAULT_SIZES); \
-	} while (0)
-
-
-static int	newLOfd(LargeObjectDesc *lobjCookie);
+static int	newLOfd(void);
 static void deleteLOfd(int fd);
 static Oid	lo_import_internal(text *filename, Oid lobjOid);
 
@@ -99,11 +91,16 @@ be_lo_open(PG_FUNCTION_ARGS)
 	elog(DEBUG4, "lo_open(%u,%d)", lobjId, mode);
 #endif
 
-	CreateFSContext();
+	/*
+	 * Allocate a large object descriptor first.  This will also create
+	 * 'fscxt' if this is the first LO opened in this transaction.
+	 */
+	fd = newLOfd();
 
-	lobjDesc = inv_open(lobjId, mode, fscxt);
+	lobjDesc = inv_open(lobjId, mode, true, fscxt);
 
-	fd = newLOfd(lobjDesc);
+	Assert(cookies[fd] == NULL);
+	cookies[fd] = lobjDesc;
 
 	PG_RETURN_INT32(fd);
 }
@@ -238,12 +235,7 @@ be_lo_creat(PG_FUNCTION_ARGS)
 {
 	Oid			lobjId;
 
-	/*
-	 * We don't actually need to store into fscxt, but create it anyway to
-	 * ensure that AtEOXact_LargeObject knows there is state to clean up
-	 */
-	CreateFSContext();
-
+	lo_cleanup_needed = true;
 	lobjId = inv_create(InvalidOid);
 
 	PG_RETURN_OID(lobjId);
@@ -254,12 +246,7 @@ be_lo_create(PG_FUNCTION_ARGS)
 {
 	Oid			lobjId = PG_GETARG_OID(0);
 
-	/*
-	 * We don't actually need to store into fscxt, but create it anyway to
-	 * ensure that AtEOXact_LargeObject knows there is state to clean up
-	 */
-	CreateFSContext();
-
+	lo_cleanup_needed = true;
 	lobjId = inv_create(lobjId);
 
 	PG_RETURN_OID(lobjId);
@@ -339,7 +326,7 @@ be_lo_unlink(PG_FUNCTION_ARGS)
 
 	/*
 	 * inv_drop does not create a need for end-of-transaction cleanup and
-	 * hence we don't need to have created fscxt.
+	 * hence we don't need to set lo_cleanup_needed.
 	 */
 	PG_RETURN_INT32(inv_drop(lobjId));
 }
@@ -419,8 +406,6 @@ lo_import_internal(text *filename, Oid lobjOid)
 	LargeObjectDesc *lobj;
 	Oid			oid;
 
-	CreateFSContext();
-
 	/*
 	 * open the file to be read in
 	 */
@@ -435,12 +420,13 @@ lo_import_internal(text *filename, Oid lobjOid)
 	/*
 	 * create an inversion object
 	 */
+	lo_cleanup_needed = true;
 	oid = inv_create(lobjOid);
 
 	/*
 	 * read in from the filesystem and write to the inversion object
 	 */
-	lobj = inv_open(oid, INV_WRITE, fscxt);
+	lobj = inv_open(oid, INV_WRITE, false, CurrentMemoryContext);
 
 	while ((nbytes = read(fd, buf, BUFSIZE)) > 0)
 	{
@@ -482,12 +468,11 @@ be_lo_export(PG_FUNCTION_ARGS)
 	LargeObjectDesc *lobj;
 	mode_t		oumask;
 
-	CreateFSContext();
-
 	/*
 	 * open the inversion object (no need to test for failure)
 	 */
-	lobj = inv_open(lobjId, INV_READ, fscxt);
+	lo_cleanup_needed = true;
+	lobj = inv_open(lobjId, INV_READ, false, CurrentMemoryContext);
 
 	/*
 	 * open the file to be written to
@@ -592,7 +577,7 @@ AtEOXact_LargeObject(bool isCommit)
 {
 	int			i;
 
-	if (fscxt == NULL)
+	if (!lo_cleanup_needed)
 		return;					/* no LO operations in this xact */
 
 	/*
@@ -614,11 +599,14 @@ AtEOXact_LargeObject(bool isCommit)
 	cookies_size = 0;
 
 	/* Release the LO memory context to prevent permanent memory leaks. */
-	MemoryContextDelete(fscxt);
+	if (fscxt)
+		MemoryContextDelete(fscxt);
 	fscxt = NULL;
 
 	/* Give inv_api.c a chance to clean up, too */
 	close_lo_relation(isCommit);
+
+	lo_cleanup_needed = false;
 }
 
 /*
@@ -663,19 +651,22 @@ AtEOSubXact_LargeObject(bool isCommit, SubTransactionId mySubid,
  *****************************************************************************/
 
 static int
-newLOfd(LargeObjectDesc *lobjCookie)
+newLOfd(void)
 {
 	int			i,
 				newsize;
 
+	lo_cleanup_needed = true;
+	if (fscxt == NULL)
+		fscxt = AllocSetContextCreate(TopMemoryContext,
+									  "Filesystem",
+									  ALLOCSET_DEFAULT_SIZES);
+
 	/* Try to find a free slot */
 	for (i = 0; i < cookies_size; i++)
 	{
 		if (cookies[i] == NULL)
-		{
-			cookies[i] = lobjCookie;
 			return i;
-		}
 	}
 
 	/* No free slot, so make the array bigger */
@@ -700,8 +691,6 @@ newLOfd(LargeObjectDesc *lobjCookie)
 		cookies_size = newsize;
 	}
 
-	Assert(cookies[i] == NULL);
-	cookies[i] = lobjCookie;
 	return i;
 }
 
@@ -727,13 +716,8 @@ lo_get_fragment_internal(Oid loOid, int64 offset, int32 nbytes)
 	int			total_read PG_USED_FOR_ASSERTS_ONLY;
 	bytea	   *result = NULL;
 
-	/*
-	 * We don't actually need to store into fscxt, but create it anyway to
-	 * ensure that AtEOXact_LargeObject knows there is state to clean up
-	 */
-	CreateFSContext();
-
-	loDesc = inv_open(loOid, INV_READ, fscxt);
+	lo_cleanup_needed = true;
+	loDesc = inv_open(loOid, INV_READ, false, CurrentMemoryContext);
 
 	/*
 	 * Compute number of bytes we'll actually read, accommodating nbytes == -1
@@ -817,10 +801,9 @@ be_lo_from_bytea(PG_FUNCTION_ARGS)
 	LargeObjectDesc *loDesc;
 	int			written PG_USED_FOR_ASSERTS_ONLY;
 
-	CreateFSContext();
-
+	lo_cleanup_needed = true;
 	loOid = inv_create(loOid);
-	loDesc = inv_open(loOid, INV_WRITE, fscxt);
+	loDesc = inv_open(loOid, INV_WRITE, false, CurrentMemoryContext);
 	written = inv_write(loDesc, VARDATA_ANY(str), VARSIZE_ANY_EXHDR(str));
 	Assert(written == VARSIZE_ANY_EXHDR(str));
 	inv_close(loDesc);
@@ -840,9 +823,8 @@ be_lo_put(PG_FUNCTION_ARGS)
 	LargeObjectDesc *loDesc;
 	int			written PG_USED_FOR_ASSERTS_ONLY;
 
-	CreateFSContext();
-
-	loDesc = inv_open(loOid, INV_WRITE, fscxt);
+	lo_cleanup_needed = true;
+	loDesc = inv_open(loOid, INV_WRITE, false, CurrentMemoryContext);
 
 	/* Permission check */
 	if (!lo_compat_privileges &&
diff --git a/src/backend/storage/large_object/inv_api.c b/src/backend/storage/large_object/inv_api.c
index bee234bffc..d4f0e28cd8 100644
--- a/src/backend/storage/large_object/inv_api.c
+++ b/src/backend/storage/large_object/inv_api.c
@@ -244,13 +244,18 @@ inv_create(Oid lobjId)
 /*
  *	inv_open -- access an existing large object.
  *
- *		Returns:
- *		  Large object descriptor, appropriately filled in.  The descriptor
- *		  and subsidiary data are allocated in the specified memory context,
- *		  which must be suitably long-lived for the caller's purposes.
+ *	Arguments:
+ *	  lobjId	 - OID of the large object to open
+ *	  long_lived - if true, the descriptor is associated with the top
+ *				   transaction
+ *	  mcxt		 - memory context to allocate the descriptor and its
+ *				   subsidiary data in
+ *
+ *	Returns:
+ *	  Large object descriptor, appropriately filled in.
  */
 LargeObjectDesc *
-inv_open(Oid lobjId, int flags, MemoryContext mcxt)
+inv_open(Oid lobjId, int flags, bool long_lived, MemoryContext mcxt)
 {
 	LargeObjectDesc *retval;
 	Snapshot	snapshot = NULL;
@@ -314,17 +319,18 @@ inv_open(Oid lobjId, int flags, MemoryContext mcxt)
 	retval = (LargeObjectDesc *) MemoryContextAlloc(mcxt,
 													sizeof(LargeObjectDesc));
 	retval->id = lobjId;
-	retval->subid = GetCurrentSubTransactionId();
+	retval->long_lived = long_lived;
+	retval->subid = long_lived ? GetCurrentSubTransactionId() : InvalidSubTransactionId;
 	retval->offset = 0;
 	retval->flags = descflags;
 
 	/*
-	 * We must register the snapshot in TopTransaction's resowner, because it
-	 * must stay alive until the LO is closed rather than until the current
-	 * portal shuts down.  Do this last to avoid uselessly leaking the
-	 * snapshot if an error is thrown above.
+	 * If the descriptor is to be long-lived, we must register the snapshot in
+	 * TopTransaction's resowner so that it stays alive until the LO is closed
+	 * rather than until the current portal shuts down.  Do this last to avoid
+	 * uselessly leaking the snapshot if an error is thrown above.
 	 */
-	if (snapshot)
+	if (snapshot && long_lived)
 		snapshot = RegisterSnapshotOnOwner(snapshot,
 										   TopTransactionResourceOwner);
 	retval->snapshot = snapshot;
@@ -341,8 +347,9 @@ inv_close(LargeObjectDesc *obj_desc)
 {
 	Assert(PointerIsValid(obj_desc));
 
-	UnregisterSnapshotFromOwner(obj_desc->snapshot,
-								TopTransactionResourceOwner);
+	if (obj_desc->long_lived)
+		UnregisterSnapshotFromOwner(obj_desc->snapshot,
+									TopTransactionResourceOwner);
 
 	pfree(obj_desc);
 }
diff --git a/src/include/storage/large_object.h b/src/include/storage/large_object.h
index ae1e2482ea..a24d5a34fb 100644
--- a/src/include/storage/large_object.h
+++ b/src/include/storage/large_object.h
@@ -15,6 +15,7 @@
 #ifndef LARGE_OBJECT_H
 #define LARGE_OBJECT_H
 
+#include "utils/resowner.h"
 #include "utils/snapshot.h"
 
 
@@ -23,6 +24,7 @@
  *
  * id is the logical OID of the large object
  * snapshot is the snapshot to use for read/write operations
+ * long_lived indicates the descriptor must survive past current portal
  * subid is the subtransaction that opened the desc (or currently owns it)
  * offset is the current seek offset within the LO
  * flags contains some flag bits
@@ -40,6 +42,7 @@ typedef struct LargeObjectDesc
 {
 	Oid			id;				/* LO's identifier */
 	Snapshot	snapshot;		/* snapshot to use */
+	bool		long_lived;
 	SubTransactionId subid;		/* owning subtransaction ID */
 	uint64		offset;			/* current seek pointer */
 	int			flags;			/* see flag bits below */
@@ -88,7 +91,7 @@ extern bool lo_compat_privileges;
 /* inversion stuff in inv_api.c */
 extern void close_lo_relation(bool isCommit);
 extern Oid	inv_create(Oid lobjId);
-extern LargeObjectDesc *inv_open(Oid lobjId, int flags, MemoryContext mcxt);
+extern LargeObjectDesc *inv_open(Oid lobjId, int flags, bool long_lived, MemoryContext mcxt);
 extern void inv_close(LargeObjectDesc *obj_desc);
 extern int	inv_drop(Oid lobjId);
 extern int64 inv_seek(LargeObjectDesc *obj_desc, int64 offset, int whence);
diff --git a/src/test/regress/input/largeobject.source b/src/test/regress/input/largeobject.source
index ff42697d11..b1e7ae9909 100644
--- a/src/test/regress/input/largeobject.source
+++ b/src/test/regress/input/largeobject.source
@@ -124,6 +124,17 @@ BEGIN;
 SELECT lo_open(loid, x'40000'::int) from lotest_stash_values;
 ABORT;
 
+DO $$
+DECLARE
+  loid oid;
+BEGIN
+  SELECT tbl.loid INTO loid FROM lotest_stash_values tbl;
+  PERFORM lo_export(loid, '@abs_builddir@/results/invalid/path');
+EXCEPTION
+  WHEN UNDEFINED_FILE THEN RAISE NOTICE 'could not open file, as expected';
+END;
+$$;
+
 -- Test truncation.
 BEGIN;
 UPDATE lotest_stash_values SET fd=lo_open(loid, CAST(x'20000' | x'40000' AS integer));
diff --git a/src/test/regress/output/largeobject.source b/src/test/regress/output/largeobject.source
index 91090f0fde..91d33b4d0c 100644
--- a/src/test/regress/output/largeobject.source
+++ b/src/test/regress/output/largeobject.source
@@ -161,6 +161,17 @@ SELECT lo_open(loid, x'40000'::int) from lotest_stash_values;
 (1 row)
 
 ABORT;
+DO $$
+DECLARE
+  loid oid;
+BEGIN
+  SELECT tbl.loid INTO loid FROM lotest_stash_values tbl;
+  PERFORM lo_export(loid, '@abs_builddir@/results/invalid/path');
+EXCEPTION
+  WHEN UNDEFINED_FILE THEN RAISE NOTICE 'could not open file, as expected';
+END;
+$$;
+NOTICE:  could not open file, as expected
 -- Test truncation.
 BEGIN;
 UPDATE lotest_stash_values SET fd=lo_open(loid, CAST(x'20000' | x'40000' AS integer));
-- 
2.30.2

