/*-------------------------------------------------------------------------
 *
 * pg_readonly is a PostgreSQL extension which allows to set a whole
 * cluster read only: no INSERT,UPDATE,DELETE and no DDL can be run.
 *
 * This program is open source, licensed under the PostgreSQL license.
 * For license terms, see the LICENSE file.
 *
 * Copyright (c) 2020, Pierre Forstmann.
 *
 *-------------------------------------------------------------------------
*/
#include "postgres.h"
#include "parser/analyze.h"
#include "nodes/nodes.h"
#include "storage/proc.h"
#include "access/xact.h"

#include "tcop/tcopprot.h"
#include "tcop/utility.h"
#include "utils/guc.h"
#include "utils/snapmgr.h"
#include "utils/memutils.h"
#if PG_VERSION_NUM <= 90600
#include "storage/lwlock.h"
#endif
#if PG_VERSION_NUM < 120000
#include "access/transam.h"
#endif

#include "storage/ipc.h"
#include "storage/spin.h"
#include "miscadmin.h"
#include "storage/procarray.h"
#include "executor/executor.h"
#include "optimizer/optimizer.h"

PG_MODULE_MAGIC;

/*
 * has set_cluster_readonly() been executed
 * in the current backend.
 */
static bool read_only_flag_has_been_set = false;

/*
 *
 * Global shared state
 *
 */
typedef struct pgroSharedState
{
	LWLock	   	*lock;			/* self protection */
	bool		cluster_is_readonly;	/* cluster read-only global flag */
} pgroSharedState;

/* Saved hook values in case of unload */
#if PG_VERSION_NUM >= 150000
static shmem_request_hook_type prev_shmem_request_hook = NULL;
#endif
static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
static ExecutorStart_hook_type prev_executor_start_hook = NULL;
static ProcessUtility_hook_type prev_process_utility_hook = NULL;

/* Links to shared memory state */
static pgroSharedState *pgro= NULL;

static bool pgro_enabled = false;

/*---- Function declarations ----*/

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

static void pgro_shmem_request(void);
static void pgro_shmem_startup(void);
static void pgro_shmem_shutdown(int code, Datum arg);
static void pgro_exec(QueryDesc *queryDesc, int eflags);
static void pgro_utility(PlannedStmt *pstmt, const char *queryString,
						 bool readOnlyTree,
						 ProcessUtilityContext context,
						 ParamListInfo params,
						 QueryEnvironment *queryEnv,
						 DestReceiver *dest, QueryCompletion *qc);

static bool pgro_set_readonly_internal(void);
static bool pgro_unset_readonly_internal(void);
static bool pgro_get_readonly_internal(void);

PG_FUNCTION_INFO_V1(pgro_set_readonly);
PG_FUNCTION_INFO_V1(pgro_unset_readonly);
PG_FUNCTION_INFO_V1(pgro_get_readonly);


/*
 * set cluster databases to read-only
 */

static bool pgro_set_readonly_internal(void)
{

	VirtualTransactionId *tvxid;
	TransactionId limitXmin = InvalidTransactionId;
	bool excludeXmin0 = false;
	bool allDbs = true;
	int excludeVacuum = 0;
	int nvxids;
	int i;
#if PG_VERSION_NUM < 190000
	pid_t pid;
#endif

	elog(LOG, "pg_readonly: killing all transactions ...");
	tvxid = GetCurrentVirtualXIDs(
                 limitXmin,
	         excludeXmin0,
                 allDbs,
                 excludeVacuum,
	         &nvxids);
	for (i=0; i < nvxids; i++)
	{
#if PG_VERSION_NUM >= 190000
		SignalRecoveryConflictWithVirtualXID(
			tvxid[i],
			RECOVERY_CONFLICT_SNAPSHOT);
		elog(LOG, "pg_readonly: virtual transaction signalled");
#else
		pid = CancelVirtualTransaction(
			tvxid[i],
			PROCSIG_RECOVERY_CONFLICT_SNAPSHOT);
		elog(LOG, "pg_readonly: PID %d signalled", pid);
#endif
	}
 	elog(LOG, "pg_readonly: ... done.");


	LWLockAcquire(pgro->lock, LW_EXCLUSIVE);
	pgro->cluster_is_readonly = true;
	LWLockRelease(pgro->lock);
	return true;
}


/*
 * set cluster databases to read write
 */

static bool pgro_unset_readonly_internal(void)
{
	LWLockAcquire(pgro->lock, LW_EXCLUSIVE);
	pgro->cluster_is_readonly = false;
	LWLockRelease(pgro->lock);
	return true;
}


/*
 * get cluster databases read-only or
 * read-write status
 */

static bool pgro_get_readonly_internal(void)
{
	bool val;

	LWLockAcquire(pgro->lock, LW_SHARED);
	val = pgro->cluster_is_readonly;
	LWLockRelease(pgro->lock);
	return val;
}

/*
 * set cluster databases to read-only
 */
