[OSS-ONLY] [Regression] Restrict dropping of BBF logins via PG port (b…

…abelfish-for-postgresql#3494) (babelfish-for-postgresql#3537) Restrict dropping of BBF logins via PG port: To identify if a role is a Babelfish login, we need to check if a login is a member of any Babelfish role. Task: BABEL-5565 Signed-off-by: Shalini Lohia [email protected]
amazon-aurora · Feb 27, 2025 · 8b6a640 · 8b6a640
1 parent 262247e
commit 8b6a640
Show file tree

Hide file tree

Showing 3 changed files with 207 additions and 8 deletions.
diff --git a/contrib/babelfishpg_tds/src/backend/tds/tdsutils.c b/contrib/babelfishpg_tds/src/backend/tds/tdsutils.c
@@ -19,6 +19,7 @@
 #include "access/genam.h"
 #include "access/htup_details.h"
 #include "access/table.h"
+#include "catalog/pg_auth_members.h"
 #include "catalog/pg_authid.h"
 #include "catalog/pg_db_role_setting.h"
 #include "catalog/pg_database.h"
@@ -35,6 +36,7 @@
 #include "utils/syscache.h"
 #include "miscadmin.h"
 #include "utils/builtins.h"
+#include "utils/catcache.h"
 
 static int	FindMatchingParam(List *params, const char *name);
 static Node *TransformParamRef(ParseState *pstate, ParamRef *pref);
@@ -937,8 +939,11 @@ check_babelfish_droprole_restrictions(char *role)
  * 	This function check the underlying assumption on the membership chain instead
  * 	sysadmin <-- dbo* <--- db_owner* <--- users/roles
  *
+ * A Babelfish login is a member of babelfish created user/role i.e., dbo, guest or user created user.
+ *
  * actual dbo and db_owner name varies across different babelfish logical databases
  */
+
 static bool
 is_babelfish_role(const char *role)
 {
@@ -949,13 +954,16 @@ is_babelfish_role(const char *role)
 	Oid			bbf_msdb_guest_oid;
 	Oid			securityadmin;
 	Oid			dbcreator;
+	CatCList		*memlist;
+	int			i;
+	bool			is_babelfish_login = false;
 
 	sysadmin_oid = get_role_oid(BABELFISH_SYSADMIN, true);	/* missing OK */
 	role_oid = get_role_oid(role, true);	/* missing OK */
 	securityadmin = get_role_oid(BABELFISH_SECURITYADMIN, true);  /* missing OK */
 	dbcreator = get_role_oid(BABELFISH_DBCREATOR, true);  /* missing OK */
 
-	if (!OidIsValid(sysadmin_oid) || !OidIsValid(role_oid) 
+	if (!OidIsValid(sysadmin_oid) || !OidIsValid(role_oid)
 			|| !OidIsValid(securityadmin) || !OidIsValid(dbcreator))
 		return false;
 
@@ -965,18 +973,37 @@ is_babelfish_role(const char *role)
 		pg_strcasecmp(role, BABELFISH_ROLE_ADMIN) == 0) /* check if it is bbf_role_admin */
 		return true;
 
+	/* Most of the Babelfish logins would be a member of one of these guest roles.*/
 	bbf_master_guest_oid = get_role_oid("master_guest", true);
 	bbf_tempdb_guest_oid = get_role_oid("tempdb_guest", true);
 	bbf_msdb_guest_oid = get_role_oid("msdb_guest", true);
-	if (OidIsValid(bbf_master_guest_oid)
-		&& OidIsValid(bbf_tempdb_guest_oid)
-		&& OidIsValid(bbf_msdb_guest_oid)
-		&& (is_member_of_role(role_oid, bbf_master_guest_oid)
-		|| is_member_of_role(role_oid, bbf_tempdb_guest_oid)
-		|| is_member_of_role(role_oid, bbf_msdb_guest_oid)))
+	if ((OidIsValid(bbf_master_guest_oid) && is_member_of_role(role_oid, bbf_master_guest_oid))
+		|| (OidIsValid(bbf_tempdb_guest_oid) && is_member_of_role(role_oid, bbf_tempdb_guest_oid))
+		|| (OidIsValid(bbf_msdb_guest_oid) && is_member_of_role(role_oid, bbf_msdb_guest_oid)))
 		return true;
 
-	return false;
+	/* Check if it's a Babelfish login, if it's not a member of any guest role.*/
+	memlist = SearchSysCacheList1(AUTHMEMMEMROLE,
+								ObjectIdGetDatum(role_oid));
+	for (i = 0; i < memlist->n_members; i++)
+	{
+		HeapTuple   tup = &memlist->members[i]->tuple;
+		Form_pg_auth_members form = (Form_pg_auth_members) GETSTRUCT(tup);
+		Oid         parent_role_oid = form->roleid;
+		const char *parent_role_name = GetUserNameFromId(parent_role_oid, false);
+
+		/* Check if the parent role is a Babelfish role */
+		if (is_member_of_role(sysadmin_oid, parent_role_oid) ||
+			is_member_of_role(securityadmin, parent_role_oid) ||
+			is_member_of_role(dbcreator, parent_role_oid) ||
+			pg_strcasecmp(parent_role_name, BABELFISH_ROLE_ADMIN) == 0)
+		{
+			is_babelfish_login = true;
+			break;
+		}
+	}
+	ReleaseSysCacheList(memlist);
+	return is_babelfish_login;
 }
 
 /*

diff --git a/test/JDBC/expected/permission_restrictions_from_pg.out b/test/JDBC/expected/permission_restrictions_from_pg.out
@@ -8,6 +8,32 @@ go
 create user permission_restrictions_tsql_user1 for login permission_restrictions_tsql_login1;
 go
 
+-- create user for a login in master, tempdb and msdb
+create login permission_restrictions_tsql_login2 with password = '123';
+go
+
+create user permission_restrictions_tsql_login2;
+go
+
+use tempdb;
+go
+
+create user permission_restrictions_tsql_login2;
+go
+
+use msdb;
+go
+
+create user permission_restrictions_tsql_login2;
+go
+
+use master
+go
+
+-- create a login with no mapped user
+create login permission_restrictions_tsql_login3 with password = '123';
+go
+
 -- psql
 create user permission_restrictions_psql_user with password '123';
 go
@@ -76,6 +102,14 @@ go
     Server SQLState: 42501)~~
 
 
+grant permission_restrictions_tsql_login2 to permission_restrictions_psql_user;
+go
+~~ERROR (Code: 0)~~
+
+~~ERROR (Message: ERROR: Babelfish-created logins/users/roles cannot be altered outside of a Babelfish session
+    Server SQLState: 42501)~~
+
+
 grant master_db_datareader to master_db_datareader;
 go
 ~~ERROR (Code: 0)~~
@@ -93,6 +127,22 @@ go
     Server SQLState: 42501)~~
 
 
+drop user permission_restrictions_tsql_login2;
+go
+~~ERROR (Code: 0)~~
+
+~~ERROR (Message: ERROR: Babelfish-created logins/users/roles cannot be dropped or altered outside of a Babelfish session
+    Server SQLState: 42501)~~
+
+
+drop user permission_restrictions_tsql_login3;
+go
+~~ERROR (Code: 0)~~
+
+~~ERROR (Message: ERROR: Babelfish-created logins/users/roles cannot be dropped or altered outside of a Babelfish session
+    Server SQLState: 42501)~~
+
+
 drop user master_permission_restrictions_tsql_user1;
 go
 ~~ERROR (Code: 0)~~
@@ -198,6 +248,14 @@ go
     Server SQLState: 42501)~~
 
 
+grant permission_restrictions_tsql_login2 to permission_restrictions_psql_user;
+go
+~~ERROR (Code: 0)~~
+
+~~ERROR (Message: ERROR: Babelfish-created logins/users/roles cannot be altered outside of a Babelfish session
+    Server SQLState: 42501)~~
+
+
 grant master_db_datareader to master_db_datareader;
 go
 ~~ERROR (Code: 0)~~
@@ -215,6 +273,22 @@ go
     Server SQLState: 42501)~~
 
 
+drop user permission_restrictions_tsql_login2;
+go
+~~ERROR (Code: 0)~~
+
+~~ERROR (Message: ERROR: Babelfish-created logins/users/roles cannot be dropped or altered outside of a Babelfish session
+    Server SQLState: 42501)~~
+
+
+drop user permission_restrictions_tsql_login3;
+go
+~~ERROR (Code: 0)~~
+
+~~ERROR (Message: ERROR: Babelfish-created logins/users/roles cannot be dropped or altered outside of a Babelfish session
+    Server SQLState: 42501)~~
+
+
 drop user master_permission_restrictions_tsql_user1;
 go
 ~~ERROR (Code: 0)~~
@@ -283,6 +357,14 @@ bool
 ~~END~~
 
 
+SELECT pg_terminate_backend(pid) FROM pg_stat_get_activity(NULL)
+WHERE sys.suser_name(usesysid) = 'permission_restrictions_tsql_login2' AND backend_type = 'client backend' AND usesysid IS NOT NULL;
+GO
+~~START~~
+bool
+~~END~~
+
+
 select pg_sleep(1);
 GO
 ~~START~~
@@ -295,8 +377,29 @@ void
 drop user permission_restrictions_tsql_user1;
 go
 
+drop user permission_restrictions_tsql_login2;
+go
+
+use tempdb
+go
+
+drop user permission_restrictions_tsql_login2;
+go
+
+use msdb
+go
+
+drop user permission_restrictions_tsql_login2;
+go
+
 drop login permission_restrictions_tsql_login1;
 go
 
+drop login permission_restrictions_tsql_login2;
+go
+
 drop login permission_restrictions_tsql_login
 go
+
+drop login permission_restrictions_tsql_login3;
+go
diff --git a/test/JDBC/input/permission_restrictions_from_pg.mix b/test/JDBC/input/permission_restrictions_from_pg.mix
@@ -8,6 +8,32 @@ go
 create user permission_restrictions_tsql_user1 for login permission_restrictions_tsql_login1;
 go
 
+-- create user for a login in master, tempdb and msdb
+create login permission_restrictions_tsql_login2 with password = '123';
+go
+
+create user permission_restrictions_tsql_login2;
+go
+
+use tempdb;
+go
+
+create user permission_restrictions_tsql_login2;
+go
+
+use msdb;
+go
+
+create user permission_restrictions_tsql_login2;
+go
+
+use master
+go
+
+-- create a login with no mapped user
+create login permission_restrictions_tsql_login3 with password = '123';
+go
+
 -- psql
 create user permission_restrictions_psql_user with password '123';
 go
@@ -39,13 +65,22 @@ go
 grant permission_restrictions_tsql_login1 to permission_restrictions_psql_user;
 go
 
+grant permission_restrictions_tsql_login2 to permission_restrictions_psql_user;
+go
+
 grant master_db_datareader to master_db_datareader;
 go
 
 -- dropping of BBF users/logins should not be allowed
 drop user permission_restrictions_tsql_login1;
 go
 
+drop user permission_restrictions_tsql_login2;
+go
+
+drop user permission_restrictions_tsql_login3;
+go
+
 drop user master_permission_restrictions_tsql_user1;
 go
 
@@ -97,13 +132,22 @@ go
 grant permission_restrictions_tsql_login1 to permission_restrictions_psql_user;
 go
 
+grant permission_restrictions_tsql_login2 to permission_restrictions_psql_user;
+go
+
 grant master_db_datareader to master_db_datareader;
 go
 
 -- dropping of BBF users/logins should not be allowed
 drop user permission_restrictions_tsql_login1;
 go
 
+drop user permission_restrictions_tsql_login2;
+go
+
+drop user permission_restrictions_tsql_login3;
+go
+
 drop user master_permission_restrictions_tsql_user1;
 go
 
@@ -147,15 +191,40 @@ SELECT pg_terminate_backend(pid) FROM pg_stat_get_activity(NULL)
 WHERE sys.suser_name(usesysid) = 'permission_restrictions_tsql_login1' AND backend_type = 'client backend' AND usesysid IS NOT NULL;
 GO
 
+SELECT pg_terminate_backend(pid) FROM pg_stat_get_activity(NULL)
+WHERE sys.suser_name(usesysid) = 'permission_restrictions_tsql_login2' AND backend_type = 'client backend' AND usesysid IS NOT NULL;
+GO
+
 select pg_sleep(1);
 GO
 
 -- tsql
 drop user permission_restrictions_tsql_user1;
 go
 
+drop user permission_restrictions_tsql_login2;
+go
+
+use tempdb
+go
+
+drop user permission_restrictions_tsql_login2;
+go
+
+use msdb
+go
+
+drop user permission_restrictions_tsql_login2;
+go
+
 drop login permission_restrictions_tsql_login1;
 go
 
+drop login permission_restrictions_tsql_login2;
+go
+
 drop login permission_restrictions_tsql_login
 go
+
+drop login permission_restrictions_tsql_login3;
+go