/*
 * 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.
 */

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

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

#include "access/genam.h"
#include "access/heapam.h"
#include "access/xact.h"
#include "catalog/dependency.h"
#include "catalog/indexing.h"
#include "catalog/objectaccess.h"
#include "catalog/pg_auth_members.h"
#include "catalog/pg_authid.h"
#include "catalog/pg_database.h"
#include "catalog/pg_db_role_setting.h"
#include "catalog/namespace.h"
#include "commands/comment.h"
#include "commands/dbcommands.h"
#include "commands/seclabel.h"
#include "commands/user.h"
#include "common/md5.h"
#include "funcapi.h"
#include "miscadmin.h"
#include "storage/lmgr.h"
#include "tcop/utility.h"
#include "utils/acl.h"
#include "utils/builtins.h"
#include "utils/fmgroids.h"
#include "utils/inval.h"
#include "utils/lsyscache.h"
#include "utils/syscache.h"
#include "utils/timestamp.h"
#if PG_MAJOR_VERSION < 1200
#include "utils/tqual.h"
#endif
#include "utils/varlena.h"

/*
 * This code has only been tested with PostgreSQL 9.1.
 *
 * It should be "deprecated" in 9.2 and following thanks to command triggers,
 * and the extension mechanism it works with didn't exist before 9.1.
 */
PG_MODULE_MAGIC;

char	   *extwlist_extensions = NULL;
char	   *extwlist_custom_path = NULL;
bool		extwlist_restrict_to_database_owner = false;

static ProcessUtility_hook_type prev_ProcessUtility = NULL;

void		_PG_init(void);
void		_PG_fini(void);

#if PG_MAJOR_VERSION < 1300
#define PROCESS_UTILITY_PROTO_ARGS PlannedStmt *pstmt,                    \
										const char *queryString,       \
										ProcessUtilityContext context, \
										ParamListInfo params,          \
										QueryEnvironment *queryEnv,    \
										DestReceiver *dest,            \
										char *completionTag

#define PROCESS_UTILITY_ARGS pstmt, queryString, context, \
                              params, queryEnv, dest, completionTag
#elif PG_MAJOR_VERSION < 1400
#define PROCESS_UTILITY_PROTO_ARGS PlannedStmt *pstmt,                    \
										const char *queryString,       \
										ProcessUtilityContext context, \
										ParamListInfo params,          \
										QueryEnvironment *queryEnv,    \
										DestReceiver *dest,            \
										QueryCompletion *qc
#define PROCESS_UTILITY_ARGS pstmt, queryString, context, \
                              params, queryEnv, dest, qc
#else
#define PROCESS_UTILITY_PROTO_ARGS PlannedStmt *pstmt,                    \
										const char *queryString,       \
										bool readOnlyTree,             \
										ProcessUtilityContext context, \
										ParamListInfo params,          \
										QueryEnvironment *queryEnv,    \
										DestReceiver *dest,            \
										QueryCompletion *qc
#define PROCESS_UTILITY_ARGS pstmt, queryString, readOnlyTree, context, \
                              params, queryEnv, dest, qc
#endif							/* PG_MAJOR_VERSION */

#define EREPORT_EXTENSION_IS_NOT_WHITELISTED(op)						\
        ereport(ERROR,                                                  \
                (errcode(ERRCODE_INVALID_PARAMETER_VALUE),              \
                 errmsg("extension \"%s\" is not whitelisted", name),   \
                 errdetail("%s the extension \"%s\" failed, "           \
                           "because it is not on the whitelist of "     \
                           "user-installable extensions.", op, name),	\
                 errhint("Your system administrator has allowed users " \
                         "to install certain extensions. "              \
						 "See: SHOW extwlist.extensions;")));


static void extwlist_ProcessUtility(PROCESS_UTILITY_PROTO_ARGS);
static void call_ProcessUtility(PROCESS_UTILITY_PROTO_ARGS,
								const char *name,
								const char *schema,
								const char *old_version,
								const char *new_version,
								const char *action);
static void call_RawProcessUtility(PROCESS_UTILITY_PROTO_ARGS);

/*
 * Check that the configured extwlist.custom_path exists, is a directory,
 * and is reachable. Empty/NULL disables the feature (no validation).
 *
 * Failing here is intentional: a misconfigured custom_path would silently
 * bypass intended extension hooks, which defeats the security purpose of
 * the whitelist. Better to refuse the configuration outright.
 */
