From 878672ff4456a6df6955a207811beb3fe60f226e Mon Sep 17 00:00:00 2001
From: Corey Huinker <corey.huinker@gmail.com>
Date: Thu, 23 Jan 2025 23:36:54 -0500
Subject: [PATCH v2 5/6] Add attnum bounds checking routines to pg_ndistinct

---
 .../statistics/extended_stats_internal.h      |  3 +
 src/backend/statistics/mvdistinct.c           | 62 +++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/src/include/statistics/extended_stats_internal.h b/src/include/statistics/extended_stats_internal.h
index df1a5308f4..d5d4cddad5 100644
--- a/src/include/statistics/extended_stats_internal.h
+++ b/src/include/statistics/extended_stats_internal.h
@@ -130,5 +130,8 @@ extern Datum import_mcvlist(HeapTuple tup, int elevel, int numattrs,
 							Oid *atttypids, int32 *atttypmods, Oid *atttypcolls,
 							int nitems, Datum *mcv_elems, bool *mcv_nulls,
 							bool *mcv_elem_nulls, float8 *freqs, float8 *base_freqs);
+extern bool pg_ndistinct_validate_items(MVNDistinct *ndistinct, int2vector *stxkeys,
+										int numexprs, int elevel);
+extern void free_pg_ndistinct(MVNDistinct *ndistinct);
 
 #endif							/* EXTENDED_STATS_INTERNAL_H */
diff --git a/src/backend/statistics/mvdistinct.c b/src/backend/statistics/mvdistinct.c
index e2a3acaf78..e9c02aaa63 100644
--- a/src/backend/statistics/mvdistinct.c
+++ b/src/backend/statistics/mvdistinct.c
@@ -630,6 +630,68 @@ pg_ndistinct_in(PG_FUNCTION_ARGS)
 	PG_RETURN_NULL();
 }
 
+/*
+ * Free allocations of an MVNDistinct
+ */
+void
+free_pg_ndistinct(MVNDistinct *ndistinct)
+{
+	for (int i = 0; i < ndistinct->nitems; i++)
+		pfree(ndistinct->items[i].attributes);
+
+	pfree(ndistinct);
+}
+
+/*
+ * Validate an MVNDistinct against the extended statistics object definition.
+ *
+ * Every MVNDistinctItem must be checked to ensure that the attnums in the
+ * attributes list correspond to attnums/expressions defined by the
+ * extended statistics object.
+ *
+ * Positive attnums are attributes which must be found in the stxkeys,
+ * while negative attnums correspond to an expr number, so the attnum 
+ * can't be below (0 - numexprs).
+ */
+bool
+pg_ndistinct_validate_items(MVNDistinct *ndistinct, int2vector *stxkeys, int numexprs, int elevel)
+{
+	int		attnum_expr_lowbound = 0 - numexprs;
+
+	for (int i = 0; i < ndistinct->nitems; i++)
+	{
+		MVNDistinctItem item = ndistinct->items[i];
+
+		for (int j = 0; j < item.nattributes; j++)
+		{
+			AttrNumber	attnum = item.attributes[j];
+			bool		ok = false;
+
+			if (attnum > 0)
+			{
+				for (int k = 0; k < stxkeys->dim1; k++)
+					if (attnum == stxkeys->values[k])
+					{
+						ok = true;
+						break;
+					}
+			}
+			else if ((attnum < 0) && (attnum >= attnum_expr_lowbound))
+				ok = true;
+
+			if (!ok)
+			{
+				ereport(elevel,
+						(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
+						errmsg("pg_ndistinct: invalid attnum for this statistics object: %d", attnum)));
+				return false;
+			}
+		}
+	}
+	return true;
+}
+
+
 /*
  * pg_ndistinct
  *		output routine for type pg_ndistinct
-- 
2.48.1