From edcab4bef4d3384763a0a45b4453cbce327f50f1 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <bossartn@amazon.com>
Date: Thu, 13 May 2021 23:04:34 +0000
Subject: [PATCH v1 1/1] Allow specifying direct role membership in
 pg_hba.conf.

---
 doc/src/sgml/client-auth.sgml        | 29 +++++++++++++++++-----------
 src/backend/libpq/hba.c              | 22 ++++++++++++++-------
 src/backend/libpq/pg_hba.conf.sample | 10 ++++++----
 src/backend/utils/adt/acl.c          | 37 ++++++++++++++++++++++++++++++++----
 src/include/utils/acl.h              |  1 +
 5 files changed, 73 insertions(+), 26 deletions(-)

diff --git a/doc/src/sgml/client-auth.sgml b/doc/src/sgml/client-auth.sgml
index 951af49e9a..3639903550 100644
--- a/doc/src/sgml/client-auth.sgml
+++ b/doc/src/sgml/client-auth.sgml
@@ -249,18 +249,21 @@ hostnogssenc  <replaceable>database</replaceable>  <replaceable>user</replaceabl
        Specifies which database user name(s) this record
        matches. The value <literal>all</literal> specifies that it
        matches all users.  Otherwise, this is either the name of a specific
-       database user, or a group name preceded by <literal>+</literal>.
-       (Recall that there is no real distinction between users and groups
-       in <productname>PostgreSQL</productname>; a <literal>+</literal> mark really means
-       <quote>match any of the roles that are directly or indirectly members
-       of this role</quote>, while a name without a <literal>+</literal> mark matches
-       only that specific role.) For this purpose, a superuser is only
+       database user, or a group name preceded by <literal>+</literal> or
+       <literal>&amp;</literal>.  Group names preceded by <literal>+</literal> match
+       roles that are directly or indirectly members of this role, and group
+       names preceded by <literal>&amp;</literal> only match roles that are directly
+       members of this role.  (Recall that there is no real distinction between
+       users and groups in <productname>PostgreSQL</productname>; a
+       <literal>+</literal> or <literal>&amp;</literal> mark really means
+       <quote>match any of the roles that are members of this role</quote>,
+       while a name without a <literal>+</literal> or <literal>&amp;</literal> mark
+       matches only that specific role.)  For this purpose, a superuser is only
        considered to be a member of a role if they are explicitly a member
-       of the role, directly or indirectly, and not just by virtue of
-       being a superuser.
-       Multiple user names can be supplied by separating them with commas.
-       A separate file containing user names can be specified by preceding the
-       file name with <literal>@</literal>.
+       of the role and not just by virtue of being a superuser.  Multiple user
+       names can be supplied by separating them with commas.  A separate file
+       containing user names can be specified by preceding the file name with
+       <literal>@</literal>.
       </para>
      </listitem>
     </varlistentry>
@@ -800,6 +803,10 @@ local   all             +support                                md5
 # The last two lines above can be combined into a single line:
 local   all             @admins,+support                        md5
 
+# Group names can also be specified using "&amp;" to indicate that only direct
+# membership should be considered
+local   all             &amp;support                                md5
+
 # The database column can also use lists and file names:
 local   db1,db2,@demodbs  all                                   md5
 </programlisting>
diff --git a/src/backend/libpq/hba.c b/src/backend/libpq/hba.c
index 60767f2957..5dd0e6ad39 100644
--- a/src/backend/libpq/hba.c
+++ b/src/backend/libpq/hba.c
@@ -581,7 +581,7 @@ tokenize_file(const char *filename, FILE *file, List **tok_lines, int elevel)
  * We check to see if it is a member of the specified role name.
  */
 static bool
-is_member(Oid userid, const char *role)
+is_member(Oid userid, const char *role, bool only_direct)
 {
 	Oid			roleid;
 
@@ -594,11 +594,14 @@ is_member(Oid userid, const char *role)
 		return false;			/* if target role not exist, say "no" */
 
 	/*
-	 * See if user is directly or indirectly a member of role. For this
-	 * purpose, a superuser is not considered to be automatically a member of
-	 * the role, so group auth only applies to explicit membership.
+	 * See if user is a member of role. For this purpose, a superuser is not
+	 * considered to be automatically a member of the role, so group auth only
+	 * applies to explicit membership.
 	 */
-	return is_member_of_role_nosuper(userid, roleid);
+	if (only_direct)
+		return is_direct_member_of_role_nosuper(userid, roleid);
+	else
+		return is_member_of_role_nosuper(userid, roleid);
 }
 
 /*
@@ -615,7 +618,12 @@ check_role(const char *role, Oid roleid, List *tokens)
 		tok = lfirst(cell);
 		if (!tok->quoted && tok->string[0] == '+')
 		{
-			if (is_member(roleid, tok->string + 1))
+			if (is_member(roleid, tok->string + 1, false))
+				return true;
+		}
+		else if (!tok->quoted && tok->string[0] == '&')
+		{
+			if (is_member(roleid, tok->string + 1, true))
 				return true;
 		}
 		else if (token_matches(tok, role) ||
@@ -656,7 +664,7 @@ check_db(const char *dbname, const char *role, Oid roleid, List *tokens)
 		else if (token_is_keyword(tok, "samegroup") ||
 				 token_is_keyword(tok, "samerole"))
 		{
-			if (is_member(roleid, dbname))
+			if (is_member(roleid, dbname, false))
 				return true;
 		}
 		else if (token_is_keyword(tok, "replication"))
diff --git a/src/backend/libpq/pg_hba.conf.sample b/src/backend/libpq/pg_hba.conf.sample
index 5f3f63eb0c..8aef6ae742 100644
--- a/src/backend/libpq/pg_hba.conf.sample
+++ b/src/backend/libpq/pg_hba.conf.sample
@@ -31,10 +31,12 @@
 # keyword does not match "replication". Access to replication
 # must be enabled in a separate record (see example below).
 #
-# USER can be "all", a user name, a group name prefixed with "+", or a
-# comma-separated list thereof.  In both the DATABASE and USER fields
-# you can also write a file name prefixed with "@" to include names
-# from a separate file.
+# USER can be "all", a user name, a group name prefixed with "+" or
+# "&", or a comma-separated list thereof.  Group names prefixed with
+# "+" match both direct and indirect members while group names
+# prefixed with "&" only match direct members.  In both the DATABASE
+# and USER fields you can also write a file name prefixed with "@" to
+# include names from a separate file.
 #
 # ADDRESS specifies the set of hosts the record matches.  It can be a
 # host name, or it is made up of an IP address and a CIDR mask that is
diff --git a/src/backend/utils/adt/acl.c b/src/backend/utils/adt/acl.c
index 67f8b29434..0096f66aeb 100644
--- a/src/backend/utils/adt/acl.c
+++ b/src/backend/utils/adt/acl.c
@@ -65,10 +65,11 @@ typedef struct
 enum RoleRecurseType
 {
 	ROLERECURSE_PRIVS = 0,		/* recurse if rolinherit */
