pgsql-switcher-function.1.patch

diff --git a/contrib/dummy_seclabel/dummy_seclabel.c b/contrib/dummy_seclabel/dummy_seclabel.c
index 8bd50a3..ff7b22c 100644
--- a/contrib/dummy_seclabel/dummy_seclabel.c
+++ b/contrib/dummy_seclabel/dummy_seclabel.c
@@ -14,12 +14,60 @@
 
 #include "commands/seclabel.h"
 #include "miscadmin.h"
+#include "utils/builtins.h"
+#include "utils/lsyscache.h"
 
 PG_MODULE_MAGIC;
 
+PG_FUNCTION_INFO_V1(dummy_client_label);
+
 /* Entrypoint of the module */
 void _PG_init(void);
 
+static const char *client_label = NULL;
+
+/* SQL function to show client label */
+Datum dummay_client_label(PG_FUNCTION_ARGS);
+
+static const char *
+dummy_get_client_label(void)
+{
+	return client_label;
+}
+
+static void
+dummy_set_client_label(const char *seclabel)
+{
+	client_label = seclabel;
+}
+
+static const char *
+dummy_get_switched_label(Oid procOid)
+{
+	char   *proname = get_func_name(procOid);
+	char   *result = NULL;
+
+	/*
+	 * It assumes a function whose name contains 'trusted' means
+	 * a label switcher function.
+	 */
+	if (strstr(proname, "trusted") != NULL)
+		result = "trusted";
+
+	pfree(proname);
+
+	return result;
+}
+
+Datum
+dummy_client_label(PG_FUNCTION_ARGS)
+{
+	if (!client_label)
+		PG_RETURN_NULL();
+
+	PG_RETURN_TEXT_P(cstring_to_text(client_label));
+}
+
 static void
 dummy_object_relabel(const ObjectAddress *object, const char *seclabel)
 {
@@ -45,5 +93,17 @@ dummy_object_relabel(const ObjectAddress *object, const char *seclabel)
 void
 _PG_init(void)
 {
-	register_label_provider("dummy", dummy_object_relabel);
+	/*
+	 * XXX - we assume this module is loaded during regression test
+	 * using LOAD command. In normal label providers which is installed
+	 * by shared_preload_libraries, the client label shall be initialized
+	 * on the authentication hook.
+	 */
+	client_label = (superuser() ? "secret" : "unclassified");
+
+	register_label_provider("dummy",
+							dummy_object_relabel,
+							dummy_get_switched_label,
+							dummy_get_client_label,
+							dummy_set_client_label);
 }
diff --git a/src/backend/access/transam/xact.c b/src/backend/access/transam/xact.c
index d2e2e11..5e71b62 100644
--- a/src/backend/access/transam/xact.c
+++ b/src/backend/access/transam/xact.c
@@ -30,6 +30,7 @@
 #include "catalog/namespace.h"
 #include "catalog/storage.h"
 #include "commands/async.h"
+#include "commands/seclabel.h"
 #include "commands/tablecmds.h"
 #include "commands/trigger.h"
 #include "executor/spi.h"
@@ -141,6 +142,7 @@ typedef struct TransactionStateData
 	int			maxChildXids;	/* allocated size of childXids[] */
 	Oid			prevUser;		/* previous CurrentUserId setting */
 	int			prevSecContext; /* previous SecurityRestrictionContext */
+	List	   *prevSecLabels;	/* previous security label of client */
 	bool		prevXactReadOnly;		/* entry-time xact r/o state */
 	bool		startedInRecovery;		/* did we start in recovery? */
 	struct TransactionStateData *parent;		/* back link to parent */
@@ -170,6 +172,7 @@ static TransactionStateData TopTransactionStateData = {
 	0,							/* allocated size of childXids[] */
 	InvalidOid,					/* previous CurrentUserId setting */
 	0,							/* previous SecurityRestrictionContext */
+	NIL,						/* previous security label of client */
 	false,						/* entry-time xact r/o state */
 	false,						/* startedInRecovery */
 	NULL						/* link to parent state block */
@@ -1694,6 +1697,9 @@ StartTransaction(void)
 	/* SecurityRestrictionContext should never be set outside a transaction */
 	Assert(s->prevSecContext == 0);
 
+	/* Save current security label of the client */
+	s->prevSecLabels = GetClientSecLabels();
+
 	/*
 	 * initialize other subsystems for new transaction
 	 */
@@ -2198,6 +2204,9 @@ AbortTransaction(void)
 	 */
 	SetUserIdAndSecContext(s->prevUser, s->prevSecContext);
 
+	/* Reset security labels of the client */
+	SetClientSecLabels(s->prevSecLabels);
+
 	/*
 	 * do abort processing
 	 */
@@ -4042,6 +4051,9 @@ AbortSubTransaction(void)
 	 */
 	SetUserIdAndSecContext(s->prevUser, s->prevSecContext);
 
+	/* Reset security labels of the client */
+	SetClientSecLabels(s->prevSecLabels);
+
 	/*
 	 * We can skip all this stuff if the subxact failed before creating a
 	 * ResourceOwner...
@@ -4181,6 +4193,7 @@ PushTransaction(void)
 	s->state = TRANS_DEFAULT;
 	s->blockState = TBLOCK_SUBBEGIN;
 	GetUserIdAndSecContext(&s->prevUser, &s->prevSecContext);
+	s->prevSecLabels = GetClientSecLabels();
 	s->prevXactReadOnly = XactReadOnly;
 
 	CurrentTransactionState = s;
diff --git a/src/backend/commands/seclabel.c b/src/backend/commands/seclabel.c
index 762bbae..7ac2547 100644
--- a/src/backend/commands/seclabel.c
+++ b/src/backend/commands/seclabel.c
@@ -35,11 +35,20 @@ static void CheckAttributeSecLabel(Relation relation);
 typedef struct
 {
 	const char *provider_name;
-	check_object_relabel_type	hook;
+	check_object_relabel_type	relabel_hook;
+	get_switched_label_type		get_switched_hook;
+	get_client_label_type		get_client_hook;
+	set_client_label_type		set_client_hook;
 } LabelProvider;
 
 static List *label_provider_list = NIL;
 
+typedef struct
+{
+	const char *provider_name;
+	const char *client_label;
+} ClientLabel;
+
 /*
  * ExecSecLabelStmt --
  *
@@ -145,7 +154,7 @@ ExecSecLabelStmt(SecLabelStmt *stmt)
 	}
 
 	/* Provider gets control here, may throw ERROR to veto new label. */
-	(*provider->hook)(&address, stmt->label);
+	(*provider->relabel_hook)(&address, stmt->label);
 
 	/* Apply new label. */
 	SetSecurityLabel(&address, provider->provider_name, stmt->label);
@@ -372,8 +381,134 @@ CheckAttributeSecLabel(Relation relation)
 						RelationGetRelationName(relation))));
 }
 
