From 3a01bad214fab20977c0df44548d9db2a83678b7 Mon Sep 17 00:00:00 2001
From: Michel Pelletier <michel@supabase.io>
Date: Fri, 29 Sep 2023 11:13:15 -0700
Subject: [PATCH] Add config var to enable or disable event trigger function.

---
 sql/pgsodium--3.1.8--3.1.9.sql | 37 ++++++++++++++++++++++++++++++++++
 src/pgsodium.c                 | 13 +++++++++++-
 test/pgsodium_schema.sql       |  2 +-
 test/tce.sql                   | 35 +++++++++++++++++++++++++++++---
 4 files changed, 82 insertions(+), 5 deletions(-)

diff --git a/sql/pgsodium--3.1.8--3.1.9.sql b/sql/pgsodium--3.1.8--3.1.9.sql
index 946e820..c4a0e96 100644
--- a/sql/pgsodium--3.1.8--3.1.9.sql
+++ b/sql/pgsodium--3.1.8--3.1.9.sql
@@ -125,3 +125,40 @@ $$
   LANGUAGE plpgsql
   SET search_path=''
 ;
+
+CREATE OR REPLACE FUNCTION pgsodium.trg_mask_update()
+RETURNS EVENT_TRIGGER AS
+$$
+DECLARE
+  r record;
+BEGIN
+  IF (SELECT bool_or(in_extension) FROM pg_event_trigger_ddl_commands()) THEN
+    RAISE NOTICE 'skipping pgsodium mask regeneration in extension';
+	RETURN;
+  ELSIF current_setting('pgsodium.enable_event_trigger') <> 'on' THEN
+    RAISE NOTICE 'skipping pgsodium mask regeneration due to false pgsodium.enable_event_trigger';
+	RETURN;
+  END IF;
+
+  FOR r IN
+    SELECT e.*
+    FROM pg_event_trigger_ddl_commands() e
+    WHERE EXISTS (
+      SELECT FROM pg_catalog.pg_class c
+      JOIN pg_catalog.pg_seclabel s ON s.classoid = c.tableoid
+                                   AND s.objoid = c.oid
+      WHERE c.tableoid = e.classid
+        AND e.objid = c.oid
+        AND s.provider = 'pgsodium'
+    )
+  LOOP
+    IF r.object_type in ('table', 'table column')
+    THEN
+      PERFORM pgsodium.update_mask(r.objid);
+    END IF;
+  END LOOP;
+END
+$$
+  LANGUAGE plpgsql
+  SET search_path=''
+;
diff --git a/src/pgsodium.c b/src/pgsodium.c
index b49b94b..ef4cc09 100644
--- a/src/pgsodium.c
+++ b/src/pgsodium.c
@@ -24,6 +24,7 @@ PG_MODULE_MAGIC;
 
 bytea      *pgsodium_secret_key;
 static char *getkey_script = NULL;
+static bool enable_event_trigger = true;
 
 /*
  * Checking the syntax of the masking rules
@@ -120,10 +121,20 @@ _PG_init (void)
 	/* Security label provider hook */
 	register_label_provider ("pgsodium", pgsodium_object_relabel);
 
-	// we're done if not preloaded, otherwise try to get internal shared key
+	// we're done if not preloaded
 	if (!process_shared_preload_libraries_in_progress)
 		return;
 
+    // Variable to enable/disable event trigger
+	DefineCustomBoolVariable("pgsodium.enable_event_trigger",
+							 "Variable to enable/disable event trigger that regenerates triggers and views.",
+							 NULL,
+							 &enable_event_trigger,
+							 true,
+							 PGC_USERSET, 0,
+							 NULL, NULL, NULL);
+
+    // try to get internal shared key
 	path = (char *) palloc0 (MAXPGPATH);
 	get_share_path (my_exec_path, sharepath);
 	snprintf (path, MAXPGPATH, "%s/extension/%s", sharepath, PG_GETKEY_EXEC);
diff --git a/test/pgsodium_schema.sql b/test/pgsodium_schema.sql
index 26ea97d..488918c 100644
--- a/test/pgsodium_schema.sql
+++ b/test/pgsodium_schema.sql
@@ -5604,7 +5604,7 @@ SELECT function_privs_are('pgsodium'::name, proname, proargtypes::regtype[]::tex
     AND oidvectortypes(proargtypes) = 'bytea';
 
 SELECT unnest(ARRAY[
-    is(md5(prosrc), 'b8b02682e0138dc894512f55587db8d4',
+    is(md5(prosrc), '7e6641f8c9f661514f123598b1ca2448',
        format('Function pgsodium.%s(%s) body should match checksum',
               proname, pg_get_function_identity_arguments(oid))
     ),
diff --git a/test/tce.sql b/test/tce.sql
index 2b4687f..30607f2 100644
--- a/test/tce.sql
+++ b/test/tce.sql
@@ -139,16 +139,45 @@ select ok(has_table_privilege('bobo', 'private.bar', 'SELECT'),
 
 select ok(has_table_privilege('bobo', 'private.other_bar', 'SELECT'),
 	'user keeps view select privs after regeneration');
-	
+
 select ok(has_table_privilege('bobo', 'private.other_bar', 'INSERT'),
 	'user keeps view insert privs after regeneration');
-	
+
 select ok(has_table_privilege('bobo', 'private.other_bar', 'UPDATE'),
 	'user keeps view update privs after regeneration');
-	
+
 select ok(has_table_privilege('bobo', 'private.other_bar', 'DELETE'),
 	'user keeps view delete privs after regeneration');
 
+SET pgsodium.enable_event_trigger = 'off';
+
+CREATE TABLE private.fooz(
+  secret text
+);
+
+SELECT lives_ok(
+  format($test$
+         SECURITY LABEL FOR pgsodium ON COLUMN private.fooz.secret
+         IS 'ENCRYPT WITH KEY ID %s'
+         $test$, :'secret_key_id'),
+  'can label column for encryption with event trigger disabled');
+
+SELECT hasnt_view('private', 'decrypted_fooz', 'Dynamic view was not created due to disabled event trigger.');
+
+SELECT hasnt_trigger('private', 'fooz', 'fooz_encrypt_secret_trigger_secret',
+    'Dynamic trigger was not created due to disabled event trigger.');
+
+SELECT lives_ok(
+    $test$SELECT pgsodium.update_mask('private.fooz'::regclass);$test$,
+    'can manually create trigger and view with event trigger disabled.');
+
+SELECT has_view('private', 'decrypted_fooz', 'Dynamic view was created manually.');
+
+SELECT has_trigger('private', 'fooz', 'fooz_encrypt_secret_trigger_secret',
+    'Dynamic trigger was created manually.');
+
+RESET pgsodium.enable_event_trigger;
+
 SET SESSION AUTHORIZATION bobo;
 SET ROLE bobo;