/*
 * PostgreSQL Extension WhiteList -- Dimitri Fontaine
 *
 * Author: Dimitri Fontaine <dimitri@2ndQuadrant.fr>
 * Licence: PostgreSQL
 * Copyright Dimitri Fontaine, 2011-2013
 *
 * For a description of the features see the README.md file from the same
 * distribution.
 */

/*
 * Some tools to read a SQL file and execute commands found in there.
 *
 * The following code comes from the PostgreSQL source tree in
 * src/backend/commands/extension.c, with some modifications to run in the
 * context of the Extension Whitelisting Extension.
 */

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include "postgres.h"

#include "pgextwlist.h"
#include "utils.h"

#include "access/htup_details.h"
#include "access/genam.h"
#include "access/heapam.h"
#include "access/skey.h"
#include "access/xact.h"
#include "catalog/indexing.h"
#include "catalog/namespace.h"
#include "catalog/pg_collation.h"
#include "catalog/pg_database.h"
#include "catalog/pg_extension.h"
#include "commands/extension.h"
#include "executor/executor.h"
#include "mb/pg_wchar.h"
#include "miscadmin.h"
#include "storage/fd.h"
#include "tcop/pquery.h"
#include "tcop/utility.h"
#include "utils/builtins.h"
#include "utils/memutils.h"
#include "utils/fmgroids.h"
#include "utils/lsyscache.h"
#include "utils/rel.h"
#include "utils/snapmgr.h"
#include "utils/syscache.h"
#if PG_MAJOR_VERSION < 1200
#include "utils/tqual.h"
#define table_open(r, l) heap_open(r, l)
#define table_close(r, l) heap_close(r, l)
#endif

/*
 * Parse contents of primary or auxiliary control file, and fill in
 * fields of *control.	We parse primary file if version == NULL,
 * else the optional auxiliary file for that version.
 *
 * Control files are supposed to be very short, half a dozen lines,
 * so we don't worry about memory allocation risks here.  Also we don't
 * worry about what encoding it's in; all values are expected to be ASCII.
 */
static void
parse_default_version_in_control_file(const char *extname,
									  char **version,
									  char **schema)
{
	char		sharepath[MAXPGPATH];
	char	   *filename;
	FILE	   *file;
	ConfigVariable *item,
			   *head = NULL,
			   *tail = NULL;

	/*
	 * Locate the file to read.
	 */
	get_share_path(my_exec_path, sharepath);
	filename = (char *) palloc(MAXPGPATH);
	snprintf(filename, MAXPGPATH, "%s/extension/%s.control", sharepath, extname);

	if ((file = AllocateFile(filename, "r")) == NULL)
	{
		/* we still need to handle the following error here */
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not open extension control file \"%s\": %m",
						filename)));
	}

	/*
	 * Parse the file content, using GUC's file parsing code.  We need not
	 * check the return value since any errors will be thrown at ERROR level.
	 */
	(void) ParseConfigFp(file, filename, 0, ERROR, &head, &tail);

	FreeFile(file);

	/*
	 * Convert the ConfigVariable list into ExtensionControlFile entries, we
	 * are only interested into the default version.
	 */
	for (item = head; item != NULL; item = item->next)
	{
		if (*version == NULL && strcmp(item->name, "default_version") == 0)
		{
			*version = pstrdup(item->value);
		}
		else if (*schema == NULL && strcmp(item->name, "schema") == 0)
		{
			*schema = pstrdup(item->value);
		}
	}

	FreeConfigVariables(head);

	pfree(filename);
}

/*
 * Returns true when extwlist.custom_path is set to a usable value, false
 * when it's NULL or empty (feature disabled).
 */
static bool
custom_path_is_set(void)
{
	return extwlist_custom_path != NULL && extwlist_custom_path[0] != '\0';
}

/*
 * We lookup scripts at the following places and run them when they exist:
 *
 *  ${extwlist_custom_path}/${extname}/${when}--${version}.sql (upgrade)
 *  ${extwlist_custom_path}/${extname}/${when}-${action}.sql (create)
 *  ${extwlist_custom_path}/${extname}/${when}-${action}--${version}.sql (rest)
 */