Datum pgro_set_readonly(PG_FUNCTION_ARGS)
{
	if (pgro_enabled == false)
	{
		ereport(ERROR, (errmsg("pg_readonly: pgro_set_readonly: pg_readonly is not enabled")));
		PG_RETURN_BOOL(false);
	}
	else
	{
		elog(DEBUG5, "pg_readonly: pgro_set_readonly: entry");
		elog(DEBUG5, "pg_readonly: pgro_set_readonly: exit");
		read_only_flag_has_been_set = true;
		PG_RETURN_BOOL(pgro_set_readonly_internal());
	}
}

/*
 * set cluster databases to read-write
 */
Datum pgro_unset_readonly(PG_FUNCTION_ARGS)
{
	if (pgro_enabled == false)
	{
		ereport(ERROR, (errmsg("pg_readonly: pgro_unset_readonly: pg_readonly is not enabled")));
		PG_RETURN_BOOL(false);
	}
	else
	{
		elog(DEBUG5, "pg_readonly: pgro_unset_readonly: entry");
		elog(DEBUG5, "pg_readonly: pgro_unset_readonly: exit");
		read_only_flag_has_been_set = false;
		PG_RETURN_BOOL(pgro_unset_readonly_internal());
	}

}

/*
 * get cluster databases status
 */
Datum pgro_get_readonly(PG_FUNCTION_ARGS)
{
	if (pgro_enabled == false)
	{
		ereport(ERROR, (errmsg("pg_readonly: pgro_get_readonly: pg_readonly is not enabled")));
		PG_RETURN_BOOL(false);
	}
	else
	{
		elog(DEBUG5, "pg_readonly: pgro_get_readonly: entry");
		elog(DEBUG5, "pg_readonly: pgro_get_readonly: exit");
		PG_RETURN_BOOL(pgro_get_readonly_internal());
	}

}

/*
 ** Estimate shared memory space needed.
 *
 **/
static Size
pgro_memsize(void)
{
	Size		size;

	size = MAXALIGN(sizeof(pgroSharedState));

	return size;
}

/*
 *
 * shmen_request_hook
 */
static void
pgro_shmem_request(void)
{
	/*
 	 * Request additional shared resources.  (These are no-ops if we're not in
 	 * the postmaster process.)  We'll allocate or attach to the shared
 	 * resources in pgls_shmem_startup().
	 */

#if PG_VERSION_NUM >= 150000
	if (prev_shmem_request_hook)
		prev_shmem_request_hook();
#endif

	RequestAddinShmemSpace(sizeof(pgroSharedState));
#if PG_VERSION_NUM >= 90600
	RequestNamedLWLockTranche("pg_readonly", 1);
#endif

}


/*
 * shmem_startup hook: allocate or attach to shared memory.
 *
 */
static void
pgro_shmem_startup(void)
{
	bool		found;

	elog(DEBUG5, "pg_readonly: pgro_shmem_startup: entry");

	if (prev_shmem_startup_hook)
		prev_shmem_startup_hook();

	/* reset in case this is a restart within the postmaster */
	pgro = NULL;


	/*
 	** Create or attach to the shared memory state
 	**/
	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);

	pgro = ShmemInitStruct("pg_readonly",
			        pgro_memsize(),
			        &found);

	if (!found)
	{
		/* First time through ... */
#if PG_VERSION_NUM <= 90600
		RequestAddinLWLocks(1);
		pgro->lock = LWLockAssign();
#else
		pgro->lock = &(GetNamedLWLockTranche("pg_readonly"))->lock;
#endif

		pgro->cluster_is_readonly = false;
	}

	LWLockRelease(AddinShmemInitLock);

	/*
         * If we're in the postmaster (or a standalone backend...), set up a shmem
         * exit hook (no current need ???)
         */
        if (!IsUnderPostmaster)
		on_shmem_exit(pgro_shmem_shutdown, (Datum) 0);

	/*
  	 * Done if some other process already completed our initialization.
  	 */
	if (found)
		return;

	elog(DEBUG5, "pg_readonly: pgro_shmem_startup: exit");

}

/*
 *
 *  shmem_shutdown hook
 *
 *  Note: we don't bother with acquiring lock, because there should be no
 *  other processes running when this is called.
 */
static void
pgro_shmem_shutdown(int code, Datum arg)
{
	elog(DEBUG5, "pg_readonly: pgro_shmem_shutdown: entry");

	/* Don't do anything during a crash. */
	if (code)
		return;

	/* Safety check ... shouldn't get here unless shmem is set up. */
	if (!pgro)
		return;

	/* currently: no action */

	elog(DEBUG5, "pg_readonly: pgro_shmem_shutdown: exit");
}


/*
 * Module load callback.
 *
 * Loading via shared_preload_libraries is required to ensure that hooks
 * are installed in every postgres process.  Without this, only the backend
 * that called LOAD would enforce read-only mode, defeating the purpose of
 * cluster-wide protection.
 *
 * pg_readonly should be listed last in shared_preload_libraries.  Hook
 * chaining means the last-loaded extension runs first.  If another
 * extension is loaded after pg_readonly, its hooks execute before ours
 * and can perform writes before we set transaction_read_only.
 *
 * process_shared_preload_libraries_in_progress is available since PG 9.4,
 * which covers all validated versions (9.5+).
 */