static void
check_extwlist_custom_path_impl(const char *newval)
{
	struct stat st;

	if (newval == NULL || newval[0] == '\0')
		return;

	if (stat(newval, &st) != 0)
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("invalid value for parameter \"extwlist.custom_path\": "
						"could not access \"%s\": %m", newval)));

	if (!S_ISDIR(st.st_mode))
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
				 errmsg("invalid value for parameter \"extwlist.custom_path\": "
						"\"%s\" is not a directory", newval)));

	if (access(newval, R_OK | X_OK) != 0)
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("invalid value for parameter \"extwlist.custom_path\": "
						"directory \"%s\" is not readable: %m", newval)));
}

static bool
check_extwlist_custom_path(char **newval, void **extra, GucSource source)
{
	check_extwlist_custom_path_impl(*newval);
	return true;
}

/*
 * _PG_init()			- library load-time initialization
 *
 * DO NOT make this static nor change its name!
 *
 * Init the module, all we have to do here is getting our GUC
 */
void
_PG_init(void)
{
	DefineCustomStringVariable("extwlist.extensions",
							   "List of extensions that are whitelisted",
							   "Separated by comma",
							   &extwlist_extensions,
							   "",
							   PGC_SUSET,
							   GUC_NOT_IN_SAMPLE,
							   NULL,
							   NULL,
							   NULL);

	DefineCustomStringVariable("extwlist.custom_path",
							   "Directory where to load custom scripts from",
							   "Must point to an existing, readable directory; "
							   "empty disables the custom-script feature.",
							   &extwlist_custom_path,
							   "",
							   PGC_SUSET,
							   GUC_NOT_IN_SAMPLE,
							   check_extwlist_custom_path,
							   NULL,
							   NULL);

	DefineCustomBoolVariable("extwlist.restrict_to_database_owner",
							 "Restrict pgextwlist superuser override to the database owner",
							 "When on, only the database owner may use pgextwlist "
							 "to install, update, drop, or comment on whitelisted "
							 "extensions. Trusted extensions are unaffected.",
							 &extwlist_restrict_to_database_owner,
							 false,
							 PGC_SUSET,
							 GUC_NOT_IN_SAMPLE,
							 NULL,
							 NULL,
							 NULL);

	EmitWarningsOnPlaceholders("extwlist");

	prev_ProcessUtility = ProcessUtility_hook;
	ProcessUtility_hook = extwlist_ProcessUtility;
}

/*
 * Extension Whitelisting includes mechanisms to run custom scripts before and
 * after the extension's provided script.
 *
 * We lookup scripts at the following places and run them when they exist:
 *
 *  ${extwlist_custom_path}/${extname}/${when}--${version}.sql (create)
 *  ${extwlist_custom_path}/${extname}/${when}--${oldversion}--${newversion}.sql (upgrade)
 *  ${extwlist_custom_path}/${extname}/${when}-${action}--${version}.sql (rest)
 *  ${extwlist_custom_path}/${extname}/${when}-${action}.sql (all actions)
 *
 * - action is expected to be one of "create", "update", "comment", or "drop"
 * - when   is expected to be either "before" or "after"
 *
 * We don't validation the extension's name before building the scripts path
 * here because the extension name we are dealing with must have already been
 * added to the whitelist, which should be enough of a validation step.
 */
static void
call_extension_scripts(const char *extname,
					   const char *schema,
					   const char *action,
					   const char *when,
					   const char *from_version,
					   const char *version)
{
	char	   *specific_custom_script;
	char	   *generic_custom_script;

	/* Nothing to do if the custom-script feature is disabled. */
	if (extwlist_custom_path == NULL || extwlist_custom_path[0] == '\0')
		return;

	/*
	 * Ensure the per-extension subdirectory exists and is reachable. ENOENT
	 * is treated as "no custom scripts for this extension" and silently
	 * skipped; anything else is a hard error to avoid bypassing the whitelist
	 * via filesystem misconfiguration.
	 */
	validate_custom_script_dir(extname);

	if (version)
	{
		specific_custom_script =
			get_specific_custom_script_filename(extname, when,
												from_version, version);

		elog(DEBUG1, "Considering custom script \"%s\"", specific_custom_script);

		if (access(specific_custom_script, F_OK) == 0)
		{
			execute_custom_script(specific_custom_script, schema);
			return;				/* skip generic script */
		}
	}

	generic_custom_script =
		get_generic_custom_script_filename(extname, action, when);

	elog(DEBUG1, "Considering custom script \"%s\"", generic_custom_script);

	if (access(generic_custom_script, F_OK) == 0)
		execute_custom_script(generic_custom_script, schema);
}