-	ROLERECURSE_MEMBERS = 1		/* recurse unconditionally */
+	ROLERECURSE_MEMBERS = 1,	/* recurse unconditionally */
+	ROLERECURSE_DIRECT = 2		/* do not recurse */
 };
-static Oid	cached_role[] = {InvalidOid, InvalidOid};
-static List *cached_roles[] = {NIL, NIL};
+static Oid	cached_role[] = {InvalidOid, InvalidOid, InvalidOid};
+static List *cached_roles[] = {NIL, NIL, NIL};
 static uint32 cached_db_hash;
 
 
@@ -4687,6 +4688,7 @@ RoleMembershipCacheCallback(Datum arg, int cacheid, uint32 hashvalue)
 	/* Force membership caches to be recomputed on next use */
 	cached_role[ROLERECURSE_PRIVS] = InvalidOid;
 	cached_role[ROLERECURSE_MEMBERS] = InvalidOid;
+	cached_role[ROLERECURSE_DIRECT] = InvalidOid;
 }
 
 
@@ -4711,7 +4713,8 @@ has_rolinherit(Oid roleid)
  * Get a list of roles that the specified roleid is a member of
  *
  * Type ROLERECURSE_PRIVS recurses only through roles that have rolinherit
- * set, while ROLERECURSE_MEMBERS recurses through all roles.  This sets
+ * set, ROLERECURSE_MEMBERS recurses through all roles, and ROLERECURSE_DIRECT
+ * only retrieves roles that roleid has direct membership in.  This sets
  * *is_admin==true if and only if role "roleid" has an ADMIN OPTION membership
  * in role "admin_of".
  *
@@ -4809,6 +4812,10 @@ roles_is_member_of(Oid roleid, enum RoleRecurseType type,
 		if (memberid == dba && OidIsValid(dba))
 			roles_list = list_append_unique_oid(roles_list,
 												ROLE_PG_DATABASE_OWNER);
+
+		/* if we are only interested in direct membership, we're done */
+		if (type == ROLERECURSE_DIRECT)
+			break;
 	}
 
 	/*
@@ -4899,6 +4906,28 @@ check_is_member_of_role(Oid member, Oid role)
 						GetUserNameFromId(role, false))));
 }
 
+/*
+ * Is member a direct member of role, not considering superuserness?
+ *
+ * This is identical to is_member_of_role_nosuper except we do not check for
+ * indirect membership in the target role.
+ */
+bool
+is_direct_member_of_role_nosuper(Oid member, Oid role)
+{
+	/* Fast path for simple case */
+	if (member == role)
+		return true;
+
+	/*
+	 * Find all the roles that member is a member of, but do not recurse, then
+	 * see if target role is any one of them.
+	 */
+	return list_member_oid(roles_is_member_of(member, ROLERECURSE_DIRECT,
+											  InvalidOid, NULL),
+							role);
+}
+
 /*
  * Is member a member of role, not considering superuserness?
  *
diff --git a/src/include/utils/acl.h b/src/include/utils/acl.h
index af771c901d..22ad4bb58c 100644
--- a/src/include/utils/acl.h
+++ b/src/include/utils/acl.h
@@ -207,6 +207,7 @@ extern int	aclmembers(const Acl *acl, Oid **roleids);
 
 extern bool has_privs_of_role(Oid member, Oid role);
 extern bool is_member_of_role(Oid member, Oid role);
+extern bool is_direct_member_of_role_nosuper(Oid member, Oid role);
 extern bool is_member_of_role_nosuper(Oid member, Oid role);
 extern bool is_admin_of_role(Oid member, Oid role);
 extern void check_is_member_of_role(Oid member, Oid role);
-- 
2.16.6