char *
get_generic_custom_script_filename(const char *name,
								   const char *action,
								   const char *when)
{
	char	   *result;

	if (!custom_path_is_set())
		return NULL;

	result = (char *) palloc(MAXPGPATH);
	snprintf(result, MAXPGPATH, "%s/%s/%s-%s.sql",
			 extwlist_custom_path, name, when, action);

	return result;
}

char *
get_specific_custom_script_filename(const char *name,
									const char *when,
									const char *from_version,
									const char *version)
{
	char	   *result;

	if (!custom_path_is_set())
		return NULL;

	result = (char *) palloc(MAXPGPATH);
	if (from_version)
		snprintf(result, MAXPGPATH, "%s/%s/%s--%s--%s.sql",
				 extwlist_custom_path, name, when, from_version, version);
	else
		snprintf(result, MAXPGPATH, "%s/%s/%s--%s.sql",
				 extwlist_custom_path, name, when, version);

	return result;
}

/*
 * Build ${extwlist.custom_path}/${extname} into a palloc'd buffer. Returns
 * NULL if custom_path is unset; the caller is responsible for pfree.
 */
char *
get_custom_script_dir(const char *extname)
{
	char	   *result;

	if (!custom_path_is_set())
		return NULL;

	result = (char *) palloc(MAXPGPATH);
	snprintf(result, MAXPGPATH, "%s/%s", extwlist_custom_path, extname);

	return result;
}

/*
 * Ensure the per-extension subdirectory under extwlist.custom_path is
 * reachable. ENOENT means the extension simply has no custom scripts and
 * is treated as a silent skip; any other stat() failure is a hard error
 * because silently dropping hooks would defeat the security purpose of
 * the whitelist.
 */
void
validate_custom_script_dir(const char *extname)
{
	char	   *dir;
	struct stat st;

	if (!custom_path_is_set())
		return;

	dir = get_custom_script_dir(extname);

	if (stat(dir, &st) != 0)
	{
		if (errno == ENOENT)
		{
			/* subdir legitimately absent; nothing to run */
			pfree(dir);
			return;
		}

		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not access custom script directory \"%s\": %m",
						dir)));
	}

	if (!S_ISDIR(st.st_mode))
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
				 errmsg("custom script path \"%s\" is not a directory", dir)));

	pfree(dir);
}

/*
 * At CREATE EXTENSION UPDATE time, we generally aren't provided with the
 * current version of the extension to upgrade, go fetch it from the catalogs.
 */
char *
get_extension_current_version(const char *extname)
{
	char	   *oldVersionName;
	Relation	extRel;
	ScanKeyData key[1];
	SysScanDesc extScan;
	HeapTuple	extTup;
	Datum		datum;
	bool		isnull;

	/*
	 * Look up the extension --- it must already exist in pg_extension
	 */
	extRel = table_open(ExtensionRelationId, AccessShareLock);

	ScanKeyInit(&key[0],
				Anum_pg_extension_extname,
				BTEqualStrategyNumber, F_NAMEEQ,
				CStringGetDatum(extname));

	extScan = systable_beginscan(extRel, ExtensionNameIndexId, true,
								 SnapshotSelf, 1, key);

	extTup = systable_getnext(extScan);

	if (!HeapTupleIsValid(extTup))
		ereport(ERROR,
				(errcode(ERRCODE_UNDEFINED_OBJECT),
				 errmsg("extension \"%s\" does not exist", extname)));

	/*
	 * Determine the existing version we are updating from
	 */
	datum = heap_getattr(extTup, Anum_pg_extension_extversion,
						 RelationGetDescr(extRel), &isnull);
	if (isnull)
		elog(ERROR, "extversion is null");
	oldVersionName = text_to_cstring(DatumGetTextPP(datum));

	systable_endscan(extScan);

	table_close(extRel, AccessShareLock);

	return oldVersionName;
}

/*
 * Read the statement's option list and set given parameters.
 */
