diff --git a/nix/pg_tle.nix b/nix/pg_tle.nix new file mode 100644 index 0000000..2713c79 --- /dev/null +++ b/nix/pg_tle.nix @@ -0,0 +1,26 @@ +{ stdenv, fetchFromGitHub, postgresql, flex }: + +stdenv.mkDerivation rec { + pname = "pg_tle"; + version = "1.0.4"; + + nativeBuildInputs = [ flex ]; + buildInputs = [ postgresql ]; + + src = fetchFromGitHub { + owner = "aws"; + repo = pname; + rev = "refs/tags/v${version}"; + hash = "sha256-W/7pLy/27VatCdzUh1NZ4K2FRMD1erfHiFV2eY2x2W0="; + }; + + makeFlags = [ "FLEX=flex" ]; + + installPhase = '' + mkdir -p $out/{lib,share/postgresql/extension} + + cp *.so $out/lib + cp *.sql $out/share/postgresql/extension + cp *.control $out/share/postgresql/extension + ''; +} diff --git a/shell.nix b/shell.nix index be04c82..1c33338 100644 --- a/shell.nix +++ b/shell.nix @@ -18,7 +18,10 @@ let }; pgWithExt = { postgresql } : let - pg = postgresql.withPackages (p: [ (supautils {inherit postgresql;}) ]); + pg = postgresql.withPackages (p: [ + (supautils { inherit postgresql; }) + (callPackage ./nix/pg_tle.nix { inherit postgresql; }) + ]); ver = builtins.head (builtins.splitVersion postgresql.version); script = '' export PATH=${pg}/bin:"$PATH" @@ -34,11 +37,11 @@ let PGTZ=UTC initdb --no-locale --encoding=UTF8 --nosync -U "$PGUSER" - options="-F -c listen_addresses=\"\" -k $PGDATA -c shared_preload_libraries=\"supautils\"" + options="-F -c listen_addresses=\"\" -k $PGDATA -c shared_preload_libraries=\"pg_tle, supautils\"" reserved_roles="supabase_storage_admin, anon, reserved_but_not_yet_created, authenticator*" reserved_memberships="pg_read_server_files, pg_write_server_files, pg_execute_server_program, role_with_reserved_membership" - privileged_extensions="hstore, postgres_fdw" + privileged_extensions="hstore, postgres_fdw, pg_tle" privileged_extensions_custom_scripts_path="$tmpdir/privileged_extensions_custom_scripts" privileged_role="privileged_role" privileged_role_allowed_configs="session_replication_role, pgrst.*, other.nested.*" @@ -49,6 +52,15 @@ let pg_ctl start -o "$options" -o "$reserved_stuff_options" -o "$placeholder_stuff_options" mkdir -p "$tmpdir/privileged_extensions_custom_scripts/hstore" + echo "do \$\$ + begin + if not exists (select from pg_extension where extname = 'pg_tle') then + return; + end if; + if exists (select from pgtle.available_extensions() where name = @extname@) then + raise notice 'extname: %, extschema: %, extversion: %, extcascade: %', @extname@, @extschema@, @extversion@, @extcascade@; + end if; + end \$\$;" > "$tmpdir/privileged_extensions_custom_scripts/before-create.sql" echo 'create table t1();' > "$tmpdir/privileged_extensions_custom_scripts/hstore/before-create.sql" echo 'drop table t1; create table t2 as values (1);' > "$tmpdir/privileged_extensions_custom_scripts/hstore/after-create.sql" diff --git a/src/privileged_extensions.c b/src/privileged_extensions.c index 950a4f1..4b2ebc4 100644 --- a/src/privileged_extensions.c +++ b/src/privileged_extensions.c @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -18,37 +19,68 @@ #include "privileged_extensions.h" #include "utils.h" -// TODO: interpolate extschema, current_role, current_database_owner -static void run_custom_script(const char *filename) { +// Prevent recursively running custom scripts +static bool running_custom_script = false; + +static void run_custom_script(const char *filename, const char *extname, + const char *extschema, const char *extversion, + bool extcascade) { + if (running_custom_script) { + return; + } + running_custom_script = true; PushActiveSnapshot(GetTransactionSnapshot()); SPI_connect(); { - char *sql_tmp1 = "do $$\n" - "begin\n" - " execute pg_read_file("; - char *sql_tmp2 = quote_literal_cstr(filename); - char *sql_tmp3 = ");\n" - "exception\n" - " when undefined_file then\n" - " -- skip\n" - "end\n" - "$$;"; - size_t sql_len = strlen(sql_tmp1) + strlen(sql_tmp2) + strlen(sql_tmp3); + char *sql_tmp01 = "do $$\n" + "begin\n" + " execute\n" + " replace(\n" + " replace(\n" + " replace(\n" + " replace(\n" + " pg_read_file(\n"; + char *sql_tmp02 = quote_literal_cstr(filename); + char *sql_tmp03 = " ),\n" + " '@extname@', "; + char *sql_tmp04 = quote_literal_cstr(quote_literal_cstr(extname)); + char *sql_tmp05 = " ),\n" + " '@extschema@', "; + char *sql_tmp06 = quote_literal_cstr( + extschema == NULL ? "null" : quote_literal_cstr(extschema)); + char *sql_tmp07 = " ),\n" + " '@extversion@', "; + char *sql_tmp08 = quote_literal_cstr( + extversion == NULL ? "null" : quote_literal_cstr(extversion)); + char *sql_tmp09 = " ), " + " '@extcascade@', "; + char *sql_tmp10 = extcascade ? "'true'" : "'false'"; + char *sql_tmp11 = " );\n" + "exception\n" + " when undefined_file then\n" + " -- skip\n" + "end\n" + "$$;"; + size_t sql_len = + strlen(sql_tmp01) + strlen(sql_tmp02) + strlen(sql_tmp03) + + strlen(sql_tmp04) + strlen(sql_tmp05) + strlen(sql_tmp06) + + strlen(sql_tmp07) + strlen(sql_tmp08) + strlen(sql_tmp09) + + strlen(sql_tmp10) + strlen(sql_tmp11); char *sql = (char *)palloc(sql_len); int rc; - snprintf(sql, sql_len, "%s%s%s", sql_tmp1, sql_tmp2, sql_tmp3); + snprintf(sql, sql_len, "%s%s%s%s%s%s%s%s%s%s%s", sql_tmp01, sql_tmp02, + sql_tmp03, sql_tmp04, sql_tmp05, sql_tmp06, sql_tmp07, + sql_tmp08, sql_tmp09, sql_tmp10, sql_tmp11); rc = SPI_execute(sql, false, 0); if (rc != SPI_OK_UTILITY) { elog(ERROR, "SPI_execute failed with error code %d", rc); } - - pfree(sql_tmp2); - pfree(sql); } SPI_finish(); PopActiveSnapshot(); + running_custom_script = false; } void handle_create_extension( @@ -59,13 +91,72 @@ void handle_create_extension( CreateExtensionStmt *stmt = (CreateExtensionStmt *)pstmt->utilityStmt; char *filename = (char *)palloc(MAXPGPATH); - // Run before-create script. + // Run global before-create script. + { + DefElem *d_schema = NULL; + DefElem *d_new_version = NULL; + DefElem *d_cascade = NULL; + char *schemaName = NULL; + char *versionName = NULL; + bool cascade = false; + ListCell *option_cell = NULL; + + foreach (option_cell, stmt->options) { + DefElem *defel = (DefElem *)lfirst(option_cell); + + if (strcmp(defel->defname, "schema") == 0) { + d_schema = defel; + schemaName = defGetString(d_schema); + } else if (strcmp(defel->defname, "new_version") == 0) { + d_new_version = defel; + versionName = defGetString(d_new_version); + } else if (strcmp(defel->defname, "cascade") == 0) { + d_cascade = defel; + cascade = defGetBoolean(d_cascade); + } + } + + switch_to_superuser(privileged_extensions_superuser); + + snprintf(filename, MAXPGPATH, "%s/before-create.sql", + privileged_extensions_custom_scripts_path); + run_custom_script(filename, stmt->extname, schemaName, versionName, + cascade); + + switch_to_original_role(); + } + + // Run per-extension before-create script. { + DefElem *d_schema = NULL; + DefElem *d_new_version = NULL; + DefElem *d_cascade = NULL; + char *schemaName = NULL; + char *versionName = NULL; + bool cascade = false; + ListCell *option_cell = NULL; + + foreach (option_cell, stmt->options) { + DefElem *defel = (DefElem *)lfirst(option_cell); + + if (strcmp(defel->defname, "schema") == 0) { + d_schema = defel; + schemaName = defGetString(d_schema); + } else if (strcmp(defel->defname, "new_version") == 0) { + d_new_version = defel; + versionName = defGetString(d_new_version); + } else if (strcmp(defel->defname, "cascade") == 0) { + d_cascade = defel; + cascade = defGetBoolean(d_cascade); + } + } + switch_to_superuser(privileged_extensions_superuser); snprintf(filename, MAXPGPATH, "%s/%s/before-create.sql", privileged_extensions_custom_scripts_path, stmt->extname); - run_custom_script(filename); + run_custom_script(filename, stmt->extname, schemaName, versionName, + cascade); switch_to_original_role(); } @@ -82,13 +173,37 @@ void handle_create_extension( run_process_utility_hook(process_utility_hook); } - // Run after-create script. + // Run per-extension after-create script. { + DefElem *d_schema = NULL; + DefElem *d_new_version = NULL; + DefElem *d_cascade = NULL; + char *schemaName = NULL; + char *versionName = NULL; + bool cascade = false; + ListCell *option_cell = NULL; + + foreach (option_cell, stmt->options) { + DefElem *defel = (DefElem *)lfirst(option_cell); + + if (strcmp(defel->defname, "schema") == 0) { + d_schema = defel; + schemaName = defGetString(d_schema); + } else if (strcmp(defel->defname, "new_version") == 0) { + d_new_version = defel; + versionName = defGetString(d_new_version); + } else if (strcmp(defel->defname, "cascade") == 0) { + d_cascade = defel; + cascade = defGetBoolean(d_cascade); + } + } + switch_to_superuser(privileged_extensions_superuser); snprintf(filename, MAXPGPATH, "%s/%s/after-create.sql", privileged_extensions_custom_scripts_path, stmt->extname); - run_custom_script(filename); + run_custom_script(filename, stmt->extname, schemaName, versionName, + cascade); switch_to_original_role(); } diff --git a/test/expected/privileged_extensions.out b/test/expected/privileged_extensions.out index 7e0b2de..f773cd8 100644 --- a/test/expected/privileged_extensions.out +++ b/test/expected/privileged_extensions.out @@ -1,5 +1,6 @@ -- non-superuser extensions role create role extensions_role login; +grant all on database postgres to extensions_role; alter default privileges for role postgres in schema public grant all on tables to extensions_role; set role extensions_role; \echo @@ -15,7 +16,7 @@ select '1=>2'::hstore; drop extension hstore; \echo --- custom scripts are run +-- per-extension custom scripts are run select * from t2; column1 --------- @@ -27,6 +28,26 @@ drop table t2; set role extensions_role; \echo +-- global extension custom scripts are run +create extension pg_tle; +reset role; +grant pgtle_admin to extensions_role; +set role extensions_role; +select pgtle.install_extension('foo', '1', '', 'select 1', '{}'); + install_extension +------------------- + t +(1 row) + +create extension foo cascade; +NOTICE: extname: foo, extschema: , extversion: , extcascade: t +drop extension pg_tle cascade; +NOTICE: drop cascades to 3 other objects +DETAIL: drop cascades to function pgtle."foo.control"() +drop cascades to function pgtle."foo--1.sql"() +drop cascades to extension foo +\echo + -- custom scripts are run even for superusers reset role; create extension hstore; diff --git a/test/sql/privileged_extensions.sql b/test/sql/privileged_extensions.sql index 1881083..b042ae5 100644 --- a/test/sql/privileged_extensions.sql +++ b/test/sql/privileged_extensions.sql @@ -1,5 +1,6 @@ -- non-superuser extensions role create role extensions_role login; +grant all on database postgres to extensions_role; alter default privileges for role postgres in schema public grant all on tables to extensions_role; set role extensions_role; \echo @@ -11,7 +12,7 @@ select '1=>2'::hstore; drop extension hstore; \echo --- custom scripts are run +-- per-extension custom scripts are run select * from t2; reset role; @@ -19,6 +20,17 @@ drop table t2; set role extensions_role; \echo +-- global extension custom scripts are run +create extension pg_tle; +reset role; +grant pgtle_admin to extensions_role; +set role extensions_role; +select pgtle.install_extension('foo', '1', '', 'select 1', '{}'); +create extension foo cascade; + +drop extension pg_tle cascade; +\echo + -- custom scripts are run even for superusers reset role; create extension hstore;