+/*
+ * GetSwitchedSecLabels
+ *
+ * It returns a list of security labels that is applied during execution
+ * of the supplied procedure. If no label provider considered the function
+ * as a label-switcher, it returns NIL.
+ */
+List *
+GetSwitchedSecLabels(Oid procOid, MemoryContext cxt)
+{
+	MemoryContext	oldcxt;
+	ListCell	   *lcp;
+	List		   *results = NIL;
+
+	oldcxt = MemoryContextSwitchTo(cxt);
+
+	foreach (lcp, label_provider_list)
+    {
+		LabelProvider  *lp = lfirst(lcp);
+		ClientLabel	   *cl;
+		const char	   *label;
+
+		if (!lp->get_switched_hook)
+			continue;
+
+		label = (*lp->get_switched_hook)(procOid);
+		if (!label)
+			continue;
+
+		cl = palloc(sizeof(ClientLabel));
+		cl->provider_name = lp->provider_name;
+		cl->client_label = label;
+
+		results = lappend(results, cl);
+	}
+
+	MemoryContextSwitchTo(oldcxt);
+
+	return results;
+}
+
+/*
+ * ProcIsLabelSwitcher
+ *
+ * True shall be returned, if one of the label provider considers the given
+ * function is a label-switcher. Elsewhere, false shall be returned; that
+ * includes a case when no label provider is installed.
+ */
+bool
+ProcIsLabelSwitcher(Oid procOid)
+{
+	List   *temp = GetSwitchedSecLabels(procOid, CurrentMemoryContext);
+	bool	result;
+
+	result = (temp != NIL ? true : false);
+
+	list_free_deep(temp);
+
+	return result;
+}
+
+/*
+ * GetClientSecLabels
+ *
+ * It returns a list of security label assigned to the client for all the
+ * label providers.
+ */
+List *
+GetClientSecLabels(void)
+{
+	ListCell   *lcp;
+	List	   *results = NIL;
+
+	foreach (lcp, label_provider_list)
+	{
+		LabelProvider  *lp = lfirst(lcp);
+		ClientLabel	   *cl;
+
+		if (!lp->get_client_hook)
+			continue;
+
+		cl = palloc(sizeof(ClientLabel));
+		cl->provider_name = lp->provider_name;
+		cl->client_label = (*lp->get_client_hook)();
+
+		results = lappend(results, cl);
+	}
+	return results;
+}
+
+/*
+ * SetClientSecLabels
+ *
+ * It restores a list of security labels saved at GetClientSecLabels() or
+ * GetSwitchedSecLabels().
+ */
+void
+SetClientSecLabels(List *saved_labels)
+{
+	ListCell   *lcp;
+	ListCell   *lcl;
+
+	foreach (lcp, label_provider_list)
+	{
+		LabelProvider  *lp = lfirst(lcp);
+
+		if (!lp->set_client_hook)
+			continue;
+
+		foreach (lcl, saved_labels)
+		{
+			ClientLabel	   *cl = lfirst(lcl);
+
+			if (strcmp(lp->provider_name, cl->provider_name) == 0)
+			{
+				(*lp->set_client_hook)(cl->client_label);
+				break;
+			}
+		}
+	}
+}
+
 void
