Use RestrictSearchPath to lock down search_path

Change other places that lock down search_path to use RestrictSearchPath instead of SPI commands.
timescale · Jun 19, 2024 · 5836445 · 5836445
1 parent 42a12ee
commit 5836445
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 18 deletions.
diff --git a/src/telemetry/replication.c b/src/telemetry/replication.c
@@ -24,9 +24,8 @@ ts_telemetry_replication_info_gather(void)
 		return info;
 
 	/* Lock down search_path */
-	res = SPI_exec("SET LOCAL search_path TO pg_catalog, pg_temp", 0);
-	if (res < 0)
-		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), (errmsg("could not set search_path"))));
+	int save_nestlevel = NewGUCNestLevel();
+	RestrictSearchPath();
 
 	res = SPI_execute("SELECT cast(count(pid) as int) from pg_catalog.pg_stat_get_wal_senders() "
 					  "WHERE pid is not null",
@@ -59,5 +58,8 @@ ts_telemetry_replication_info_gather(void)
 	if (res != SPI_OK_FINISH)
 		elog(ERROR, "SPI_finish failed: %s", SPI_result_code_string(res));
 
+	/* Restore search_path */
+	AtEOXact_GUC(false, save_nestlevel);
+
 	return info;
 }
diff --git a/src/telemetry/telemetry.c b/src/telemetry/telemetry.c
@@ -360,9 +360,8 @@ add_errors_by_sqlerrcode(JsonbParseState *parse_state)
 		elog(ERROR, "could not connect to SPI");
 
 	/* Lock down search_path */
-	res = SPI_exec("SET LOCAL search_path TO pg_catalog, pg_temp", 0);
-	if (res < 0)
-		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), (errmsg("could not set search_path"))));
+	int save_nestlevel = NewGUCNestLevel();
+	RestrictSearchPath();
 
 	command = makeStringInfo();
 
@@ -398,6 +397,9 @@ add_errors_by_sqlerrcode(JsonbParseState *parse_state)
 		MemoryContextSwitchTo(spi_context);
 	}
 
+	/* Restore search_path */
+	AtEOXact_GUC(false, save_nestlevel);
+
 	res = SPI_finish();
 
 	Assert(res == SPI_OK_FINISH);
@@ -462,9 +464,8 @@ add_job_stats_by_job_type(JsonbParseState *parse_state)
 		elog(ERROR, "could not connect to SPI");
 
 	/* Lock down search_path */
-	res = SPI_exec("SET LOCAL search_path TO pg_catalog, pg_temp", 0);
-	if (res < 0)
-		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), (errmsg("could not set search_path"))));
+	int save_nestlevel = NewGUCNestLevel();
+	RestrictSearchPath();
 
 	command = makeStringInfo();
 
@@ -524,6 +525,10 @@ add_job_stats_by_job_type(JsonbParseState *parse_state)
 		add_job_stats_internal(parse_state, TextDatumGetCString(jobtype_datum), &stats);
 		MemoryContextSwitchTo(spi_context);
 	}
+
+	/* Restore search_path */
+	AtEOXact_GUC(false, save_nestlevel);
+
 	res = SPI_finish();
 	Assert(res == SPI_OK_FINISH);
 }
@@ -795,8 +800,8 @@ add_query_result_dict(JsonbParseState *state, const char *query)
 		elog(ERROR, "could not connect to SPI");
 
 	/* Lock down search_path */
-	res = SPI_execute("SET LOCAL search_path TO pg_catalog, pg_temp", false, 0);
-	Ensure(res >= 0, "could not set search path");
+	int save_nestlevel = NewGUCNestLevel();
+	RestrictSearchPath();
 
 	res = SPI_execute(query, true, 0);
 	Ensure(res >= 0, "could not execute query");
@@ -833,6 +838,10 @@ add_query_result_dict(JsonbParseState *state, const char *query)
 		}
 		pushJsonbValue(&state, WJB_END_OBJECT, NULL);
 	}
+
+	/* Restore search_path */
+	AtEOXact_GUC(false, save_nestlevel);
+
 	MemoryContextSwitchTo(spi_context);
 	res = SPI_finish();
 	Assert(res == SPI_OK_FINISH);

diff --git a/tsl/src/continuous_aggs/materialize.c b/tsl/src/continuous_aggs/materialize.c
@@ -64,12 +64,10 @@ continuous_agg_update_materialization(Hypertable *mat_ht, const ContinuousAgg *c
 {
 	InternalTimeRange combined_materialization_range = new_materialization_range;
 	bool materialize_invalidations_separately = range_length(invalidation_range) > 0;
-	int res;
 
 	/* Lock down search_path */
-	res = SPI_exec("SET LOCAL search_path TO pg_catalog, pg_temp", 0);
-	if (res < 0)
-		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), (errmsg("could not set search_path"))));
+	int save_nestlevel = NewGUCNestLevel();
+	RestrictSearchPath();
 
 	/* pin the start of new_materialization to the end of new_materialization,
 	 * we are not allowed to materialize beyond that point
@@ -131,6 +129,9 @@ continuous_agg_update_materialization(Hypertable *mat_ht, const ContinuousAgg *c
 									internal_time_range_to_time_range(new_materialization_range),
 									chunk_id);
 	}
+
+	/* Restore search_path */
+	AtEOXact_GUC(false, save_nestlevel);
 }
 
 static bool

diff --git a/tsl/src/continuous_aggs/refresh.c b/tsl/src/continuous_aggs/refresh.c
@@ -771,9 +771,8 @@ continuous_agg_refresh_internal(const ContinuousAgg *cagg,
 		elog(ERROR, "SPI_connect failed: %s", SPI_result_code_string(rc));
 
 	/* Lock down search_path */
-	rc = SPI_exec("SET LOCAL search_path TO pg_catalog, pg_temp", 0);
-	if (rc < 0)
-		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), (errmsg("could not set search_path"))));
+	int save_nestlevel = NewGUCNestLevel();
+	RestrictSearchPath();
 
 	/* Like regular materialized views, require owner to refresh. */
 	if (!object_ownercheck(RelationRelationId, cagg->relid, GetUserId()))
@@ -862,6 +861,9 @@ continuous_agg_refresh_internal(const ContinuousAgg *cagg,
 	{
 		emit_up_to_date_notice(cagg, callctx);
 
+		/* Restore search_path */
+		AtEOXact_GUC(false, save_nestlevel);
+
 		rc = SPI_finish();
 		if (rc != SPI_OK_FINISH)
 			elog(ERROR, "SPI_finish failed: %s", SPI_result_code_string(rc));
@@ -882,6 +884,9 @@ continuous_agg_refresh_internal(const ContinuousAgg *cagg,
 	if (!process_cagg_invalidations_and_refresh(cagg, &refresh_window, callctx, INVALID_CHUNK_ID))
 		emit_up_to_date_notice(cagg, callctx);
 
+	/* Restore search_path */
+	AtEOXact_GUC(false, save_nestlevel);
+
 	rc = SPI_finish();
 	if (rc != SPI_OK_FINISH)
 		elog(ERROR, "SPI_finish failed: %s", SPI_result_code_string(rc));