void
_PG_init(void)
{

	elog(DEBUG5, "pg_readonly: _PG_init(): entry");

	if (!process_shared_preload_libraries_in_progress)
	{
		ereport(WARNING,
				(errmsg("pg_readonly must be loaded via shared_preload_libraries")));
		pgro_enabled = false;
	}
	else
		pgro_enabled = true;


	/*
 	** Install hooks
	*/

	if (pgro_enabled)
	{
#if PG_VERSION_NUM >= 150000
		prev_shmem_request_hook = shmem_request_hook;
		shmem_request_hook = pgro_shmem_request;
#else
		pgro_shmem_request();
#endif
		prev_shmem_startup_hook = shmem_startup_hook;
		shmem_startup_hook = pgro_shmem_startup;
		prev_executor_start_hook = ExecutorStart_hook;
 		ExecutorStart_hook = pgro_exec;
		prev_process_utility_hook = ProcessUtility_hook;
		ProcessUtility_hook = pgro_utility;
	}

	elog(DEBUG5, "pg_readonly: _PG_init(): exit");
}


/*
 *  Module unload callback
 */
void
_PG_fini(void)
{
	elog(DEBUG5, "pg_readonly: _PG_fini(): entry");

	/* Uninstall hooks. */
	shmem_startup_hook = prev_shmem_startup_hook;
	ExecutorStart_hook = prev_executor_start_hook;
	ProcessUtility_hook = prev_process_utility_hook;

	elog(DEBUG5, "pg_readonly: _PG_fini(): exit");
}


/*
 * Set transaction_read_only for the current transaction via the GUC
 * machinery.  Using GUC_ACTION_LOCAL means the value is automatically
 * reverted at transaction end — no manual restore is needed.
 *
 * Once set, the user cannot revert to read-write mode within the same
 * transaction: check_transaction_read_only() in variable.c rejects the
 * read-only -> read-write transition after the first snapshot is taken.
 */
static void
pgro_set_xact_readonly(void)
{
	if (XactReadOnly)
		return;

	set_config_option("transaction_read_only", "on",
					  PGC_USERSET, PGC_S_SESSION,
					  GUC_ACTION_LOCAL, true, 0, false);
}

#if PG_VERSION_NUM < 160000
/*
 * Recursively walk a plan tree looking for volatile functions in
 * targetlists and quals.  Throws an error if any are found.
 */
static void
walk_plan(Plan *plan)
{
	if (plan == NULL)
		return;

	if (contain_volatile_functions((Node *) plan->targetlist))
		ereport(ERROR,
				(errmsg("pg_readonly: cannot execute query containing "
						"volatile functions because cluster is read-only")));

	if (contain_volatile_functions((Node *) plan->qual))
		ereport(ERROR,
				(errmsg("pg_readonly: cannot execute query containing "
						"volatile functions because cluster is read-only")));

	walk_plan(plan->lefttree);
	walk_plan(plan->righttree);
}
#endif

/*
 * ExecutorStart hook.
 *
 * When the cluster is read-only, set transaction_read_only = on via proper
 * GUC machinery.  The downstream standard_ExecutorStart will then call
 * ExecCheckXactReadOnly(), which enforces all the read-only checks:
 * DML blocking, temp table exemptions, modifying CTE detection, etc.
 */
static void
pgro_exec(QueryDesc *queryDesc, int eflags)
{
	if (pgro_get_readonly_internal())
	{
		pgro_set_xact_readonly();
#if PG_VERSION_NUM < 160000
		/*
		 * Block SELECT-callable volatile functions that could write to the
		 * database (e.g., lo_create()).  transaction_read_only does not
		 * prevent these.
		 */
		if (queryDesc->plannedstmt && queryDesc->plannedstmt->planTree)
			walk_plan(queryDesc->plannedstmt->planTree);
#endif
	}

	if (prev_executor_start_hook)
		(*prev_executor_start_hook)(queryDesc, eflags);
	else
		standard_ExecutorStart(queryDesc, eflags);
}

/*
 * ProcessUtility hook.
 *
 * When the cluster is read-only, set transaction_read_only = on via proper
 * GUC machinery.  The downstream standard_ProcessUtility uses
 * ClassifyUtilityCommandAsReadOnly() + PreventCommandIfReadOnly() to block
 * DDL and other write commands — no custom whitelist needed.
 */
static void
pgro_utility(PlannedStmt *pstmt, const char *queryString,
			 bool readOnlyTree,
			 ProcessUtilityContext context,
			 ParamListInfo params,
			 QueryEnvironment *queryEnv,
			 DestReceiver *dest, QueryCompletion *qc)
{
	if (pgro_get_readonly_internal())
		pgro_set_xact_readonly();

	if (prev_process_utility_hook)
		(*prev_process_utility_hook)(pstmt, queryString, readOnlyTree,
									 context, params, queryEnv,
									 dest, qc);
	else
		standard_ProcessUtility(pstmt, queryString, readOnlyTree,
								context, params, queryEnv,
								dest, qc);
}