static bool
extension_is_whitelisted(const char *name)
{
	bool		whitelisted = false;
	char	   *rawnames = pstrdup(extwlist_extensions);
	List	   *extensions;
	ListCell   *lc;

	if (!SplitIdentifierString(rawnames, ',', &extensions))
	{
		/* syntax error in extension name list */
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
				 errmsg("parameter \"extwlist.extensions\" must be a list of extension names")));
	}
	foreach(lc, extensions)
	{
		char	   *curext = (char *) lfirst(lc);

		if (strcmp(name, curext) == 0)
		{
			whitelisted = true;
			break;
		}
	}
	return whitelisted;
}

/*
 * Return true if the current user owns the current database. Used to gate
 * pgextwlist's superuser-override path when extwlist.restrict_to_database_owner
 * is enabled. PG 16 removed pg_database_ownercheck() in favour of the generic
 * object_ownercheck(), so we dispatch on version.
 */
static bool
current_user_is_database_owner(void)
{
#if PG_MAJOR_VERSION >= 1600
	return object_ownercheck(DatabaseRelationId, MyDatabaseId, GetUserId());
#else
	return pg_database_ownercheck(MyDatabaseId, GetUserId());
#endif
}

/*
 * ProcessUtility hook
 */
static void
extwlist_ProcessUtility(PROCESS_UTILITY_PROTO_ARGS)
{
	char	   *name = NULL;
	char	   *schema = NULL;
	char	   *old_version = NULL;
	char	   *new_version = NULL;
	Node	   *parsetree = pstmt->utilityStmt;

	/*
	 * Don't try to make life hard for our friendly superusers. Also, if a
	 * valid transaction is not ongoing then return early.
	 */
	if (!IsTransactionState() || superuser())
	{
		call_RawProcessUtility(PROCESS_UTILITY_ARGS);
		return;
	}

	switch (nodeTag(parsetree))
	{
		case T_CreateExtensionStmt:
			{
				CreateExtensionStmt *stmt = (CreateExtensionStmt *) parsetree;

				name = stmt->extname;
				fill_in_extension_properties(name, stmt->options,
											 &schema, &old_version, &new_version);

				if (extension_is_whitelisted(name))
				{
					if (extwlist_restrict_to_database_owner &&
						!current_user_is_database_owner())
						ereport(ERROR,
								(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
								 errmsg("extension \"%s\" installation via pgextwlist "
										"is restricted to the database owner",
										name)));

					call_ProcessUtility(PROCESS_UTILITY_ARGS,
										name, schema,
										old_version, new_version, "create");
					return;
				}
				break;
			}

		case T_AlterExtensionStmt:
			{
				AlterExtensionStmt *stmt = (AlterExtensionStmt *) parsetree;

				name = stmt->extname;
				fill_in_extension_properties(name, stmt->options,
											 &schema, &old_version, &new_version);

				/* fetch old_version from the catalogs, actually */
				old_version = get_extension_current_version(name);

				if (extension_is_whitelisted(name))
				{
					if (extwlist_restrict_to_database_owner &&
						!current_user_is_database_owner())
						ereport(ERROR,
								(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
								 errmsg("extension \"%s\" update via pgextwlist "
										"is restricted to the database owner",
										name)));

					call_ProcessUtility(PROCESS_UTILITY_ARGS,
										name, schema,
										old_version, new_version, "update");
					return;
				}
				break;
			}

		case T_DropStmt:
			if (((DropStmt *) parsetree)->removeType == OBJECT_EXTENSION)
			{
				/* DROP EXTENSION can target several of them at once */
				bool		all_in_whitelist = true;
				ListCell   *lc;

				foreach(lc, ((DropStmt *) parsetree)->objects)
				{
					/*
					 * For deconstructing the object list into actual names,
					 * see the get_object_address_unqualified() function in
					 * src/backend/catalog/objectaddress.c
					 */
					bool		whitelisted = false;
					List	   *objname = lfirst(lc);
#if PG_MAJOR_VERSION < 1500
					name = strVal((Value *) objname);
#else
					name = strVal(castNode(String, objname));
#endif

					whitelisted = extension_is_whitelisted(name);
					all_in_whitelist = all_in_whitelist && whitelisted;
				}

				/*
				 * If we have a mix of whitelisted and non-whitelisted
				 * extensions in a single DROP EXTENSION command, better play
				 * safe and do the DROP without superpowers.
				 *
				 * So we only give superpowers when all extensions are in the
				 * whitelist.
				 */
				if (all_in_whitelist)
				{
					if (extwlist_restrict_to_database_owner &&
						!current_user_is_database_owner())
						ereport(ERROR,
								(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
								 errmsg("extension drop via pgextwlist "
										"is restricted to the database owner")));

					call_ProcessUtility(PROCESS_UTILITY_ARGS,
										NULL, "",	/* schema must not be NULL */
										NULL, NULL, "drop");
					return;
				}
			}
			break;

		case T_CommentStmt:
			{
				CommentStmt *stmt = (CommentStmt *) parsetree;

				if (stmt->objtype == OBJECT_EXTENSION)
				{
#if PG_MAJOR_VERSION < 1500
					name = strVal((Value *) stmt->object);
#else
					name = strVal(castNode(String, stmt->object));
#endif

					if (extension_is_whitelisted(name))
					{
						if (extwlist_restrict_to_database_owner &&
							!current_user_is_database_owner())
							ereport(ERROR,
									(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
									 errmsg("comment on extension \"%s\" via pgextwlist "
											"is restricted to the database owner",
											name)));

						call_ProcessUtility(PROCESS_UTILITY_ARGS,
											name, "",	/* schema must not be
														 * NULL */
											NULL, NULL, "comment");
						return;
					}
				}
				break;
			}
			/* We intentionally don't support that command. */
		case T_AlterExtensionContentsStmt:
		default:
			break;
	}

	/*
	 * We can only fall here if we don't want to support the command, so pass
	 * control over to the usual processing.
	 */
	call_RawProcessUtility(PROCESS_UTILITY_ARGS);
}