void
fill_in_extension_properties(const char *extname,
							 List *options,
							 char **schema,
							 char **old_version,
							 char **new_version)
{
	ListCell   *lc;
	DefElem    *d_schema = NULL;
	DefElem    *d_new_version = NULL;
	DefElem    *d_old_version = NULL;

	/*
	 * Read the statement option list, taking care not to issue any errors
	 * here ourselves if at all possible: let the core code handle them.
	 */
	foreach(lc, options)
	{
		DefElem    *defel = (DefElem *) lfirst(lc);

		if (strcmp(defel->defname, "schema") == 0)
		{
			d_schema = defel;
		}
		else if (strcmp(defel->defname, "new_version") == 0)
		{
			d_new_version = defel;
		}
		else if (strcmp(defel->defname, "old_version") == 0)
		{
			d_old_version = defel;
		}
		else
		{
			/* intentionally don't try and catch errors here */
		}
	}

	if (d_schema && d_schema->arg)
		*schema = strVal(d_schema->arg);

	if (d_old_version && d_old_version->arg)
		*old_version = strVal(d_old_version->arg);

	if (d_new_version && d_new_version->arg)
		*new_version = strVal(d_new_version->arg);

	if (*new_version == NULL || *schema == NULL)
		/* fetch the default_version from the extension's control file */
		parse_default_version_in_control_file(extname, new_version, schema);

	/* schema might be given neither in the statement nor the control file */
	if (*schema == NULL)
	{
		/*
		 * Use the current default creation namespace, which is the first
		 * explicit entry in the search_path.
		 */
		Oid			schemaOid;
		List	   *search_path = fetch_search_path(false);

		if (search_path == NIL) /* nothing valid in search_path? */
			ereport(ERROR,
					(errcode(ERRCODE_UNDEFINED_SCHEMA),
					 errmsg("no schema has been selected to create in")));
		schemaOid = linitial_oid(search_path);
		*schema = get_namespace_name(schemaOid);
		if (*schema == NULL)	/* recently-deleted namespace? */
			ereport(ERROR,
					(errcode(ERRCODE_UNDEFINED_SCHEMA),
					 errmsg("no schema has been selected to create in")));

		list_free(search_path);
	}
}

/*
 * Read an SQL script file into a string, and convert to database encoding
 */
static char *
read_custom_script_file(const char *filename)
{
	int			src_encoding,
				dest_encoding = GetDatabaseEncoding();
	bytea	   *content;
	char	   *src_str;
	char	   *dest_str;
	int			len;
	FILE	   *fp;
	struct stat fst;
	size_t		nbytes;

	/*
	 * read_binary_file was made static in 9.5 so we'll reimplement the logic
	 * here
	 */
	if ((fp = AllocateFile(filename, PG_BINARY_R)) == NULL)
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not open file \"%s\" for reading: %m",
						filename)));

	if (fstat(fileno(fp), &fst) < 0)
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not stat file \"%s\" %m",
						filename)));

	content = (bytea *) palloc((Size) fst.st_size + VARHDRSZ);
	nbytes = fread(VARDATA(content), 1, (size_t) fst.st_size, fp);

	if (ferror(fp))
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not read file \"%s\": %m", filename)));

	FreeFile(fp);
	SET_VARSIZE(content, nbytes + VARHDRSZ);

	/* use database encoding */
	src_encoding = dest_encoding;

	/* make sure that source string is valid in the expected encoding */
	len = VARSIZE_ANY_EXHDR(content);
	src_str = VARDATA_ANY(content);
	pg_verify_mbstr_len(src_encoding, src_str, len, false);

	/* convert the encoding to the database encoding */
	dest_str = (char *) pg_do_encoding_conversion((unsigned char *) src_str,
												  len,
												  src_encoding,
												  dest_encoding);

	/* if no conversion happened, we have to arrange for null termination */
	if (dest_str == src_str)
	{
		dest_str = (char *) palloc(len + 1);
		memcpy(dest_str, src_str, len);
		dest_str[len] = '\0';
	}

	return dest_str;
}

/*
 * Execute given SQL string.
 *
 * filename is used only to report errors.
 *
 * Note: it's tempting to just use SPI to execute the string, but that does
 * not work very well.	The really serious problem is that SPI will parse,
 * analyze, and plan the whole string before executing any of it; of course
 * this fails if there are any planable statements referring to objects
 * created earlier in the script.  A lesser annoyance is that SPI insists
 * on printing the whole string as errcontext in case of any error, and that
 * could be very long.
 */
