From beef392f5ab7b1b16ddfd43433ce8b7868b62a83 Mon Sep 17 00:00:00 2001
From: Gledis Zeneli <43916939+gledis69@users.noreply.github.com>
Date: Sat, 28 May 2022 00:22:00 +0300
Subject: [PATCH] Fix memory error with citus_add_node reported by valgrind
 test (#5967)

The error comes due to the datum jsonb in pg_dist_metadata_node.metadata being 0 in some scenarios. This is likely due to not copying the data when receiving a datum from a tuple and pg deciding to deallocate that memory when the table that the tuple was from is closed.
Also fix another place in the code that might have been susceptible to this issue.
I tested on both multi-vg and multi-1-vg and the test were successful.
---
 src/backend/distributed/commands/role.c           | 3 ++-
 src/backend/distributed/metadata/metadata_cache.c | 7 +++++++
 src/backend/distributed/planner/multi_explain.c   | 2 +-
 src/backend/distributed/test/shard_rebalancer.c   | 2 +-
 4 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/src/backend/distributed/commands/role.c b/src/backend/distributed/commands/role.c
index af0a3a856..294e511a9 100644
--- a/src/backend/distributed/commands/role.c
+++ b/src/backend/distributed/commands/role.c
@@ -337,6 +337,7 @@ ExtractEncryptedPassword(Oid roleOid)
 
 	Datum passwordDatum = heap_getattr(tuple, Anum_pg_authid_rolpassword,
 									   pgAuthIdDescription, &isNull);
+	char *passwordCstring = TextDatumGetCString(passwordDatum);
 
 	table_close(pgAuthId, AccessShareLock);
 	ReleaseSysCache(tuple);
@@ -346,7 +347,7 @@ ExtractEncryptedPassword(Oid roleOid)
 		return NULL;
 	}
 
-	return pstrdup(TextDatumGetCString(passwordDatum));
+	return pstrdup(passwordCstring);
 }
 
 
diff --git a/src/backend/distributed/metadata/metadata_cache.c b/src/backend/distributed/metadata/metadata_cache.c
index 26371b45a..acf883422 100644
--- a/src/backend/distributed/metadata/metadata_cache.c
+++ b/src/backend/distributed/metadata/metadata_cache.c
@@ -70,6 +70,7 @@
 #include "utils/datum.h"
 #include "utils/elog.h"
 #include "utils/hsearch.h"
+#include "utils/jsonb.h"
 #if PG_VERSION_NUM >= PG_VERSION_13
 #include "common/hashfn.h"
 #endif
@@ -4910,6 +4911,12 @@ DistNodeMetadata(void)
 							"could not find any entries in pg_dist_metadata")));
 	}
 
+	/*
+	 * Copy the jsonb result before closing the table
+	 * since that memory can be freed.
+	 */
+	metadata = JsonbPGetDatum(DatumGetJsonbPCopy(metadata));
+
 	systable_endscan(scanDescriptor);
 	table_close(pgDistNodeMetadata, AccessShareLock);
 
diff --git a/src/backend/distributed/planner/multi_explain.c b/src/backend/distributed/planner/multi_explain.c
index a807085af..c63d359ec 100644
--- a/src/backend/distributed/planner/multi_explain.c
+++ b/src/backend/distributed/planner/multi_explain.c
@@ -386,7 +386,7 @@ static void
 ExplainPropertyBytes(const char *qlabel, int64 bytes, ExplainState *es)
 {
 	Datum textDatum = DirectFunctionCall1(pg_size_pretty, Int64GetDatum(bytes));
-	ExplainPropertyText(qlabel, text_to_cstring(DatumGetTextP(textDatum)), es);
+	ExplainPropertyText(qlabel, TextDatumGetCString(textDatum), es);
 }
 
 
diff --git a/src/backend/distributed/test/shard_rebalancer.c b/src/backend/distributed/test/shard_rebalancer.c
index f3640f415..9836c132a 100644
--- a/src/backend/distributed/test/shard_rebalancer.c
+++ b/src/backend/distributed/test/shard_rebalancer.c
@@ -577,7 +577,7 @@ JsonFieldValueString(Datum jsonDocument, const char *key)
 		return NULL;
 	}
 
-	char *valueString = text_to_cstring(DatumGetTextP(valueTextDatum));
+	char *valueString = TextDatumGetCString(valueTextDatum);
 	return valueString;
 }