/*
 * Change current user and security context as if running a SECURITY DEFINER
 * procedure owned by a superuser, hard coded as the bootstrap user.
 */
static void
call_ProcessUtility(PROCESS_UTILITY_PROTO_ARGS,
					const char *name,
					const char *schema,
					const char *old_version,
					const char *new_version,
					const char *action)
{
	Oid			save_userid;
	int			save_sec_context;

	GetUserIdAndSecContext(&save_userid, &save_sec_context);

	SetUserIdAndSecContext(BOOTSTRAP_SUPERUSERID,
						   save_sec_context
						   | SECURITY_LOCAL_USERID_CHANGE
						   | SECURITY_RESTRICTED_OPERATION);

	if (action)
	{
		/* "drop extension" can list several extensions, walk them here */
		if (strcmp(action, "drop") == 0)
		{
			Node	   *parsetree = pstmt->utilityStmt;
			ListCell   *lc;
			char	   *name = NULL;

			foreach(lc, ((DropStmt *) parsetree)->objects)
			{
				List	   *objname = lfirst(lc);
#if PG_MAJOR_VERSION < 1500
				name = strVal((Value *) objname);
#else
				name = strVal(castNode(String, objname));
#endif
				call_extension_scripts(name, schema, action,
									   "before", old_version, new_version);
			}
		}
		else
			call_extension_scripts(name, schema, action,
								   "before", old_version, new_version);
	}

	call_RawProcessUtility(PROCESS_UTILITY_ARGS);

	if (action)
	{
		if (strcmp(action, "drop") == 0)
		{
			Node	   *parsetree = pstmt->utilityStmt;
			ListCell   *lc;
			char	   *name = NULL;

			foreach(lc, ((DropStmt *) parsetree)->objects)
			{
				List	   *objname = lfirst(lc);
#if PG_MAJOR_VERSION < 1500
				name = strVal((Value *) objname);
#else
				name = strVal(castNode(String, objname));
#endif
				call_extension_scripts(name, schema, action,
									   "after", old_version, new_version);
			}
		}
		else
			call_extension_scripts(name, schema, action,
								   "after", old_version, new_version);
	}

	SetUserIdAndSecContext(save_userid, save_sec_context);
}

static void
call_RawProcessUtility(PROCESS_UTILITY_PROTO_ARGS)
{
	if (prev_ProcessUtility)
		prev_ProcessUtility(PROCESS_UTILITY_ARGS);
	else
		standard_ProcessUtility(PROCESS_UTILITY_ARGS);
}
