[DPE-6898] User->databases pg_hba rules (#885)

* Restrict each user to their allowed databases

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Fix unit tests

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Fix sync users on replicas

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Fix unit test

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Add default landscape user permission

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Increase sleep time in pg_hba test, fix user->database mapping for upgrade from stable and skip event trigger function code when not a superuser

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Improve users list check

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Fix raft reinitialisation in tests

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Decrease the amount of API calls by one

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Check users list directly

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Tweak test fast interval

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Improvements to avoid replica restart while syncing from primary

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Fix linting

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

---------

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

Author: marceloneppel

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -35,7 +35,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 52
+LIBPATCH = 54
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -113,6 +113,10 @@ class PostgreSQLGetPostgreSQLVersionError(Exception):
     """Exception raised when retrieving PostgreSQL version fails."""
 
 
+class PostgreSQLListAccessibleDatabasesForUserError(Exception):
+    """Exception raised when retrieving the accessible databases for a user fails."""
+
+
 class PostgreSQLListGroupsError(Exception):
     """Exception raised when retrieving PostgreSQL groups list fails."""
 
@@ -190,6 +194,11 @@ def create_access_groups(self) -> None:
         try:
             with self._connect_to_database() as connection, connection.cursor() as cursor:
                 for group in ACCESS_GROUPS:
+                    cursor.execute(
+                        SQL("SELECT TRUE FROM pg_roles WHERE rolname={};").format(Literal(group))
+                    )
+                    if cursor.fetchone() is not None:
+                        continue
                     cursor.execute(
                         SQL("CREATE ROLE {} NOLOGIN;").format(
                             Identifier(group),
@@ -622,15 +631,22 @@ def is_tls_enabled(self, check_current_host: bool = False) -> bool:
             # Connection errors happen when PostgreSQL has not started yet.
             return False
 
-    def list_access_groups(self) -> Set[str]:
+    def list_access_groups(self, current_host=False) -> Set[str]:
         """Returns the list of PostgreSQL database access groups.
 
+        Args:
+            current_host: whether to check the current host
+                instead of the primary host.
+
         Returns:
             List of PostgreSQL database access groups.
         """
         connection = None
+        host = self.current_host if current_host else None
         try:
-            with self._connect_to_database() as connection, connection.cursor() as cursor:
+            with self._connect_to_database(
+                database_host=host
+            ) as connection, connection.cursor() as cursor:
                 cursor.execute(
                     "SELECT groname FROM pg_catalog.pg_group WHERE groname LIKE '%_access';"
                 )
@@ -643,16 +659,69 @@ def list_access_groups(self) -> Set[str]:
             if connection is not None:
                 connection.close()
 
-    def list_users(self) -> Set[str]:
+    def list_accessible_databases_for_user(self, user: str, current_host=False) -> Set[str]:
+        """Returns the list of accessible databases for a specific user.
+
+        Args:
+            user: the user to check.
+            current_host: whether to check the current host
+                instead of the primary host.
+
+        Returns:
+            List of accessible database (the ones where
+                the user has the CONNECT privilege).
+        """
+        connection = None
+        host = self.current_host if current_host else None
+        try:
+            with self._connect_to_database(
+                database_host=host
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL(
+                        "SELECT TRUE FROM pg_catalog.pg_user WHERE usename = {} AND usesuper;"
+                    ).format(Literal(user))
+                )
+                if cursor.fetchone() is not None:
+                    return {"all"}
+                cursor.execute(
+                    SQL(
+                        "SELECT datname FROM pg_catalog.pg_database WHERE has_database_privilege({}, datname, 'CONNECT') AND NOT datistemplate;"
+                    ).format(Literal(user))
+                )
+                databases = cursor.fetchall()
+                return {database[0] for database in databases}
+        except psycopg2.Error as e:
+            logger.error(f"Failed to list accessible databases for user {user}: {e}")
+            raise PostgreSQLListAccessibleDatabasesForUserError() from e
+        finally:
+            if connection is not None:
+                connection.close()
+
+    def list_users(self, group: Optional[str] = None, current_host=False) -> Set[str]:
         """Returns the list of PostgreSQL database users.
 
+        Args:
+            group: optional group to filter the users.
+            current_host: whether to check the current host
+                instead of the primary host.
+
         Returns:
             List of PostgreSQL database users.
         """
         connection = None
+        host = self.current_host if current_host else None
         try:
-            with self._connect_to_database() as connection, connection.cursor() as cursor:
-                cursor.execute("SELECT usename FROM pg_catalog.pg_user;")
+            with self._connect_to_database(
+                database_host=host
+            ) as connection, connection.cursor() as cursor:
+                if group:
+                    query = SQL(
+                        "SELECT usename FROM (SELECT UNNEST(grolist) AS user_id FROM pg_catalog.pg_group WHERE groname = {}) AS g JOIN pg_catalog.pg_user AS u ON g.user_id = u.usesysid;"
+                    ).format(Literal(group))
+                else:
+                    query = "SELECT usename FROM pg_catalog.pg_user;"
+                cursor.execute(query)
                 usernames = cursor.fetchall()
                 return {username[0] for username in usernames}
         except psycopg2.Error as e:
@@ -662,19 +731,27 @@ def list_users(self) -> Set[str]:
             if connection is not None:
                 connection.close()
 
-    def list_users_from_relation(self) -> Set[str]:
+    def list_users_from_relation(self, current_host=False) -> Set[str]:
         """Returns the list of PostgreSQL database users that were created by a relation.
 
+        Args:
+            current_host: whether to check the current host
+                instead of the primary host.
+
         Returns:
             List of PostgreSQL database users.
         """
         connection = None
+        host = self.current_host if current_host else None
         try:
-            with self._connect_to_database() as connection, connection.cursor() as cursor:
+            with self._connect_to_database(
+                database_host=host
+            ) as connection, connection.cursor() as cursor:
                 cursor.execute(
                     "SELECT usename "
                     "FROM pg_catalog.pg_user "
-                    "WHERE usename LIKE 'relation_id_%' OR usename LIKE 'relation-%';"
+                    "WHERE usename LIKE 'relation_id_%' OR usename LIKE 'relation-%' "
+                    "OR usename LIKE 'pgbouncer_auth_relation_id_%' OR usename LIKE '%_user_%_%';"
                 )
                 usernames = cursor.fetchall()
                 return {username[0] for username in usernames}
@@ -700,42 +777,108 @@ def list_valid_privileges_and_roles(self) -> Tuple[Set[str], Set[str]]:
                 "superuser",
             }, {role[0] for role in cursor.fetchall() if role[0]}
 
-    def set_up_database(self, temp_location: Optional[str] = None) -> None:
+    def set_up_database(self) -> None:
         """Set up postgres database with the right permissions."""
         connection = None
-        cursor = None
         try:
-            connection = self._connect_to_database()
-            cursor = connection.cursor()
-
-            if temp_location is not None:
-                cursor.execute("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';")
+            with self._connect_to_database(
+                database="template1"
+            ) as connection, connection.cursor() as cursor:
+                # Create database function and event trigger to identify users created by PgBouncer.
+                cursor.execute(
+                    "SELECT TRUE FROM pg_event_trigger WHERE evtname = 'update_pg_hba_on_create_schema';"
+                )
                 if cursor.fetchone() is None:
-                    cursor.execute(f"CREATE TABLESPACE temp LOCATION '{temp_location}';")
-                    cursor.execute("GRANT CREATE ON TABLESPACE temp TO public;")
-
-            cursor.execute("SELECT TRUE FROM pg_roles WHERE rolname='admin';")
-            if cursor.fetchone() is None:
-                # Allow access to the postgres database only to the system users.
-                cursor.execute("REVOKE ALL PRIVILEGES ON DATABASE postgres FROM PUBLIC;")
-                cursor.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC;")
-                for user in self.system_users:
-                    cursor.execute(
-                        SQL("GRANT ALL PRIVILEGES ON DATABASE postgres TO {};").format(
-                            Identifier(user)
+                    cursor.execute("""
+CREATE OR REPLACE FUNCTION update_pg_hba()
+    RETURNS event_trigger
+    LANGUAGE plpgsql
+    AS $$
+        DECLARE
+          hba_file TEXT;
+          copy_command TEXT;
+          connection_type TEXT;
+          rec record;
+          insert_value TEXT;
+          changes INTEGER = 0;
+        BEGIN
+          -- Don't execute on replicas.
+          IF NOT pg_is_in_recovery() THEN
+            -- Load the current authorisation rules.
+            DROP TABLE IF EXISTS pg_hba;
+            CREATE TEMPORARY TABLE pg_hba (lines TEXT);
+            SELECT setting INTO hba_file FROM pg_settings WHERE name = 'hba_file';
+            IF hba_file IS NOT NULL THEN
+                copy_command='COPY pg_hba FROM ''' || hba_file || '''' ;
+                EXECUTE copy_command;
+                -- Build a list of the relation users and the databases they can access.
+                DROP TABLE IF EXISTS relation_users;
+                CREATE TEMPORARY TABLE relation_users AS
+                  SELECT t.user, STRING_AGG(DISTINCT t.database, ',') AS databases FROM( SELECT u.usename AS user, CASE WHEN u.usesuper THEN 'all' ELSE d.datname END AS database FROM ( SELECT usename, usesuper FROM pg_catalog.pg_user WHERE usename NOT IN ('backup', 'monitoring', 'operator', 'postgres', 'replication', 'rewind')) AS u JOIN ( SELECT datname FROM pg_catalog.pg_database WHERE NOT datistemplate ) AS d ON has_database_privilege(u.usename, d.datname, 'CONNECT') ) AS t GROUP BY 1;
+                IF (SELECT COUNT(lines) FROM pg_hba WHERE lines LIKE 'hostssl %') > 0 THEN
+                  connection_type := 'hostssl';
+                ELSE
+                  connection_type := 'host';
+                END IF;
+                -- Add the new users to the pg_hba file.
+                FOR rec IN SELECT * FROM relation_users
+                LOOP
+                  insert_value := connection_type || ' ' || rec.databases || ' ' || rec.user || ' 0.0.0.0/0 md5';
+                  IF (SELECT COUNT(lines) FROM pg_hba WHERE lines = insert_value) = 0 THEN
+                    INSERT INTO pg_hba (lines) VALUES (insert_value);
+                    changes := changes + 1;
+                  END IF;
+                END LOOP;
+                -- Remove users that don't exist anymore from the pg_hba file.
+                FOR rec IN SELECT h.lines FROM pg_hba AS h LEFT JOIN relation_users AS r ON SPLIT_PART(h.lines, ' ', 3) = r.user WHERE r.user IS NULL AND (SPLIT_PART(h.lines, ' ', 3) LIKE 'relation_id_%' OR SPLIT_PART(h.lines, ' ', 3) LIKE 'pgbouncer_auth_relation_id_%' OR SPLIT_PART(h.lines, ' ', 3) LIKE '%_user_%_%')
+                LOOP
+                  DELETE FROM pg_hba WHERE lines = rec.lines;
+                  changes := changes + 1;
+                END LOOP;
+                -- Apply the changes to the pg_hba file.
+                IF changes > 0 THEN
+                  copy_command='COPY pg_hba TO ''' || hba_file || '''' ;
+                  EXECUTE copy_command;
+                  PERFORM pg_reload_conf();
+                END IF;
+            END IF;
+          END IF;
+        END;
+    $$;
+                    """)
+                    cursor.execute("""
+CREATE EVENT TRIGGER update_pg_hba_on_create_schema
+    ON ddl_command_end
+    WHEN TAG IN ('CREATE SCHEMA')
+    EXECUTE FUNCTION update_pg_hba();
+                    """)
+                    cursor.execute("""
+CREATE EVENT TRIGGER update_pg_hba_on_drop_schema
+    ON ddl_command_end
+    WHEN TAG IN ('DROP SCHEMA')
+    EXECUTE FUNCTION update_pg_hba();
+                    """)
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                cursor.execute("SELECT TRUE FROM pg_roles WHERE rolname='admin';")
+                if cursor.fetchone() is None:
+                    # Allow access to the postgres database only to the system users.
+                    cursor.execute("REVOKE ALL PRIVILEGES ON DATABASE postgres FROM PUBLIC;")
+                    cursor.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC;")
+                    for user in self.system_users:
+                        cursor.execute(
+                            SQL("GRANT ALL PRIVILEGES ON DATABASE postgres TO {};").format(
+                                Identifier(user)
+                            )
                         )
+                    self.create_user(
+                        PERMISSIONS_GROUP_ADMIN,
+                        extra_user_roles=["pg_read_all_data", "pg_write_all_data"],
                     )
-                self.create_user(
-                    PERMISSIONS_GROUP_ADMIN,
-                    extra_user_roles=["pg_read_all_data", "pg_write_all_data"],
-                )
-                cursor.execute("GRANT CONNECT ON DATABASE postgres TO admin;")
+                    cursor.execute("GRANT CONNECT ON DATABASE postgres TO admin;")
         except psycopg2.Error as e:
             logger.error(f"Failed to set up databases: {e}")
             raise PostgreSQLDatabasesSetupError() from e
         finally:
-            if cursor is not None:
-                cursor.close()
             if connection is not None:
                 connection.close()
 

scripts/cluster_topology_observer.py

@@ -22,11 +22,51 @@ class UnreachableUnitsError(Exception):
     """Cannot reach any known cluster member."""
 
 
-def dispatch(run_cmd, unit, charm_dir):
-    """Use the input juju-run command to dispatch a :class:`ClusterTopologyChangeEvent`."""
-    dispatch_sub_cmd = "JUJU_DISPATCH_PATH=hooks/cluster_topology_change {}/dispatch"
+def check_for_authorisation_rules_changes(run_cmd, unit, charm_dir, previous_authorisation_rules):
+    """Check for changes in the authorisation rules.
+
+    If changes are detected, dispatch an event to handle them.
+    """
+    # Read contents from the pg_hba.conf file.
+    with open("/var/snap/charmed-postgresql/common/var/lib/postgresql/pg_hba.conf") as file:
+        current_authorisation_rules = file.read()
+        current_authorisation_rules = [
+            line for line in current_authorisation_rules.splitlines() if not line.startswith("#")
+        ]
+    # If it's the first time the authorisation rules were retrieved, then store it and use
+    # it for subsequent checks.
+    if not previous_authorisation_rules:
+        previous_authorisation_rules = current_authorisation_rules
+    # If the authorisation rules changed, dispatch a charm event to handle this change.
+    elif current_authorisation_rules != previous_authorisation_rules:
+        previous_authorisation_rules = current_authorisation_rules
+        dispatch(run_cmd, unit, charm_dir, "authorisation_rules_change")
+    return previous_authorisation_rules
+
+
+def check_for_cluster_topology_changes(
+    run_cmd, unit, charm_dir, previous_cluster_topology, current_cluster_topology
+):
+    """Check for changes in the cluster topology.
+
+    If changes are detected, dispatch an event to handle them.
+    """
+    # If it's the first time the cluster topology was retrieved, then store it and use
+    # it for subsequent checks.
+    if not previous_cluster_topology:
+        previous_cluster_topology = current_cluster_topology
+    # If the cluster topology changed, dispatch a charm event to handle this change.
+    elif current_cluster_topology != previous_cluster_topology:
+        previous_cluster_topology = current_cluster_topology
+        dispatch(run_cmd, unit, charm_dir, "cluster_topology_change")
+    return previous_cluster_topology
+
+
+def dispatch(run_cmd, unit, charm_dir, custom_event):
+    """Use the input juju-run command to dispatch a custom event."""
+    dispatch_sub_cmd = "JUJU_DISPATCH_PATH=hooks/{} {}/dispatch"
     # Input is generated by the charm
-    subprocess.run([run_cmd, "-u", unit, dispatch_sub_cmd.format(charm_dir)])  # noqa: S603
+    subprocess.run([run_cmd, "-u", unit, dispatch_sub_cmd.format(custom_event, charm_dir)])  # noqa: S603
 
 
 def main():
@@ -37,6 +77,7 @@ def main():
     patroni_urls, run_cmd, unit, charm_dir = sys.argv[1:]
 
     previous_cluster_topology = {}
+    previous_authorisation_rules = []
     urls = [urljoin(url, PATRONI_CLUSTER_STATUS_ENDPOINT) for url in patroni_urls.split(",")]
     member_name = unit.replace("/", "-")
     while True:
@@ -63,23 +104,27 @@ def main():
             raise UnreachableUnitsError("Unable to reach cluster members")
         current_cluster_topology = {}
         urls = []
+        is_primary = False
         for member in cluster_status["members"]:
             current_cluster_topology[member["name"]] = member["role"]
             member_url = urljoin(member["api_url"], PATRONI_CLUSTER_STATUS_ENDPOINT)
             # Call the current unit first
             if member["name"] == member_name:
                 urls.insert(0, member_url)
+                # Check if the current member is the primary.
+                if member["role"] == "leader":
+                    is_primary = True
             else:
                 urls.append(member_url)
 
-        # If it's the first time the cluster topology was retrieved, then store it and use
-        # it for subsequent checks.
-        if not previous_cluster_topology:
-            previous_cluster_topology = current_cluster_topology
-        # If the cluster topology changed, dispatch a charm event to handle this change.
-        elif current_cluster_topology != previous_cluster_topology:
-            previous_cluster_topology = current_cluster_topology
-            dispatch(run_cmd, unit, charm_dir)
+        previous_cluster_topology = check_for_cluster_topology_changes(
+            run_cmd, unit, charm_dir, previous_cluster_topology, current_cluster_topology
+        )
+
+        if is_primary:
+            previous_authorisation_rules = check_for_authorisation_rules_changes(
+                run_cmd, unit, charm_dir, previous_authorisation_rules
+            )
 
         # Wait some time before checking again for a cluster topology change.
         sleep(30)