-register_label_provider(const char *provider_name, check_object_relabel_type hook)
+register_label_provider(const char *provider_name,
+						check_object_relabel_type relabel_hook,
+						get_switched_label_type get_switched_hook,
+						get_client_label_type get_client_hook,
+						set_client_label_type set_client_hook)
 {
 	LabelProvider  *provider;
 	MemoryContext	oldcxt;
@@ -381,7 +516,10 @@ register_label_provider(const char *provider_name, check_object_relabel_type hoo
 	oldcxt = MemoryContextSwitchTo(TopMemoryContext);
 	provider = palloc(sizeof(LabelProvider));
 	provider->provider_name = pstrdup(provider_name);
-	provider->hook = hook;
+	provider->relabel_hook = relabel_hook;
+	provider->get_switched_hook = get_switched_hook;
+	provider->get_client_hook = get_client_hook;
+	provider->set_client_hook = set_client_hook;
 	label_provider_list = lappend(label_provider_list, provider);
 	MemoryContextSwitchTo(oldcxt);
 }
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index de2e66b..cc42702 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -24,6 +24,7 @@
 #include "catalog/pg_operator.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
+#include "commands/seclabel.h"
 #include "executor/executor.h"
 #include "executor/functions.h"
 #include "miscadmin.h"
@@ -3721,6 +3722,10 @@ inline_function(Oid funcid, Oid result_type, List *args,
 	if (pg_proc_aclcheck(funcid, GetUserId(), ACL_EXECUTE) != ACLCHECK_OK)
 		return NULL;
 
+	/* Check whether the function is not label switcher */
+	if (ProcIsLabelSwitcher(funcid))
+		return NULL;
+
 	/*
 	 * Make a temporary memory context, so that we don't leak all the stuff
 	 * that parsing might create.
@@ -4153,6 +4158,10 @@ inline_set_returning_function(PlannerInfo *root, RangeTblEntry *rte)
 	if (pg_proc_aclcheck(func_oid, GetUserId(), ACL_EXECUTE) != ACLCHECK_OK)
 		return NULL;
 
+	/* Check whether the function is not label switcher */
+	if (ProcIsLabelSwitcher(func_oid))
+		return NULL;
+
 	/*
 	 * OK, let's take a look at the function's pg_proc entry.
 	 */
diff --git a/src/backend/utils/fmgr/fmgr.c b/src/backend/utils/fmgr/fmgr.c
index 1c9d2c2..7671db3 100644
--- a/src/backend/utils/fmgr/fmgr.c
+++ b/src/backend/utils/fmgr/fmgr.c
@@ -18,6 +18,7 @@
 #include "access/tuptoaster.h"
 #include "catalog/pg_language.h"
 #include "catalog/pg_proc.h"
+#include "commands/seclabel.h"
 #include "executor/functions.h"
 #include "executor/spi.h"
 #include "lib/stringinfo.h"
@@ -230,7 +231,8 @@ fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 	 */
 	if (!ignore_security &&
 		(procedureStruct->prosecdef ||
-		 !heap_attisnull(procedureTuple, Anum_pg_proc_proconfig)))
+		 !heap_attisnull(procedureTuple, Anum_pg_proc_proconfig) ||
+		 ProcIsLabelSwitcher(functionId)))
 	{
 		finfo->fn_addr = fmgr_security_definer;
 		finfo->fn_stats = TRACK_FUNC_ALL;		/* ie, never track */
@@ -857,6 +859,7 @@ struct fmgr_security_definer_cache
 	FmgrInfo	flinfo;			/* lookup info for target function */
 	Oid			userid;			/* userid to set, or InvalidOid */
 	ArrayType  *proconfig;		/* GUC values to set, or NULL */
+	List	   *seclabels;		/* security label to switch, or NIL */
 };
 
 /*
@@ -878,6 +881,7 @@ fmgr_security_definer(PG_FUNCTION_ARGS)
 	Oid			save_userid;
 	int			save_sec_context;
 	volatile int save_nestlevel;
+	List	   *save_sec_labels;
 	PgStat_FunctionCallUsage fcusage;
 
 	if (!fcinfo->flinfo->fn_extra)
@@ -914,6 +918,8 @@ fmgr_security_definer(PG_FUNCTION_ARGS)
 			MemoryContextSwitchTo(oldcxt);
 		}
 
+		fcache->seclabels = GetSwitchedSecLabels(fcinfo->flinfo->fn_oid,
+												 fcinfo->flinfo->fn_mcxt);
 		ReleaseSysCache(tuple);
 
 		fcinfo->flinfo->fn_extra = fcache;
@@ -940,6 +946,14 @@ fmgr_security_definer(PG_FUNCTION_ARGS)
 						GUC_ACTION_SAVE);
 	}
 
+	if (fcache->seclabels != NIL)
+	{
+		save_sec_labels = GetClientSecLabels();
+		SetClientSecLabels(fcache->seclabels);
+	}
+	else
+		save_sec_labels = NIL;	/* keep compiler quiet */
+
 	/*
 	 * We don't need to restore GUC or userid settings on error, because the
 	 * ensuing xact or subxact abort will do that.	The PG_TRY block is only
@@ -978,6 +992,11 @@ fmgr_security_definer(PG_FUNCTION_ARGS)
 		AtEOXact_GUC(true, save_nestlevel);
 	if (OidIsValid(fcache->userid))
 		SetUserIdAndSecContext(save_userid, save_sec_context);
+	if (fcache->seclabels)
+	{
+		SetClientSecLabels(save_sec_labels);
+		list_free_deep(save_sec_labels);
+	}
 
 	return result;
 }
diff --git a/src/include/commands/seclabel.h b/src/include/commands/seclabel.h
index 4c3854e..b8985ba 100644
--- a/src/include/commands/seclabel.h
+++ b/src/include/commands/seclabel.h
@@ -12,6 +12,7 @@
 #include "catalog/objectaddress.h"
 #include "nodes/primnodes.h"
 #include "nodes/parsenodes.h"
+#include "utils/memutils.h"
 
 /*
  * Internal APIs
@@ -29,7 +30,22 @@ extern void ExecSecLabelStmt(SecLabelStmt *stmt);
 
 typedef void (*check_object_relabel_type)(const ObjectAddress *object,
 										  const char *seclabel);
-extern void register_label_provider(const char *provider,
-								    check_object_relabel_type hook);
 
+/*
+ * Trusted-Procedure Support
+ */
+extern List	   *GetSwitchedSecLabels(Oid procOid, MemoryContext cxt);
+extern bool		ProcIsLabelSwitcher(Oid procOid);
+extern List	   *GetClientSecLabels(void);
+extern void		SetClientSecLabels(List *savedLabels);
+
+typedef const char *(*get_switched_label_type)(Oid procOid);
+typedef const char *(*get_client_label_type)(void);
+typedef void (*set_client_label_type)(const char *seclabel);
+
+extern void register_label_provider(const char *provider,
+								    check_object_relabel_type relabel_hook,
+									get_switched_label_type get_switch_hook,
+									get_client_label_type get_client_hook,
+									set_client_label_type set_client_hook);
 #endif	/* SECLABEL_H */