static void
execute_sql_string(const char *sql, const char *filename)
{
	List	   *raw_parsetree_list;
	DestReceiver *dest;
	ListCell   *lc1;
	MemoryContext temp_ctx = AllocSetContextCreate(CurrentMemoryContext,
												   "temp_script_context",
												   ALLOCSET_DEFAULT_SIZES);
	MemoryContext prev_ctx = MemoryContextSwitchTo(temp_ctx);

	/*
	 * Parse the SQL string into a list of raw parse trees.
	 */
	raw_parsetree_list = pg_parse_query(sql);

	/* All output from SELECTs goes to the bit bucket */
	dest = CreateDestReceiver(DestNone);

	/*
	 * Do parse analysis, rule rewrite, planning, and execution for each raw
	 * parsetree.  We must fully execute each query before beginning parse
	 * analysis on the next one, since there may be interdependencies.
	 */
	foreach(lc1, raw_parsetree_list)
	{
		RawStmt    *parsetree = lfirst_node(RawStmt, lc1);
		List	   *stmt_list;
		ListCell   *lc2;

#if PG_MAJOR_VERSION >= 1500
		stmt_list = pg_analyze_and_rewrite_fixedparams(parsetree,
													   sql,
													   NULL,
													   0,
													   NULL
			);
#else
		stmt_list = pg_analyze_and_rewrite(parsetree,
										   sql,
										   NULL,
										   0
										   ,NULL
			);
#endif
		stmt_list = pg_plan_queries(stmt_list,
#if PG_MAJOR_VERSION >= 1300
									sql,
#endif
									0,
									NULL);

		foreach(lc2, stmt_list)
		{
			PlannedStmt *stmt = lfirst_node(PlannedStmt, lc2);

			if (IsA(stmt, TransactionStmt))
				ereport(ERROR,
						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
						 errmsg("transaction control statements are not allowed within an extension script")));

			CommandCounterIncrement();

			PushActiveSnapshot(GetTransactionSnapshot());

			if (IsA(stmt, PlannedStmt) &&
				((PlannedStmt *) stmt)->utilityStmt == NULL)
			{
				QueryDesc  *qdesc;

				qdesc = CreateQueryDesc((PlannedStmt *) stmt, sql,
										GetActiveSnapshot(), NULL,
										dest, NULL, NULL, 0);

				ExecutorStart(qdesc, 0);
				ExecutorRun(qdesc, ForwardScanDirection, 0
#if PG_MAJOR_VERSION < 1800
							,true
#endif
					);
				ExecutorFinish(qdesc);
				ExecutorEnd(qdesc);

				FreeQueryDesc(qdesc);
			}
			else
			{
				ProcessUtility(stmt, sql,
#if PG_MAJOR_VERSION >= 1400
							   false,	/* no need to copy */
#endif
							   PROCESS_UTILITY_QUERY,
							   NULL, NULL, dest, NULL);
			}

			PopActiveSnapshot();
		}
	}

	/* Be sure to advance the command counter after the last script command */
	CommandCounterIncrement();
	MemoryContextSwitchTo(prev_ctx);
	MemoryContextDelete(temp_ctx);
}

/*
 * get_current_database_owner_name
 *
 * select rolname
 *   from pg_roles u join pg_database d on d.datdba = u.oid
 *  where datname = current_database();
 */
static char *
get_current_database_owner_name(void)
{
	HeapTuple	dbtuple;
	Oid			owner;

	dbtuple = SearchSysCache1(DATABASEOID, MyDatabaseId);
	if (HeapTupleIsValid(dbtuple))
	{
		owner = ((Form_pg_database) GETSTRUCT(dbtuple))->datdba;
		ReleaseSysCache(dbtuple);
	}
	else
		return NULL;

	return GetUserNameFromId(owner, false);
}

/*
 * Reject @-substitution values containing characters that could be used to
 * break out of quoting in the resulting script (CVE-2023-39417; upstream
 * fix cd5f2a3570). The check is applied to the raw, unquoted name because
 * quote_identifier() output always contains '"', which would trip the
 * check spuriously.
 */
static void
check_substitution_value(const char *placeholder, const char *value)
{
	static const char forbidden[] = "\"$'\\";

	if (value == NULL)
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
				 errmsg("invalid value for \"%s\" substitution",
						placeholder)));

	if (strpbrk(value, forbidden) != NULL)
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
				 errmsg("invalid value for \"%s\" substitution",
						placeholder),
				 errdetail("Value must not contain any of: %s", forbidden),
				 errhint("Rename the object referenced by %s.",
						 placeholder)));
}

