From cfdbf4a7f6da17096e8ff1a2534452647de282f1 Mon Sep 17 00:00:00 2001
From: Nikita Glukhov <n.gluhov@postgrespro.ru>
Date: Mon, 29 Aug 2022 23:56:45 +0300
Subject: [PATCH v10 06/10] Add safe input function for jsonb

---
 src/backend/utils/adt/jsonb.c | 26 ++++++++++++++++++--------
 src/include/utils/jsonb.h     |  2 ++
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/backend/utils/adt/jsonb.c b/src/backend/utils/adt/jsonb.c
index f700c5b4c93..1da519daaef 100644
--- a/src/backend/utils/adt/jsonb.c
+++ b/src/backend/utils/adt/jsonb.c
@@ -46,7 +46,6 @@ typedef struct JsonbAggState
 	Oid			val_output_func;
 } JsonbAggState;
 
-static inline Datum jsonb_from_cstring(char *json, int len, bool unique_keys);
 static size_t checkStringLen(size_t len);
 static void jsonb_in_object_start(void *pstate);
 static void jsonb_in_object_end(void *pstate);
@@ -77,7 +76,7 @@ jsonb_in(PG_FUNCTION_ARGS)
 {
 	char	   *json = PG_GETARG_CSTRING(0);
 
-	return jsonb_from_cstring(json, strlen(json), false);
+	PG_RETURN_DATUM(jsonb_from_cstring(json, strlen(json), false, NULL));
 }
 
 /*
@@ -101,7 +100,7 @@ jsonb_recv(PG_FUNCTION_ARGS)
 	else
 		elog(ERROR, "unsupported jsonb version number %d", version);
 
-	return jsonb_from_cstring(str, nbytes, false);
+	return jsonb_from_cstring(str, nbytes, false, NULL);
 }
 
 /*
@@ -147,7 +146,7 @@ jsonb_from_text(text *js, bool unique_keys)
 {
 	return jsonb_from_cstring(VARDATA_ANY(js),
 							  VARSIZE_ANY_EXHDR(js),
-							  unique_keys);
+							  unique_keys, NULL);
 }
 
 /*
@@ -239,12 +238,13 @@ jsonb_typeof(PG_FUNCTION_ARGS)
  *
  * Uses the json parser (with hooks) to construct a jsonb.
  */
-static inline Datum
-jsonb_from_cstring(char *json, int len, bool unique_keys)
+Datum
+jsonb_from_cstring(char *json, int len, bool unique_keys, bool *error)
 {
 	JsonLexContext *lex;
 	JsonbInState state;
 	JsonSemAction sem;
+	JsonParseErrorType result;
 
 	memset(&state, 0, sizeof(state));
 	memset(&sem, 0, sizeof(sem));
@@ -261,10 +261,20 @@ jsonb_from_cstring(char *json, int len, bool unique_keys)
 	sem.scalar = jsonb_in_scalar;
 	sem.object_field_start = jsonb_in_object_field_start;
 
-	pg_parse_json_or_ereport(lex, &sem);
+	result = pg_parse_json(lex, &sem);
+	if (result != JSON_SUCCESS)
+	{
+		if (error)
+		{
+			*error = true;
+			return (Datum) 0;
+		}
+
+		json_ereport_error(result, lex);
+	}
 
 	/* after parsing, the item member has the composed jsonb structure */
-	PG_RETURN_POINTER(JsonbValueToJsonb(state.res));
+	return JsonbPGetDatum(JsonbValueToJsonb(state.res));
 }
 
 static size_t
diff --git a/src/include/utils/jsonb.h b/src/include/utils/jsonb.h
index bae466b5234..d5574faa79e 100644
--- a/src/include/utils/jsonb.h
+++ b/src/include/utils/jsonb.h
@@ -420,6 +420,8 @@ extern void JsonbHashScalarValueExtended(const JsonbValue *scalarVal,
 
 /* jsonb.c support functions */
 extern Datum jsonb_from_text(text *js, bool unique_keys);
+extern Datum jsonb_from_cstring(char *json, int len, bool unique_keys,
+								bool *error);
 extern char *JsonbToCString(StringInfo out, JsonbContainer *in,
 							int estimated_len);
 extern char *JsonbToCStringIndent(StringInfo out, JsonbContainer *in,
-- 
2.17.1