diff --git a/src/test/regress/input/security_label.source b/src/test/regress/input/security_label.source
index 810a721..5fdb466 100644
--- a/src/test/regress/input/security_label.source
+++ b/src/test/regress/input/security_label.source
@@ -24,6 +24,9 @@ CREATE DOMAIN seclabel_domain AS text;
 ALTER TABLE seclabel_tbl1 OWNER TO seclabel_user1;
 ALTER TABLE seclabel_tbl2 OWNER TO seclabel_user2;
 
+DROP FUNCTION IF EXISTS dummy_trusted_func();
+DROP FUNCTION IF EXISTS dummy_regular_func();
+
 RESET client_min_messages;
 
 --
@@ -37,6 +40,9 @@ SECURITY LABEL ON TABLE seclabel_tbl3 IS 'unclassified';			-- fail
 -- Load dummy external security provider
 LOAD '@libdir@/dummy_seclabel@DLSUFFIX@';
 
+CREATE OR REPLACE FUNCTION dummy_client_label() RETURNS text LANGUAGE 'c'
+    AS '@libdir@/dummy_seclabel@DLSUFFIX@', 'dummy_client_label';
+
 --
 -- Test of SECURITY LABEL statement with a plugin
 --
@@ -70,7 +76,18 @@ SELECT objtype, objname, provider, label FROM pg_seclabels
 SECURITY LABEL ON LANGUAGE plpgsql IS NULL;						-- OK
 SECURITY LABEL ON SCHEMA public IS NULL;						-- OK
 