/*
 * Execute given script
 */
void
execute_custom_script(const char *filename, const char *schemaName)
{
	int			save_nestlevel;
	StringInfoData pathbuf;

	elog(DEBUG1, "Executing custom script \"%s\"", filename);

	/*
	 * Force client_min_messages and log_min_messages to be at least WARNING,
	 * so that we won't spam the user with useless NOTICE messages from common
	 * script actions like creating shell types.
	 *
	 * We use the equivalent of a function SET option to allow the setting to
	 * persist for exactly the duration of the script execution.  guc.c also
	 * takes care of undoing the setting on error.
	 */
	save_nestlevel = NewGUCNestLevel();

	if (client_min_messages < WARNING)
		(void) set_config_option("client_min_messages", "warning",
								 PGC_USERSET, PGC_S_SESSION,
								 GUC_ACTION_SAVE, true, 0, false
			);

	/*
	 * log_min_messages was changed from a scalar int (PG <= 18) to something
	 * more complex in PG 19. Set it unconditionally here for simplicity.
	 */
	(void) set_config_option("log_min_messages", "warning",
							 PGC_SUSET, PGC_S_SESSION,
							 GUC_ACTION_SAVE, true, 0, false);

	/*
	 * Set up the search path to contain the target schema, then the schemas
	 * of any prerequisite extensions, and nothing else.  In particular this
	 * makes the target schema be the default creation target namespace.
	 *
	 * Note: it might look tempting to use PushOverrideSearchPath for this,
	 * but we cannot do that.  We have to actually set the search_path GUC in
	 * case the extension script examines or changes it.  In any case, the
	 * GUC_ACTION_SAVE method is just as convenient.
	 */
	initStringInfo(&pathbuf);
	appendStringInfoString(&pathbuf, quote_identifier(schemaName));

	(void) set_config_option("search_path", pathbuf.data,
							 PGC_USERSET, PGC_S_SESSION,
							 GUC_ACTION_SAVE, true, 0, false);

	PG_TRY();
	{
		char	   *c_sql = read_custom_script_file(filename);
		Datum		t_sql;

		/* We use various functions that want to operate on text datums */
		t_sql = CStringGetTextDatum(c_sql);

		/*
		 * Reduce any lines beginning with "\echo" to empty.  This allows
		 * scripts to contain messages telling people not to run them via
		 * psql, which has been found to be necessary due to old habits.
		 */
		t_sql = DirectFunctionCall4Coll(textregexreplace,
										C_COLLATION_OID,
										t_sql,
										CStringGetTextDatum("^\\\\echo.*$"),
										CStringGetTextDatum(""),
										CStringGetTextDatum("ng"));

		/*
		 * substitute the target schema name for occurrences of @extschema@.
		 */
		check_substitution_value("@extschema@", schemaName);
		t_sql = DirectFunctionCall3Coll(replace_text,
										C_COLLATION_OID,
										t_sql,
										CStringGetTextDatum("@extschema@"),
										CStringGetTextDatum(schemaName));

		/*
		 * substitute the current user name for occurrences of @current_user@
		 */
		{
			char	   *cur_user = GetUserNameFromId(GetUserId(), false);

			check_substitution_value("@current_user@", cur_user);
			t_sql = DirectFunctionCall3Coll(replace_text,
											C_COLLATION_OID,
											t_sql,
											CStringGetTextDatum("@current_user@"),
											CStringGetTextDatum(cur_user));
		}

		/*
		 * substitute the database owner for occurrences of @database_owner@
		 */
		{
			char	   *db_owner = get_current_database_owner_name();

			check_substitution_value("@database_owner@", db_owner);
			t_sql = DirectFunctionCall3Coll(replace_text,
											C_COLLATION_OID,
											t_sql,
											CStringGetTextDatum("@database_owner@"),
											CStringGetTextDatum(db_owner));
		}

		/* And now back to C string */
		c_sql = text_to_cstring(DatumGetTextPP(t_sql));

		execute_sql_string(c_sql, filename);
	}
	PG_CATCH();
	{
		PG_RE_THROW();
	}
	PG_END_TRY();

	/*
	 * Restore the GUC variables we set above.
	 */
	AtEOXact_GUC(true, save_nestlevel);
}
