diff --git a/pgxn/neon/Makefile b/pgxn/neon/Makefile index 42f2a8efda..c87ae59fd6 100644 --- a/pgxn/neon/Makefile +++ b/pgxn/neon/Makefile @@ -16,6 +16,7 @@ OBJS = \ neon_walreader.o \ pagestore_smgr.o \ relsize_cache.o \ + unstable_extensions.o \ walproposer.o \ walproposer_pg.o \ control_plane_connector.o \ diff --git a/pgxn/neon/neon.c b/pgxn/neon/neon.c index f8ec725c18..dc87d79e87 100644 --- a/pgxn/neon/neon.c +++ b/pgxn/neon/neon.c @@ -30,6 +30,7 @@ #include "neon.h" #include "control_plane_connector.h" #include "logical_replication_monitor.h" +#include "unstable_extensions.h" #include "walsender_hooks.h" #if PG_MAJORVERSION_NUM >= 16 #include "storage/ipc.h" @@ -424,6 +425,7 @@ _PG_init(void) LogicalFuncs_Custom_XLogReaderRoutines = NeonOnDemandXLogReaderRoutines; SlotFuncs_Custom_XLogReaderRoutines = NeonOnDemandXLogReaderRoutines; + InitUnstableExtensionsSupport(); InitLogicalReplicationMonitor(); InitControlPlaneConnector(); diff --git a/pgxn/neon/neon_pgversioncompat.c b/pgxn/neon/neon_pgversioncompat.c index a0dbddde4b..6d139002fe 100644 --- a/pgxn/neon/neon_pgversioncompat.c +++ b/pgxn/neon/neon_pgversioncompat.c @@ -42,3 +42,59 @@ InitMaterializedSRF(FunctionCallInfo fcinfo, bits32 flags) MemoryContextSwitchTo(old_context); } #endif + +#if PG_MAJORVERSION_NUM < 16 +/* + * Some infrastructure for checking malloc/strdup/realloc calls + */ +void * +guc_malloc(int elevel, size_t size) +{ + void *data; + + /* Avoid unportable behavior of malloc(0) */ + if (size == 0) + size = 1; + data = malloc(size); + if (data == NULL) + ereport(elevel, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"))); + return data; +} + +void * +guc_realloc(int elevel, void *old, size_t size) +{ + void *data; + + /* Avoid unportable behavior of realloc(NULL, 0) */ + if (old == NULL && size == 0) + size = 1; + data = realloc(old, size); + if (data == NULL) + ereport(elevel, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"))); + return data; +} + +char * +guc_strdup(int elevel, const char *src) +{ + char *data; + + data = strdup(src); + if (data == NULL) + ereport(elevel, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"))); + return data; +} + +void +guc_free(void *ptr) +{ + free(ptr); +} +#endif diff --git a/pgxn/neon/neon_pgversioncompat.h b/pgxn/neon/neon_pgversioncompat.h index 6b4b355672..84b21ab923 100644 --- a/pgxn/neon/neon_pgversioncompat.h +++ b/pgxn/neon/neon_pgversioncompat.h @@ -83,6 +83,11 @@ InitBufferTag(BufferTag *tag, const RelFileNode *rnode, #define DropRelationAllLocalBuffers DropRelFileNodeAllLocalBuffers +void *guc_malloc(int elevel, size_t size); +void *guc_realloc(int elevel, void *old, size_t size); +char *guc_strdup(int elevel, const char *src); +void guc_free(void *ptr); + #else /* major version >= 16 */ #define USE_RELFILELOCATOR diff --git a/pgxn/neon/unstable_extensions.c b/pgxn/neon/unstable_extensions.c new file mode 100644 index 0000000000..ece77aaf91 --- /dev/null +++ b/pgxn/neon/unstable_extensions.c @@ -0,0 +1,164 @@ +#include +#include +#include + +#include "postgres.h" + +#include "nodes/plannodes.h" +#include "nodes/parsenodes.h" +#include "tcop/utility.h" +#include "utils/errcodes.h" +#include "utils/guc.h" + +#include "neon_pgversioncompat.h" +#include "unstable_extensions.h" + +typedef struct UnstableExtensions +{ + size_t extc; + char *extv[FLEXIBLE_ARRAY_MEMBER]; +} UnstableExtensions; + +static bool allow_unstable_extensions = false; +static char *unstable_extensions_str = NULL; +static UnstableExtensions *unstable_extensions = NULL; + +static ProcessUtility_hook_type PreviousProcessUtilityHook = NULL; + +static bool +check_unstable_extensions(char **newval, void **extra, GucSource source) +{ + char *n; + UnstableExtensions *exts; + char *curr_ext; + size_t max_extc; + + *extra = NULL; + + if (*newval == NULL || (*newval)[0] == '\0') + return true; + + /* At this point, we know we have at least 1 extension like: ext */ + max_extc = 1; + for (size_t i = 0; i < strlen(*newval); ++i) + { + if ((*newval)[i] == ',') + ++max_extc; + } + + exts = guc_malloc(ERROR, sizeof(*exts) + max_extc * sizeof(char *)); + exts->extc = 0; + + n = guc_strdup(ERROR, *newval); + while ((curr_ext = strsep(&n, ","))) + { + /* In the event, we receive a config like: ",,ext" */ + if (curr_ext[0] == '\0') + continue; + + exts->extv[exts->extc++] = guc_strdup(ERROR, curr_ext); + } + + guc_free(n); + + *extra = exts; + + return true; +} + +static void +assign_unstable_extensions(const char *newval, void *extra) +{ + unstable_extensions = extra; +} + +static void +CheckUnstableExtension( + PlannedStmt *pstmt, + const char *queryString, + bool readOnlyTree, + ProcessUtilityContext context, + ParamListInfo params, + QueryEnvironment *queryEnv, + DestReceiver *dest, + QueryCompletion *qc) +{ + Node *parseTree = pstmt->utilityStmt; + + if (allow_unstable_extensions || unstable_extensions == NULL) + goto process; + + switch (nodeTag(parseTree)) + { + case T_CreateExtensionStmt: + { + CreateExtensionStmt *stmt = castNode(CreateExtensionStmt, parseTree); + + for (size_t i = 0; i < unstable_extensions->extc; ++i) + { + if (strcmp(unstable_extensions->extv[i], stmt->extname) == 0) + ereport(ERROR, + (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE), + errmsg("using an unstable extension (%s) is currently prohibited", stmt->extname), + errhint("Set neon.allow_unstable_extensions to true"))); + } + + break; + } + default: + goto process; + } + +process: + if (PreviousProcessUtilityHook) + { + PreviousProcessUtilityHook( + pstmt, + queryString, + readOnlyTree, + context, + params, + queryEnv, + dest, + qc); + } + else + { + standard_ProcessUtility( + pstmt, + queryString, + readOnlyTree, + context, + params, + queryEnv, + dest, + qc); + } +} + +void +InitUnstableExtensionsSupport(void) +{ + DefineCustomBoolVariable( + "neon.allow_unstable_extensions", + "Allow unstable extensions to be installed and used", + NULL, + &allow_unstable_extensions, + false, + PGC_USERSET, + 0, + NULL, NULL, NULL); + + DefineCustomStringVariable( + "neon.unstable_extensions", + "Allow unstable extensions to be installed and used", + NULL, + &unstable_extensions_str, + NULL, + PGC_SUSET, + 0, + check_unstable_extensions, assign_unstable_extensions, NULL); + + PreviousProcessUtilityHook = ProcessUtility_hook; + ProcessUtility_hook = CheckUnstableExtension; +} diff --git a/pgxn/neon/unstable_extensions.h b/pgxn/neon/unstable_extensions.h new file mode 100644 index 0000000000..3c695e9fb2 --- /dev/null +++ b/pgxn/neon/unstable_extensions.h @@ -0,0 +1,6 @@ +#ifndef __NEON_UNSTABLE_EXTENSIONS_H__ +#define __NEON_UNSTABLE_EXTENSIONS_H__ + +void InitUnstableExtensionsSupport(void); + +#endif diff --git a/test_runner/regress/test_unstable_extensions.py b/test_runner/regress/test_unstable_extensions.py new file mode 100644 index 0000000000..06a62ccfd8 --- /dev/null +++ b/test_runner/regress/test_unstable_extensions.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import pytest +from psycopg2.errors import InsufficientPrivilege + +if TYPE_CHECKING: + from fixtures.neon_fixtures import NeonEnv + + +def test_unstable_extensions_installation(neon_simple_env: NeonEnv): + """ + Test that the unstable extension support within the neon extension can + block extension installation. + """ + env = neon_simple_env + + neon_unstable_extensions = "pg_prewarm,amcheck" + + endpoint = env.endpoints.create( + "main", + config_lines=[ + "neon.allow_unstable_extensions=false", + f"neon.unstable_extensions='{neon_unstable_extensions}'", + ], + ) + endpoint.respec(skip_pg_catalog_updates=False) + endpoint.start() + + with endpoint.cursor() as cursor: + cursor.execute("SELECT current_setting('neon.unstable_extensions')") + result = cursor.fetchone() + assert result is not None + setting = cast("str", result[0]) + assert setting == neon_unstable_extensions + + with pytest.raises(InsufficientPrivilege): + cursor.execute("CREATE EXTENSION pg_prewarm") + + with pytest.raises(InsufficientPrivilege): + cursor.execute("CREATE EXTENSION amcheck") + + # Make sure that we can install a "stable" extension + cursor.execute("CREATE EXTENSION pageinspect") + + cursor.execute("BEGIN") + cursor.execute("SET neon.allow_unstable_extensions TO true") + cursor.execute("CREATE EXTENSION pg_prewarm") + cursor.execute("COMMIT")