+-- test for label switcher function
+RESET SESSION AUTHORIZATION;
+
+CREATE OR REPLACE FUNCTION dummy_trusted_func() RETURNS text LANGUAGE 'sql'
+    AS 'SELECT dummy_client_label()';
+CREATE OR REPLACE FUNCTION dummy_regular_func() RETURNS text LANGUAGE 'sql'
+    AS 'SELECT dummy_client_label()';
+
+SELECT dummy_trusted_func(), dummy_regular_func();
+
 -- clean up objects
+RESET SESSION AUTHORIZATION;
 DROP FUNCTION seclabel_four();
 DROP DOMAIN seclabel_domain;
 DROP VIEW seclabel_view1;
@@ -78,6 +95,9 @@ DROP TABLE seclabel_tbl1;
 DROP TABLE seclabel_tbl2;
 DROP USER seclabel_user1;
 DROP USER seclabel_user2;
+DROP FUNCTION dummy_trusted_func();
+DROP FUNCTION dummy_regular_func();
+DROP FUNCTION dummy_client_label();
 
 -- make sure we don't have any leftovers
 SELECT objtype, objname, provider, label FROM pg_seclabels
diff --git a/src/test/regress/output/security_label.source b/src/test/regress/output/security_label.source
index 4bc803d..3356edb 100644
--- a/src/test/regress/output/security_label.source
+++ b/src/test/regress/output/security_label.source
@@ -17,6 +17,8 @@ CREATE FUNCTION seclabel_four() RETURNS integer AS $$SELECT 4$$ language sql;
 CREATE DOMAIN seclabel_domain AS text;
 ALTER TABLE seclabel_tbl1 OWNER TO seclabel_user1;
 ALTER TABLE seclabel_tbl2 OWNER TO seclabel_user2;
