3 files changed, 26 insertions, 1 deletions
diff --git a/contrib/ltree/ltree.h b/contrib/ltree/ltree.h
index 563970a67a4..1b1305b4839 100644
--- a/contrib/ltree/ltree.h
+++ b/contrib/ltree/ltree.h
@@ -5,6 +5,7 @@
 
 #include "fmgr.h"
 #include "tsearch/ts_locale.h"
+#include "utils/memutils.h"
 
 typedef struct
 {
@@ -111,6 +112,8 @@ typedef struct
 
 #define HDRSIZEQT		MAXALIGN(VARHDRSZ + sizeof(int32))
 #define COMPUTESIZE(size,lenofoperand)	( HDRSIZEQT + (size) * sizeof(ITEM) + (lenofoperand) )
+#define LTXTQUERY_TOO_BIG(size,lenofoperand) \
+	((size) > (MaxAllocSize - HDRSIZEQT - (lenofoperand)) / sizeof(ITEM))
 #define GETQUERY(x)  (ITEM*)( (char*)(x)+HDRSIZEQT )
 #define GETOPERAND(x)	( (char*)GETQUERY(x) + ((ltxtquery*)x)->size * sizeof(ITEM) )
 
diff --git a/contrib/ltree/ltree_io.c b/contrib/ltree/ltree_io.c
index 3e88b81c16e..d64debb5f49 100644
--- a/contrib/ltree/ltree_io.c
+++ b/contrib/ltree/ltree_io.c
@@ -8,6 +8,7 @@
 #include <ctype.h>
 
 #include "ltree.h"
+#include "utils/memutils.h"
 #include "crc32.h"
 
 PG_FUNCTION_INFO_V1(ltree_in);
@@ -64,6 +65,11 @@ ltree_in(PG_FUNCTION_ARGS)
 		ptr += charlen;
 	}
 
+	if (num + 1 > MaxAllocSize / sizeof(nodeitem))
+		ereport(ERROR,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+			 errmsg("number of levels (%d) exceeds the maximum allowed (%d)",
+					num + 1, (int) (MaxAllocSize / sizeof(nodeitem)))));
 	list = lptr = (nodeitem *) palloc(sizeof(nodeitem) * (num + 1));
 	ptr = buf;
 	while (*ptr)
@@ -228,6 +234,11 @@ lquery_in(PG_FUNCTION_ARGS)
 	}
 
 	num++;
+	if (num > MaxAllocSize / ITEMSIZE)
+		ereport(ERROR,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+			 errmsg("number of levels (%d) exceeds the maximum allowed (%d)",
+					num, (int) (MaxAllocSize / ITEMSIZE))));
 	curqlevel = tmpql = (lquery_level *) palloc0(ITEMSIZE * num);
 	ptr = buf;
 	while (*ptr)
diff --git a/contrib/ltree/ltxtquery_io.c b/contrib/ltree/ltxtquery_io.c
index 583ff2aaebb..982186581a3 100644
--- a/contrib/ltree/ltxtquery_io.c
+++ b/contrib/ltree/ltxtquery_io.c
@@ -9,6 +9,7 @@
 
 #include "crc32.h"
 #include "ltree.h"
+#include "miscadmin.h"
 
 PG_FUNCTION_INFO_V1(ltxtq_in);
 Datum		ltxtq_in(PG_FUNCTION_ARGS);
@@ -212,6 +213,9 @@ makepol(QPRS_STATE *state)
 	int32		lenstack = 0;
 	uint16		flag = 0;
 
+	/* since this function recurses, it could be driven to stack overflow */
+	check_stack_depth();
+
 	while ((type = gettoken_query(state, &val, &lenval, &strval, &flag)) != END)
 	{
 		switch (type)
@@ -276,6 +280,9 @@ makepol(QPRS_STATE *state)
 static void
 findoprnd(ITEM *ptr, int32 *pos)
 {
+	/* since this function recurses, it could be driven to stack overflow. */
+	check_stack_depth();
+
 	if (ptr[*pos].type == VAL || ptr[*pos].type == VALTRUE)
 	{
 		ptr[*pos].left = 0;
@@ -340,8 +347,12 @@ queryin(char *buf)
 				 errmsg("syntax error"),
 				 errdetail("Empty query.")));
 
-	/* make finish struct */
+	if (LTXTQUERY_TOO_BIG(state.num, state.sumlen))
+		ereport(ERROR,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				 errmsg("ltxtquery is too large")));
 	commonlen = COMPUTESIZE(state.num, state.sumlen);
+
 	query = (ltxtquery *) palloc(commonlen);
 	SET_VARSIZE(query, commonlen);
 	query->size = state.num;