+DROP FUNCTION IF EXISTS dummy_trusted_func();
+DROP FUNCTION IF EXISTS dummy_regular_func();
 RESET client_min_messages;
 --
 -- Test of SECURITY LABEL statement without a plugin
@@ -30,7 +32,9 @@ ERROR:  no security label providers have been loaded
 SECURITY LABEL ON TABLE seclabel_tbl3 IS 'unclassified';			-- fail
 ERROR:  no security label providers have been loaded
 -- Load dummy external security provider
-LOAD '@abs_builddir@/dummy_seclabel@DLSUFFIX@';
+LOAD '@libdir@/dummy_seclabel@DLSUFFIX@';
+CREATE OR REPLACE FUNCTION dummy_client_label() RETURNS text LANGUAGE 'c'
+    AS '@libdir@/dummy_seclabel@DLSUFFIX@', 'dummy_client_label';
 --
 -- Test of SECURITY LABEL statement with a plugin
 --
@@ -75,7 +79,20 @@ SELECT objtype, objname, provider, label FROM pg_seclabels
 
 SECURITY LABEL ON LANGUAGE plpgsql IS NULL;						-- OK
 SECURITY LABEL ON SCHEMA public IS NULL;						-- OK
+-- test for label switcher function
+RESET SESSION AUTHORIZATION;
+CREATE OR REPLACE FUNCTION dummy_trusted_func() RETURNS text LANGUAGE 'sql'
+    AS 'SELECT dummy_client_label()';
+CREATE OR REPLACE FUNCTION dummy_regular_func() RETURNS text LANGUAGE 'sql'
+    AS 'SELECT dummy_client_label()';
+SELECT dummy_trusted_func(), dummy_regular_func();
+ dummy_trusted_func | dummy_regular_func 
+--------------------+--------------------
+ trusted            | secret
+(1 row)
+
 -- clean up objects
+RESET SESSION AUTHORIZATION;
 DROP FUNCTION seclabel_four();
 DROP DOMAIN seclabel_domain;
 DROP VIEW seclabel_view1;
@@ -83,6 +100,9 @@ DROP TABLE seclabel_tbl1;
 DROP TABLE seclabel_tbl2;
 DROP USER seclabel_user1;
 DROP USER seclabel_user2;
+DROP FUNCTION dummy_trusted_func();
+DROP FUNCTION dummy_regular_func();
+DROP FUNCTION dummy_client_label();
 -- make sure we don't have any leftovers
 SELECT objtype, objname, provider, label FROM pg_seclabels
 	ORDER BY objtype, objname;