diff --git a/docs/en_US/oauth2.rst b/docs/en_US/oauth2.rst index a8cc372b795..1eaa9235a17 100644 --- a/docs/en_US/oauth2.rst +++ b/docs/en_US/oauth2.rst @@ -430,3 +430,190 @@ For providers that don't support OIDC discovery, configure all endpoints manuall }] In this mode, user identity is retrieved only from the userinfo endpoint. + +.. _oauth2_postgresql: + +Using a pgAdmin OAuth token with PostgreSQL +------------------------------------------- + +Users signed in to pgAdmin through an OIDC provider can use +``oauth_pgadmin_token_mode`` to supply an access token for PostgreSQL OAuth +authentication. This uses an OAuth access token, not an OIDC ID token. + +Prerequisites +~~~~~~~~~~~~~ + +* Configure OIDC login for pgAdmin as described on this page and sign in + through that OIDC provider. +* Configure PostgreSQL OAuth authentication, including a token validator, + the expected issuer and scopes, and authorization for the database role. +* Use PostgreSQL 18 or later. The libpq loaded by pgAdmin must be version + 18.x, with a supported Psycopg implementation as described below. + +Configure the connection +~~~~~~~~~~~~~~~~~~~~~~~~ + +In the *Parameters* tab of the Server dialog, add *OAuth pgAdmin token mode* +(``oauth_pgadmin_token_mode``). Select ``direct`` or ``exchange`` to enable +this integration. + +.. note:: + + For shared servers, each user must explicitly enable + ``oauth_pgadmin_token_mode`` in their own server settings. + The owner's selection is not inherited. + +Both ``direct`` and ``exchange`` require ``oauth_client_id`` and +``oauth_issuer``. The issuer must match the PostgreSQL OAuth configuration +and the issuer used for the pgAdmin login token. + +For example: + +.. code-block:: text + + oauth_pgadmin_token_mode = exchange + oauth_client_id = postgres-cluster-01 + oauth_issuer = https://identity.example.com/realms/postgres + +Replace the example values with those configured for your deployment. Set +the *Username* field to the database role authorized by the server's OAuth +configuration. + +Use a trusted database endpoint. For TCP connections, configure certificate +and hostname verification, for example with ``sslmode=verify-full`` and a +trusted root certificate. Database TLS settings are separate from those for +the OIDC provider's HTTPS endpoints. + +Direct mode +~~~~~~~~~~~ + +In ``direct`` mode, pgAdmin forwards the access token from the current user's +login session to PostgreSQL. The database validator must accept its issuer, +audience and scopes and authorize the requested database role. The audience +depends on the OIDC provider's configuration; it is not necessarily the +pgAdmin OAuth client identifier. + +.. warning:: + + The database server receives the login access token, which its operator + may be able to reuse wherever that token is accepted. Use direct mode + only with trusted servers and a policy explicitly permitting this reuse. + Prefer exchange mode when the OIDC provider can issue a token restricted + to the target database. + +Exchange mode +~~~~~~~~~~~~~ + +In ``exchange`` mode, pgAdmin exchanges the login access token and supplies +the resulting token to PostgreSQL. If exchange fails, the connection attempt +fails; pgAdmin does not fall back to direct mode. + +The token endpoint belongs to the OIDC provider used for pgAdmin login. +``OAUTH2_TOKEN_URL`` takes precedence over the ``token_endpoint`` in the +OIDC provider's discovery metadata. HTTPS is required and token endpoint +redirects are not followed. + +.. list-table:: Token exchange request + :header-rows: 1 + :widths: 30 70 + + * - Parameter + - Value + * - ``grant_type`` + - ``urn:ietf:params:oauth:grant-type:token-exchange`` + * - ``audience`` + - The server connection's ``oauth_client_id``. + * - ``subject_token`` + - The pgAdmin login access token, refreshed when required and possible. + * - ``subject_token_type`` + - ``urn:ietf:params:oauth:token-type:access_token`` + * - ``scope`` + - The server connection's ``oauth_scope``, when non-empty. + Otherwise omitted; the OIDC provider determines the scopes. + +Token exchange uses the scopes explicitly configured in the server +connection's ``oauth_scope`` parameter. If it is unset or blank, pgAdmin +omits ``scope`` from the exchange request and the OIDC provider determines +the scopes. Exchange occurs before connecting to PostgreSQL, so it cannot +use scopes advertised by the database server. Configure ``oauth_scope`` +when the OIDC provider's defaults do not meet the database's requirements. + +pgAdmin authenticates the exchange request using its registered login +OIDC client's credentials. This integration supports +``client_secret_basic`` and ``client_secret_post``. A client secret is +required; unauthenticated exchange is not supported. + +.. important:: + + The OIDC provider must authorize pgAdmin's login client to exchange + tokens for the requested audience and issue the intended permissions. + The PostgreSQL validator must enforce that audience and those + permissions. + +The OIDC provider must return a non-empty ``access_token`` with a +``token_type`` of ``Bearer``. Exchanged tokens are not cached or stored in the +pgAdmin login session. Each connection attempt using this mode performs a +new exchange. + +TLS verification +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Both direct and exchange modes require TLS certificate verification for +the login OIDC provider. Enable ``OAUTH2_SSL_CERT_VERIFICATION`` in +the ``OAUTH2_CONFIG``. + +For a private certificate authority, configure a trusted CA bundle for the +OAuth client. The exchange request reuses the login OAuth client's +verification setting, including a configured CA-bundle path. + +Supported client implementations +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +These requirements apply to the libpq loaded in the pgAdmin process: + +.. list-table:: Psycopg implementation requirements + :header-rows: 1 + :widths: 20 80 + + * - Implementation + - Requirements + * - ``python`` + - Reuses Psycopg's libpq handle. Requires a discoverable libpq 18.x + installation compatible with pgAdmin's Python runtime. + * - ``c`` + - Linux only. Requires libpq 18.x and an unambiguous loaded libpq + candidate. Missing or ambiguous candidates are rejected. + * - ``binary`` + - Not supported by this integration, even with bundled libpq 18.x. + +To select the Python implementation, set ``PSYCOPG_IMPL=python`` in the +pgAdmin process environment and restart the backend. This affects all +connections in that process. It is not a connection parameter or a +``PGADMIN_CONFIG_`` setting and does not install or upgrade libpq. + +Windows and macOS require the Python implementation and a compatible, +discoverable libpq 18.x installation. The installed pgAdmin package must +support that configuration. In containers, configure the environment and +libraries inside the container. + +Token refresh and reconnects +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +When expiration information is available, pgAdmin attempts to refresh the +login access token if it has expired or will expire within 30 seconds. A +successful refresh updates the login session. If the refresh token is +unavailable or refreshing fails, sign in again before retrying the connection. + +Token retrieval, refresh and exchange require an authenticated pgAdmin +request. Reconnects made through such requests are supported. Independent +background connection attempts without the user's request context are not; +initiate reconnection through pgAdmin. Existing database connections do not +require a new token exchange for each query. + +Automatic token refresh is not supported by this integration when the +login OIDC provider uses +``OAUTH2_CLIENT_AUTH_METHOD="workload_identity"``. Direct mode can use +the login access token until it expires or enters the 30-second refresh +window. Sign in to pgAdmin again before retrying the connection. +Exchange mode requires client-secret authentication and does not support +workload-identity clients. \ No newline at end of file diff --git a/docs/en_US/server_dialog.rst b/docs/en_US/server_dialog.rst index bab8f2effe1..4341cece407 100644 --- a/docs/en_US/server_dialog.rst +++ b/docs/en_US/server_dialog.rst @@ -156,6 +156,64 @@ icon that is located to the right of each of the following fields. means that administrators of the server may be able to access those files; appropriate caution should be taken before choosing to use this feature. +Use the *OAuth pgAdmin token mode* parameter +(``oauth_pgadmin_token_mode``) to choose how pgAdmin supplies an OAuth access +token when PostgreSQL requests OAuth authentication. This option is available +for PostgreSQL 18 and later. + +.. list-table:: OAuth pgAdmin token modes + :header-rows: 1 + :widths: 20 80 + + * - Value + - Behaviour + * - ``disabled`` + - Do not supply a token from the pgAdmin login session. Other connection + and authentication settings continue to apply. This is the default. + * - ``direct`` + - Supply the OAuth access token obtained when the user signed in to + pgAdmin directly to PostgreSQL. + * - ``exchange`` + - Exchange the user's pgAdmin login access token at the OIDC provider's + token endpoint and supply the returned token to PostgreSQL. Use this + mode when the OIDC provider is configured to issue a token restricted + to the target database. + +For both ``direct`` and ``exchange``, also add the following parameters: + +* ``oauth_issuer``: the expected OAuth issuer for the PostgreSQL connection. + Use the issuer configured for the database and the corresponding pgAdmin + login OIDC provider. +* ``oauth_client_id``: the OAuth client identifier configured for the database + connection. In ``exchange`` mode, pgAdmin sends this value as the requested + ``audience`` and authenticates using the pgAdmin login OIDC client's + credentials. + +For ``exchange`` mode, optionally set ``oauth_scope`` to a space-separated +list of scopes to request. If unset or blank, the OIDC provider determines +the scopes. The exchange occurs before connecting to PostgreSQL and cannot +use scopes advertised by the database server. + +The user must sign in to pgAdmin using the configured OIDC provider. +The PostgreSQL server must be configured to validate the token and authorize +the requested database role. These parameters do not configure OAuth +authentication on the PostgreSQL server. + +.. warning:: + + ``direct`` mode sends the access token obtained during pgAdmin login to + the PostgreSQL server. The server operator may be able to reuse that token + at other services that accept it. Use this mode only with trusted database + servers and a deliberately compatible token audience and scope policy. + Prefer ``exchange`` with a database-specific token where supported. + +.. note:: + + This integration requires the libpq loaded by pgAdmin to be version 18.x, + independently of the PostgreSQL server version. Supported Psycopg + implementations, OIDC provider requirements, TLS settings and reconnect + limitations are described in :ref:`oauth2_postgresql`. + Click the *SSH Tunnel* tab to continue. .. image:: images/server_ssh_tunnel.png diff --git a/web/pgadmin/authenticate/oauth2.py b/web/pgadmin/authenticate/oauth2.py index eadbfeb647f..446313431e8 100644 --- a/web/pgadmin/authenticate/oauth2.py +++ b/web/pgadmin/authenticate/oauth2.py @@ -812,6 +812,8 @@ def get_user_profile(self): self.oauth2_current_client, provider, client ) + session["oauth2_provider"] = self.oauth2_current_client + session['pass_enc_key'] = session['oauth2_token']['access_token'] # As above, the shipped config.py sets this to None, so only stash diff --git a/web/pgadmin/browser/server_groups/servers/__init__.py b/web/pgadmin/browser/server_groups/servers/__init__.py index b1d5e9cbef6..acb72cc06d2 100644 --- a/web/pgadmin/browser/server_groups/servers/__init__.py +++ b/web/pgadmin/browser/server_groups/servers/__init__.py @@ -63,6 +63,8 @@ SENSITIVE_CONN_KEYS = frozenset({ 'passfile', 'sslcert', 'sslkey', 'sslrootcert', 'sslcrl', 'sslcrldir', + # Each user must explicitly enable use of their own login token. + 'oauth_pgadmin_token_mode', }) @@ -282,12 +284,11 @@ def get_shared_server_properties(server, sharedserver): server.tunnel_identity_file = sharedserver.tunnel_identity_file server.tunnel_prompt_password = sharedserver.tunnel_prompt_password - # Override per-user connection_params keys. Use the - # SharedServer value whenever it is present, regardless of - # whether the owner's Server has the same key. Work on a copy: - # the dict may be a MutableDict still attached to the owner's - # session-bound row, and mutating it in place would mark that - # row dirty for the next commit. + # Apply each user's own connection parameters. For these keys, + # never inherit the owner's file paths or OAuth token mode. + # Work on a copy: the dict may be a MutableDict still attached to + # the owner's session-bound row, and mutating it in place would mark + # that row dirty for the next commit. s_conn = dict(getattr(server, 'connection_params', None) or {}) ss_conn = getattr(sharedserver, 'connection_params', None) or {} @@ -295,8 +296,8 @@ def get_shared_server_properties(server, sharedserver): if key in ss_conn: s_conn[key] = ss_conn[key] elif key in s_conn: - # Owner has this key but non-owner doesn't — - # remove it so the owner's path doesn't leak. + # No per-user value exists. Remove the owner's value + # so file paths and OAuth token mode are not inherited. del s_conn[key] server.connection_params = s_conn @@ -525,9 +526,9 @@ def create_shared_server(data, gid): db.session.rollback() user = User.query.filter_by(id=data.user_id).first() - # Strip owner's sensitive file paths from - # connection_params — each user should configure - # their own SSL/passfile paths. + # Exclude per-user connection parameters from the owner's + # settings. Each user must configure their own file paths + # and explicitly enable use of their pgAdmin login token. safe_conn_params = {} if data.connection_params: safe_conn_params = { @@ -1744,6 +1745,13 @@ def connect(self, gid, sid, is_qt=False, server=None): manager.update(server) conn = manager.connection() + connection_params = manager.connection_params or {} + + use_pgadmin_oauth = connection_params.get( + "oauth_pgadmin_token_mode", + "disabled", + ) in ("direct", "exchange") + # Get enc key crypt_key_present, crypt_key = get_crypt_key() if not crypt_key_present: @@ -1772,8 +1780,13 @@ def connect(self, gid, sid, is_qt=False, server=None): except Exception as e: current_app.logger.exception(e) return internal_server_error(errormsg=str(e)) - if 'password' not in data and (server.kerberos_conn is False or - server.kerberos_conn is None): + + if use_pgadmin_oauth: + # The pgAdmin OAuth bearer token is the database credential. + password = None + save_password = False + elif 'password' not in data and (server.kerberos_conn is False or + server.kerberos_conn is None): passfile_param = None if hasattr(server, 'connection_params') and \ @@ -1824,9 +1837,22 @@ def connect(self, gid, sid, is_qt=False, server=None): server_types=ServerType.types() ) except Exception as e: + error_message = getattr(e, 'message', str(e)) + + if use_pgadmin_oauth and not prompt_tunnel_password: + return make_json_response( + status=400, + success=0, + errormsg=error_message, + ) + return self.get_response_for_password( - server, 401, not server.save_password, prompt_tunnel_password, - getattr(e, 'message', str(e))) + server, + 401, + not use_pgadmin_oauth and not server.save_password, + prompt_tunnel_password, + error_message, + ) if not status: current_app.logger.error( @@ -1836,9 +1862,20 @@ def connect(self, gid, sid, is_qt=False, server=None): if errmsg.find('Ticket expired') != -1: return internal_server_error(errmsg) + if use_pgadmin_oauth and not prompt_tunnel_password: + return make_json_response( + status=400, + success=0, + errormsg=errmsg, + ) + return self.get_response_for_password( - server, 401, not server.save_password, - prompt_tunnel_password, errmsg) + server, + 401, + not use_pgadmin_oauth and not server.save_password, + prompt_tunnel_password, + errmsg, + ) else: if save_password and config.ALLOW_SAVE_PASSWORD: try: diff --git a/web/pgadmin/browser/server_groups/servers/static/js/server.ui.js b/web/pgadmin/browser/server_groups/servers/static/js/server.ui.js index e4b6618ebcd..d93aee8ab7a 100644 --- a/web/pgadmin/browser/server_groups/servers/static/js/server.ui.js +++ b/web/pgadmin/browser/server_groups/servers/static/js/server.ui.js @@ -176,6 +176,12 @@ export function getConnectionParameters() { }, { 'value': 'oauth_scope', 'label': gettext('OAuth scope'), 'vartype': 'string', 'min_server_version': '18' + }, { + 'value': 'oauth_pgadmin_token_mode', + 'label': gettext('OAuth pgAdmin token mode'), + 'vartype': 'enum', + 'enumvals': ['disabled', 'direct', 'exchange'], + 'min_server_version': '18' }]; conParams.sort(function (a, b) { diff --git a/web/pgadmin/browser/server_groups/servers/tests/test_check_ssh_mock_connect.py b/web/pgadmin/browser/server_groups/servers/tests/test_check_ssh_mock_connect.py index 1623d1cef3e..3f85b15dafd 100644 --- a/web/pgadmin/browser/server_groups/servers/tests/test_check_ssh_mock_connect.py +++ b/web/pgadmin/browser/server_groups/servers/tests/test_check_ssh_mock_connect.py @@ -14,6 +14,10 @@ import json from psycopg import OperationalError from pgadmin.utils.constants import TWO_PARAM_STRING +from contextlib import ExitStack +from types import SimpleNamespace + +import pgadmin.browser.server_groups.servers as servers_module class ServersSSHConnectTestCase(BaseTestGenerator): @@ -106,3 +110,186 @@ def __init__(self, name, id, username, use_ssh_tunnel, response = self.connect_to_server(url, self.server) self.assertIn(response.status_code, [401, 428]) + + +class ServersOAuthSSHConnectTestCase(BaseTestGenerator): + """Test database OAuth errors independently of SSH password prompts.""" + + scenarios = [ + ( + name, + dict(parameters), + ) + for name, parameters in ServersSSHConnectTestCase.scenarios + ] + + def setUp(self): + servers_utils.restore_server_fixture_on_cleanup(self) + + def _check_failure_response( + self, mode, needs_tunnel_password, raises_exception): + error_message = 'Test connection failure.' + connection_params = { + 'oauth_pgadmin_token_mode': mode, + 'oauth_client_id': 'postgres-cluster', + 'oauth_issuer': 'https://identity.example.test', + 'oauth_scope': 'openid', + } + + server = SimpleNamespace( + id=1, + name='OAuth SSH test server', + username='database-user', + user_id=1, + shared=None, + service=None, + password=None, + save_password=False, + kerberos_conn=False, + connection_params=connection_params, + use_ssh_tunnel=1 if needs_tunnel_password else 0, + tunnel_host='127.0.0.1', + tunnel_port=22, + tunnel_username='ssh-user', + tunnel_authentication=0, + tunnel_identity_file=None, + tunnel_password=( + 'test-encrypted-tunnel-password' + if needs_tunnel_password else None + ), + tunnel_keep_alive=0, + tunnel_prompt_password=0, + ) + + manager = MagicMock() + manager.password = None + manager.server_type = 'pg' + manager.sversion = 180000 + manager.connection_params = dict(connection_params) + manager.tunnel_created = False + + connection = manager.connection.return_value + connection.connected.return_value = False + + if raises_exception: + connection.connect.side_effect = OperationalError( + error_message, + ) + else: + connection.connect.return_value = ( + False, + error_message, + ) + + payload = dict(self.server) + payload.pop('tunnel_password', None) + + # Include the field to bypass the initial database-password + # prompt. None also avoids encrypting a dummy password. + payload['password'] = None + payload['save_password'] = False + + url = self.url + TWO_PARAM_STRING.format( + utils.SERVER_GROUP, + server.id, + ) + + with ExitStack() as stack: + get_server = stack.enter_context( + patch.object( + servers_module, + 'get_server', + return_value=server, + ) + ) + + # These tests exercise an owned server. + stack.enter_context( + patch.object( + servers_module, + '_is_non_owner', + return_value=False, + ) + ) + + # Isolate the failure-response logic from the rules that + # determine whether SSH credentials can be requested. + stack.enter_context( + patch.object( + servers_module.ServerNode, + 'is_prompt_tunnel_password', + return_value=needs_tunnel_password, + ) + ) + + get_driver = stack.enter_context( + patch.object(servers_module, 'get_driver') + ) + get_driver.return_value.connection_manager.return_value = manager + + stack.enter_context( + patch.object( + servers_module, + 'get_crypt_key', + return_value=(True, 'test-crypt-key'), + ) + ) + + password_response = stack.enter_context( + patch.object( + servers_module.ServerNode, + 'get_response_for_password', + autospec=True, + return_value=({'test_password_prompt': True}, 401), + ) + ) + + response = self.tester.post( + url, + data=json.dumps(payload), + content_type='html/json', + ) + + get_server.assert_called_once_with(server.id) + connection.connect.assert_called_once() + + use_pgadmin_oauth = mode in ('direct', 'exchange') + + if use_pgadmin_oauth and not needs_tunnel_password: + self.assertEqual(response.status_code, 400) + self.assertEqual( + response.get_json()['errormsg'], + error_message, + ) + password_response.assert_not_called() + else: + self.assertEqual(response.status_code, 401) + password_response.assert_called_once() + + # The trailing arguments are the same whether the helper + # is an instance method or a static method. + arguments = password_response.call_args.args + self.assertIs(arguments[-5], server) + self.assertEqual( + arguments[-4:], + ( + 401, + not use_pgadmin_oauth, + needs_tunnel_password, + error_message, + ), + ) + + def runTest(self): + for mode in ('disabled', 'direct', 'exchange'): + for needs_tunnel_password in (False, True): + for raises_exception in (False, True): + with self.subTest( + mode=mode, + needs_tunnel_password=needs_tunnel_password, + raises_exception=raises_exception): + self._check_failure_response( + mode, + needs_tunnel_password, + raises_exception, + ) diff --git a/web/pgadmin/browser/server_groups/servers/tests/test_shared_server_oauth.py b/web/pgadmin/browser/server_groups/servers/tests/test_shared_server_oauth.py new file mode 100644 index 00000000000..8d3a7f863af --- /dev/null +++ b/web/pgadmin/browser/server_groups/servers/tests/test_shared_server_oauth.py @@ -0,0 +1,128 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +"""Shared-server users must opt in to using their own OAuth login token.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from pgadmin.browser.server_groups import servers as servers_module +from pgadmin.utils.route import BaseTestGenerator + + +class SharedServerOAuthConsentTest(BaseTestGenerator): + scenarios = [ + ('New shared user does not inherit ' + mode, { + 'operation': 'create', + 'owner_mode': mode, + 'user_mode': None, + }) + for mode in ('direct', 'exchange') + ] + [ + ('Owner {0}, shared user {1}'.format(owner_mode, user_mode), { + 'operation': 'merge', + 'owner_mode': owner_mode, + 'user_mode': user_mode, + }) + for owner_mode, user_mode in ( + ('direct', None), + ('exchange', None), + ('direct', 'disabled'), + ('exchange', 'disabled'), + ('disabled', 'direct'), + ('disabled', 'exchange'), + ('direct', 'exchange'), + ('exchange', 'direct'), + ) + ] + + def setUp(self): + # Bypass the PostgreSQL connection setup. + # Use real, detached model objects and mock only persistence. + unittest.TestCase.setUp(self) + + def _create_shared_server(self, owner): + with patch.object(servers_module, 'db') as database, \ + patch.object(servers_module, 'User') as user_model, \ + patch.object(servers_module, 'current_user', + SimpleNamespace(id=2)): + user_model.query.filter_by.return_value.first.return_value = ( + SimpleNamespace(username='server-owner') + ) + + servers_module.ServerModule.create_shared_server(owner, 1) + + database.session.add.assert_called_once() + database.session.commit.assert_called_once_with() + shared = database.session.add.call_args.args[0] + + self.assertIsInstance(shared, servers_module.SharedServer) + self.assertEqual(shared.user_id, 2) + self.assertEqual(shared.osid, owner.id) + return shared + + def runTest(self): + owner = servers_module.Server( + id=1, + user_id=1, + name='OAuth shared server', + host='database.example.test', + port=5432, + maintenance_db='postgres', + shared=True, + shared_username='database-user', + connection_params={ + 'oauth_pgadmin_token_mode': self.owner_mode, + 'oauth_client_id': 'postgres-cluster', + 'oauth_issuer': 'https://identity.example.test', + 'sslmode': 'verify-full', + }, + ) + owner_params = owner.connection_params + owner_snapshot = dict(owner_params) + + if self.operation == 'create': + shared = self._create_shared_server(owner) + self.assertNotIn( + 'oauth_pgadmin_token_mode', shared.connection_params, + ) + else: + shared = servers_module.SharedServer( + osid=owner.id, + user_id=2, + connection_params=( + {'oauth_pgadmin_token_mode': self.user_mode} + if self.user_mode is not None else {} + ), + ) + + shared_snapshot = dict(shared.connection_params) + result = servers_module.ServerModule.get_shared_server_properties( + owner, shared, + ) + + expected = dict(owner_snapshot) + if self.user_mode is None: + expected.pop('oauth_pgadmin_token_mode') + else: + expected['oauth_pgadmin_token_mode'] = self.user_mode + + self.assertEqual(result.connection_params, expected) + self.assertEqual( + result.connection_params.get('oauth_pgadmin_token_mode', + 'disabled'), + self.user_mode if self.user_mode is not None else 'disabled', + ) + + # The merge works on a copy, never on either user's stored dict. + self.assertIsNot(result.connection_params, owner_params) + self.assertIsNot(result.connection_params, shared.connection_params) + self.assertEqual(owner_params, owner_snapshot) + self.assertEqual(shared.connection_params, shared_snapshot) diff --git a/web/pgadmin/utils/driver/psycopg3/connection.py b/web/pgadmin/utils/driver/psycopg3/connection.py index fe9fb13c4f9..4740b8c57f7 100644 --- a/web/pgadmin/utils/driver/psycopg3/connection.py +++ b/web/pgadmin/utils/driver/psycopg3/connection.py @@ -48,6 +48,12 @@ from io import StringIO from pgadmin.utils.locker import ConnectionLocker from pgadmin.utils.driver import get_driver +from pgadmin.utils.pg_oauth2 import ( + install_oauth_hook, + get_postgres_oauth_token, + oauth_token_context, + OAuthTokenError, +) # On Windows, Psycopg is not compatible with the default ProactorEventLoop. @@ -489,6 +495,32 @@ def _decode_password(self, encpass, manager, password, crypt_key): ).format(str(e)), password return False, '', password + def _oauth_context(self): + """ + Prepare OAuth credentials and scope them to a connection attempt. + + Disabled mode temporarily clears any inherited OAuth token. + """ + connection_params = self.manager.connection_params or {} + mode = connection_params.get("oauth_pgadmin_token_mode", "disabled") + + token = None + + if mode != "disabled": + try: + install_oauth_hook() + except RuntimeError as exc: + raise OAuthTokenError(str(exc)) from None + + token = get_postgres_oauth_token( + mode, + connection_params.get("oauth_client_id"), + connection_params.get("oauth_scope"), + connection_params.get("oauth_issuer"), + ) + + return oauth_token_context(token) + def connect(self, **kwargs): if self.conn: if self.conn.closed: @@ -585,25 +617,40 @@ def connect(self, **kwargs): connection_string = manager.create_connection_string( database, user, password) - if self.async_: - autocommit = True - if 'auto_commit' in kwargs: - autocommit = kwargs['auto_commit'] - - async def connectdbserver(): - return await psycopg.AsyncConnection.connect( - connection_string, - cursor_factory=AsyncDictCursor, - autocommit=autocommit, - prepare_threshold=manager.prepare_threshold - ) - pg_conn = asyncio.run(connectdbserver()) - pg_conn.server_cursor_factory = AsyncDictServerCursor - else: - pg_conn = psycopg.Connection.connect( - connection_string, - cursor_factory=DictCursor, - prepare_threshold=manager.prepare_threshold) + try: + with self._oauth_context(): + if self.async_: + autocommit = True + if 'auto_commit' in kwargs: + autocommit = kwargs['auto_commit'] + + async def connectdbserver(): + return await psycopg.AsyncConnection.connect( + connection_string, + cursor_factory=AsyncDictCursor, + autocommit=autocommit, + prepare_threshold=manager.prepare_threshold + ) + + pg_conn = asyncio.run(connectdbserver()) + pg_conn.server_cursor_factory = ( + AsyncDictServerCursor + ) + else: + pg_conn = psycopg.Connection.connect( + connection_string, + cursor_factory=DictCursor, + prepare_threshold=manager.prepare_threshold, + ) + except OAuthTokenError as exc: + manager.stop_ssh_tunnel() + current_app.logger.warning( + "PostgreSQL OAuth authentication " + "failed for server %s: %s", + manager.sid, + exc, + ) + return False, str(exc) except psycopg.Error as e: manager.stop_ssh_tunnel() @@ -1784,8 +1831,19 @@ def reset(self): connection_string = manager.create_connection_string( self.db, manager.user, password) - pg_conn = psycopg.connect(connection_string, - cursor_factory=DictCursor) + with self._oauth_context(): + pg_conn = psycopg.connect( + connection_string, + cursor_factory=DictCursor, + ) + except OAuthTokenError as exc: + current_app.logger.warning( + "PostgreSQL OAuth authentication failed " + "while resetting the connection for server %s: %s", + manager.sid, + exc, + ) + return False, str(exc) except psycopg.Error as e: if hasattr(e, 'pgerror'): msg = e.pgerror @@ -2010,8 +2068,11 @@ def cancel_transaction(self, conn_id, did=None): connection_string = self.manager.create_connection_string( self.db, self.manager.user, password) - pg_conn = psycopg.connect(connection_string, - cursor_factory=DictCursor) + with self._oauth_context(): + pg_conn = psycopg.connect( + connection_string, + cursor_factory=DictCursor, + ) # Get the cursor and run the query cur = pg_conn.cursor() @@ -2019,7 +2080,9 @@ def cancel_transaction(self, conn_id, did=None): # Close the connection pg_conn.close() - + except OAuthTokenError as exc: + status = False + msg = str(exc) except psycopg.Error as e: status = False if hasattr(e, 'pgerror'): diff --git a/web/pgadmin/utils/driver/psycopg3/server_manager.py b/web/pgadmin/utils/driver/psycopg3/server_manager.py index 8134e74c4e1..0a51b248b4f 100644 --- a/web/pgadmin/utils/driver/psycopg3/server_manager.py +++ b/web/pgadmin/utils/driver/psycopg3/server_manager.py @@ -707,6 +707,10 @@ def create_connection_string(self, database, user, password=None): # Loop through all the connection parameters set in the server dialog. if self.connection_params and isinstance(self.connection_params, dict): for key, value in self.connection_params.items(): + # pgAdmin-only parameter, not a libpq connection option + if key == "oauth_pgadmin_token_mode": + continue + with_complete_path = False orig_value = value # Getting complete file path if the key is one of the below. diff --git a/web/pgadmin/utils/driver/psycopg3/tests/test_connection_oauth.py b/web/pgadmin/utils/driver/psycopg3/tests/test_connection_oauth.py new file mode 100644 index 00000000000..e5a10432292 --- /dev/null +++ b/web/pgadmin/utils/driver/psycopg3/tests/test_connection_oauth.py @@ -0,0 +1,745 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +"""Tests for pgAdmin OAuth integration in the Psycopg 3 connection.""" + +import unittest +from contextlib import ExitStack, nullcontext +from unittest.mock import AsyncMock, MagicMock, patch +from pgadmin.utils.route import BaseTestGenerator + +import psycopg +from flask import Flask + +from pgadmin.utils import pg_oauth2 +import pgadmin.utils.driver.psycopg3.connection as connection_module + +ISSUER = 'https://keycloak.example.test/realms/postgres' + + +class _OAuthUnitTestMixin: + """Run each OAuth unit test as a separate regression scenario.""" + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + + cls.scenarios = [ + (name[5:], {'case_method': name}) + for name in sorted(dir(cls)) + if name.startswith('case_') and callable(getattr(cls, name)) + ] + + def setUp(self): + # These tests mock database connections. + unittest.TestCase.setUp(self) + + def runTest(self): + getattr(self, self.case_method)() + + +class TestOAuthConnection( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test pgAdmin OAuth handling during Psycopg connection creation.""" + + def setUp(self): + self.app = Flask(__name__) + self.app.config.update( + SECRET_KEY='connection-oauth-unit-test-secret', + TESTING=True, + ) + + self.token_handle = pg_oauth2._oauth_token.set(None) + + def tearDown(self): + pg_oauth2._oauth_token.reset(self.token_handle) + + @staticmethod + def _make_manager(mode='disabled', client_id=None, oauth_scope=None): + manager = MagicMock() + + manager.sid = 1 + manager.db = 'postgres' + manager.user = 'database-user' + manager.password = None + manager.role = None + + manager.use_ssh_tunnel = 0 + manager.tunnel_created = False + manager.kerberos_conn = False + + manager.passexec = None + manager.prepare_threshold = None + + manager.connection_params = { + 'oauth_pgadmin_token_mode': mode, + 'oauth_issuer': ISSUER, + } + + if client_id is not None: + manager.connection_params['oauth_client_id'] = client_id + if oauth_scope is not None: + manager.connection_params['oauth_scope'] = oauth_scope + + manager.get_connection_param_value.return_value = None + manager.create_connection_string.return_value = ( + 'host=database.example.test ' + 'dbname=postgres ' + 'user=database-user' + ) + + return manager + + @staticmethod + def _make_connection(manager, async_=False): + # Construct the object without invoking Connection.__init__(). + # Its constructor obtains unrelated global driver state that is + # not part of these tests. + connection = object.__new__(connection_module.Connection) + + connection.manager = manager + connection.conn_id = 'DB:postgres' + connection.db = 'postgres' + connection.conn = None + + connection.password = None + connection.reconnecting = False + connection.wasConnected = False + connection.async_ = 1 if async_ else 0 + + # Connection initialization executes several SQL statements. + # These tests stop at the Psycopg connection boundary. + connection._initialize = MagicMock( + return_value=(True, None) + ) + + return connection + + @staticmethod + def _make_psycopg_connection(): + pg_connection = MagicMock() + pg_connection.closed = False + return pg_connection + + @staticmethod + def _apply_common_patches(stack): + stack.enter_context( + patch.object( + connection_module, + 'get_crypt_key', + return_value=(True, 'crypt-key') + ) + ) + stack.enter_context( + patch.object( + connection_module, + 'get_complete_file_path', + return_value=None + ) + ) + stack.enter_context( + patch.object( + connection_module, + 'ConnectionLocker', + return_value=nullcontext() + ) + ) + stack.enter_context( + patch.object( + connection_module, + 'gettext', + side_effect=lambda message, *args, **kwargs: message + ) + ) + stack.enter_context( + patch.object( + connection_module, + '_', + side_effect=lambda message, *args, **kwargs: message + ) + ) + + def case_oauth_disabled_does_not_access_pgadmin_token(self): + """ + Existing non-OAuth connections must not install the OAuth hook + or access the pgAdmin OAuth session. + """ + manager = self._make_manager(mode='disabled') + connection = self._make_connection(manager) + pg_connection = self._make_psycopg_connection() + + def connect_without_oauth(*args, **kwargs): + self.assertIsNone(pg_oauth2._oauth_token.get()) + return pg_connection + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + install_hook = stack.enter_context( + patch.object( + connection_module, + 'install_oauth_hook' + ) + ) + get_token = stack.enter_context( + patch.object( + connection_module, + 'get_postgres_oauth_token' + ) + ) + psycopg_connect = stack.enter_context( + patch.object( + connection_module.psycopg.Connection, + 'connect', + side_effect=connect_without_oauth + ) + ) + + with pg_oauth2.oauth_token_context('outer-token'): + status, message = connection.connect() + + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'outer-token', + ) + + self.assertTrue(status) + self.assertIsNone(message) + + install_hook.assert_not_called() + get_token.assert_not_called() + psycopg_connect.assert_called_once() + + connection._initialize.assert_called_once() + manager.create_connection_string.assert_called_once_with( + 'postgres', + 'database-user', + None + ) + + def case_oauth_enabled_without_token_returns_error(self): + """ + An OAuth-enabled server must fail clearly before calling + Psycopg when the pgAdmin session has no access token. + """ + manager = self._make_manager(mode='direct') + connection = self._make_connection(manager) + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + install_hook = stack.enter_context( + patch.object( + connection_module, + 'install_oauth_hook' + ) + ) + get_token = stack.enter_context( + patch.object( + connection_module, + 'get_postgres_oauth_token', + side_effect=pg_oauth2.OAuthTokenError( + 'No current pgAdmin OAuth access token' + ' is available. Sign in to pgAdmin again.' + ) + ) + ) + psycopg_connect = stack.enter_context( + patch.object( + connection_module.psycopg.Connection, + 'connect' + ) + ) + + status, message = connection.connect() + + self.assertFalse(status) + self.assertEqual( + message, + 'No current pgAdmin OAuth access token ' + 'is available. Sign in to pgAdmin again.' + ) + + get_token.assert_called_once_with('direct', None, None, ISSUER) + psycopg_connect.assert_not_called() + connection._initialize.assert_not_called() + manager.stop_ssh_tunnel.assert_called_once_with() + + def case_sync_connection_is_created_inside_token_context(self): + """ + The synchronous Psycopg connection must be created while the + pgAdmin access token is present in the ContextVar. + """ + manager = self._make_manager(mode='direct') + connection = self._make_connection(manager) + pg_connection = self._make_psycopg_connection() + + def connect_with_oauth(*args, **kwargs): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'pgadmin-access-token' + ) + return pg_connection + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + install_hook = stack.enter_context( + patch.object( + connection_module, + 'install_oauth_hook' + ) + ) + get_token = stack.enter_context( + patch.object( + connection_module, + 'get_postgres_oauth_token', + return_value='pgadmin-access-token' + ) + ) + psycopg_connect = stack.enter_context( + patch.object( + connection_module.psycopg.Connection, + 'connect', + side_effect=connect_with_oauth + ) + ) + + status, message = connection.connect() + + self.assertTrue(status) + self.assertIsNone(message) + self.assertIsNone(pg_oauth2._oauth_token.get()) + + install_hook.assert_called_once_with() + get_token.assert_called_once_with('direct', None, None, ISSUER) + psycopg_connect.assert_called_once() + + connection._initialize.assert_called_once() + self.assertIs(connection.conn, pg_connection) + + def case_async_connection_is_created_inside_token_context(self): + """ + The asynchronous Psycopg connection must inherit the token + context while asyncio.run() executes the connection coroutine. + """ + manager = self._make_manager(mode='direct') + connection = self._make_connection( + manager, + async_=True + ) + pg_connection = self._make_psycopg_connection() + + async def connect_with_oauth(*args, **kwargs): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'pgadmin-access-token' + ) + return pg_connection + + async_connect = AsyncMock( + side_effect=connect_with_oauth + ) + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + install_hook = stack.enter_context( + patch.object( + connection_module, + 'install_oauth_hook' + ) + ) + get_token = stack.enter_context( + patch.object( + connection_module, + 'get_postgres_oauth_token', + return_value='pgadmin-access-token' + ) + ) + stack.enter_context( + patch.object( + connection_module.psycopg.AsyncConnection, + 'connect', + new=async_connect + ) + ) + + status, message = connection.connect() + + self.assertTrue(status) + self.assertIsNone(message) + self.assertIsNone(pg_oauth2._oauth_token.get()) + + install_hook.assert_called_once_with() + get_token.assert_called_once_with('direct', None, None, ISSUER) + async_connect.assert_awaited_once() + + connection._initialize.assert_called_once() + self.assertIs(connection.conn, pg_connection) + self.assertIs( + pg_connection.server_cursor_factory, + connection_module.AsyncDictServerCursor + ) + + def case_token_context_is_restored_after_connection_failure(self): + """ + A Psycopg failure must not leave the failed connection's token + available to later libpq callbacks. + """ + manager = self._make_manager(mode='direct') + connection = self._make_connection(manager) + + def fail_inside_oauth_context(*args, **kwargs): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'pgadmin-access-token' + ) + raise psycopg.OperationalError( + 'test connection failure' + ) + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + stack.enter_context( + patch.object( + connection_module, + 'install_oauth_hook' + ) + ) + get_token = stack.enter_context( + patch.object( + connection_module, + 'get_postgres_oauth_token', + return_value='pgadmin-access-token' + ) + ) + psycopg_connect = stack.enter_context( + patch.object( + connection_module.psycopg.Connection, + 'connect', + side_effect=fail_inside_oauth_context + ) + ) + + status, message = connection.connect() + + self.assertFalse(status) + self.assertIn('test connection failure', message) + self.assertIsNone(pg_oauth2._oauth_token.get()) + + get_token.assert_called_once_with('direct', None, None, ISSUER) + psycopg_connect.assert_called_once() + connection._initialize.assert_not_called() + self.assertIsNone(connection.conn) + + def case_exchange_mode_passes_client_id(self): + """ + Token exchange mode must pass the configured OAuth client ID + when obtaining the PostgreSQL token. + """ + manager = self._make_manager( + mode='exchange', + client_id='postgres-cluster', + oauth_scope='openid', + ) + connection = self._make_connection(manager) + pg_connection = self._make_psycopg_connection() + + def connect_with_exchanged_token(*args, **kwargs): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'exchanged-token' + ) + return pg_connection + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + install_hook = stack.enter_context( + patch.object( + connection_module, + 'install_oauth_hook' + ) + ) + get_token = stack.enter_context( + patch.object( + connection_module, + 'get_postgres_oauth_token', + return_value='exchanged-token' + ) + ) + psycopg_connect = stack.enter_context( + patch.object( + connection_module.psycopg.Connection, + 'connect', + side_effect=connect_with_exchanged_token + ) + ) + + status, message = connection.connect() + + self.assertTrue(status) + self.assertIsNone(message) + + install_hook.assert_called_once_with() + get_token.assert_called_once_with( + 'exchange', + 'postgres-cluster', + 'openid', + ISSUER, + ) + psycopg_connect.assert_called_once() + connection._initialize.assert_called_once() + + def case_missing_connection_params_uses_non_oauth_path(self): + """ + A connection manager without connection parameters must behave as a + normal non-OAuth connection. + """ + manager = self._make_manager() + manager.connection_params = None + + connection = self._make_connection(manager) + pg_connection = self._make_psycopg_connection() + + def connect_without_oauth(*args, **kwargs): + self.assertIsNone(pg_oauth2._oauth_token.get()) + return pg_connection + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + install_hook = stack.enter_context( + patch.object( + connection_module, + 'install_oauth_hook' + ) + ) + get_token = stack.enter_context( + patch.object( + connection_module, + 'get_postgres_oauth_token' + ) + ) + psycopg_connect = stack.enter_context( + patch.object( + connection_module.psycopg.Connection, + 'connect', + side_effect=connect_without_oauth + ) + ) + + status, message = connection.connect() + + self.assertTrue(status) + self.assertIsNone(message) + + install_hook.assert_not_called() + get_token.assert_not_called() + psycopg_connect.assert_called_once() + + self.assertIs(connection.conn, pg_connection) + connection._initialize.assert_called_once() + + @patch.object(connection_module, 'install_oauth_hook') + @patch.object(connection_module, 'get_postgres_oauth_token') + @patch.object(connection_module.psycopg, 'connect') + def case_reset_creates_connection_inside_token_context( + self, psycopg_connect, get_token, install_hook): + manager = self._make_manager( + mode='exchange', + client_id='postgres-cluster', + oauth_scope='openid', + ) + connection = self._make_connection(manager) + connection._decrypt_password = MagicMock( + return_value=(True, '', None), + ) + + pg_connection = self._make_psycopg_connection() + pg_connection.info.backend_pid = 23456 + get_token.return_value = 'exchanged-token' + + def connect_with_token(*args, **kwargs): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'exchanged-token', + ) + return pg_connection + + psycopg_connect.side_effect = connect_with_token + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + with pg_oauth2.oauth_token_context('outer-token'): + status, message = connection.reset() + + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'outer-token', + ) + + self.assertTrue(status) + self.assertIsNone(message) + self.assertIs(connection.conn, pg_connection) + self.assertIsNone(pg_oauth2._oauth_token.get()) + + install_hook.assert_called_once_with() + get_token.assert_called_once_with( + 'exchange', + 'postgres-cluster', + 'openid', + ISSUER, + ) + psycopg_connect.assert_called_once_with( + manager.create_connection_string.return_value, + cursor_factory=connection_module.DictCursor, + ) + + @patch.object(connection_module, 'install_oauth_hook') + @patch.object(connection_module, 'get_postgres_oauth_token') + @patch.object(connection_module.psycopg, 'connect') + def case_cancel_creates_connection_inside_token_context( + self, psycopg_connect, get_token, install_hook): + manager = self._make_manager( + mode='exchange', + client_id='postgres-cluster', + oauth_scope='openid', + ) + connection = self._make_connection(manager) + + connection._Connection__backend_pid = 12345 + manager.connection.return_value = connection + + pg_connection = self._make_psycopg_connection() + get_token.return_value = 'exchanged-token' + + def connect_with_token(*args, **kwargs): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'exchanged-token', + ) + return pg_connection + + psycopg_connect.side_effect = connect_with_token + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + with pg_oauth2.oauth_token_context('outer-token'): + status, message = connection.cancel_transaction( + connection.conn_id, + ) + + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'outer-token', + ) + + self.assertTrue(status, message) + self.assertIsNone(pg_oauth2._oauth_token.get()) + + install_hook.assert_called_once_with() + get_token.assert_called_once_with( + 'exchange', + 'postgres-cluster', + 'openid', + ISSUER, + ) + psycopg_connect.assert_called_once_with( + manager.create_connection_string.return_value, + cursor_factory=connection_module.DictCursor, + ) + pg_connection.cursor.return_value.execute.assert_called_once_with( + 'SELECT pg_cancel_backend(12345);', + ) + pg_connection.close.assert_called_once_with() + + def case_auxiliary_connections_handle_oauth_token_errors(self): + cases = ( + ('reset', pg_oauth2.OAuthTokenError), + ('reset', pg_oauth2.OAuthRequestContextRequired), + ('cancel', pg_oauth2.OAuthTokenError), + ('cancel', pg_oauth2.OAuthRequestContextRequired), + ) + + for operation, error_type in cases: + with self.subTest( + operation=operation, + error_type=error_type.__name__): + manager = self._make_manager( + mode='exchange', + client_id='postgres-cluster', + oauth_scope='openid', + ) + connection = self._make_connection(manager) + connection._decrypt_password = MagicMock( + return_value=(True, '', None), + ) + connection._Connection__backend_pid = 12345 + manager.connection.return_value = connection + + error_message = 'OAuth credentials unavailable.' + + with self.app.test_request_context('/'): + with ExitStack() as stack: + self._apply_common_patches(stack) + + install_hook = stack.enter_context( + patch.object( + connection_module, + 'install_oauth_hook', + ) + ) + get_token = stack.enter_context( + patch.object( + connection_module, + 'get_postgres_oauth_token', + side_effect=error_type(error_message), + ) + ) + psycopg_connect = stack.enter_context( + patch.object( + connection_module.psycopg, + 'connect', + ) + ) + + if operation == 'reset': + status, message = connection.reset() + else: + status, message = ( + connection.cancel_transaction( + connection.conn_id, + ) + ) + + self.assertFalse(status) + self.assertEqual(message, error_message) + self.assertIsNone(pg_oauth2._oauth_token.get()) + + install_hook.assert_called_once_with() + get_token.assert_called_once_with( + 'exchange', + 'postgres-cluster', + 'openid', + ISSUER, + ) + psycopg_connect.assert_not_called() diff --git a/web/pgadmin/utils/driver/psycopg3/tests/test_reset_error.py b/web/pgadmin/utils/driver/psycopg3/tests/test_reset_error.py index b4f0a53e2a9..20caec22f44 100644 --- a/web/pgadmin/utils/driver/psycopg3/tests/test_reset_error.py +++ b/web/pgadmin/utils/driver/psycopg3/tests/test_reset_error.py @@ -32,6 +32,7 @@ class ResetConnectionTest(BaseTestGenerator): def runTest(self): manager = MagicMock(sid=1, kerberos_conn=False) + manager.connection_params = {} manager.create_connection_string.return_value = 'dbname=testdb' conn = Connection(manager, 'test-conn-id', 'testdb') pg_conn = MagicMock() diff --git a/web/pgadmin/utils/pg_oauth2.py b/web/pgadmin/utils/pg_oauth2.py new file mode 100644 index 00000000000..c1368997327 --- /dev/null +++ b/web/pgadmin/utils/pg_oauth2.py @@ -0,0 +1,1131 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +""" +PostgreSQL 18 libpq OAuth bearer-token integration for pgAdmin. + +The libpq OAuth hook is process-global. The OAuth access token is therefore +not obtained from Flask's request/session context inside the C callback. + +Instead, connection.py establishes the token in a ContextVar immediately +before creating a Psycopg connection. The libpq callback reads that +ContextVar when authentication requests a bearer token. +""" + +import base64 +import ctypes +import ctypes.util +import json +import logging +import os +import re +import sys +import threading +import time +from authlib.integrations.requests_client import OAuth2Session +from collections.abc import Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, Iterator, Optional +from urllib.parse import urlsplit + +import requests +from flask import has_request_context, session + +_install_lock = threading.Lock() +logger = logging.getLogger(__name__) + + +# PostgreSQL 18: +# +# typedef enum +# { +# PQAUTHDATA_PROMPT_OAUTH_DEVICE, +# PQAUTHDATA_OAUTH_BEARER_TOKEN +# } PGauthData; +PQAUTHDATA_OAUTH_BEARER_TOKEN = 1 + +_OAUTH_REFRESH_LEEWAY_SECONDS = 30 + + +class PGoauthBearerRequest(ctypes.Structure): + """ + PostgreSQL 18 PGoauthBearerRequest. + + The async and cleanup members are function pointers in C. They are + represented as c_void_p here because we assign the corresponding + ctypes callbacks explicitly below. + """ + + # This ctypes structure follows PostgreSQL 18's public libpq ABI. + # Review its layout and callback signatures before enabling another + # major version. + + _fields_ = [ + ("openid_configuration", ctypes.c_char_p), + ("scope", ctypes.c_char_p), + ("async_", ctypes.c_void_p), + ("cleanup", ctypes.c_void_p), + ("token", ctypes.c_char_p), + ("user", ctypes.c_void_p), + ] + + +_AUTH_HOOK = ctypes.CFUNCTYPE( + ctypes.c_int, + ctypes.c_int, # PGauthData + ctypes.c_void_p, # PGconn * + ctypes.c_void_p, # void * +) + + +_CLEANUP_HOOK = ctypes.CFUNCTYPE( + None, + ctypes.c_void_p, # PGconn * + ctypes.c_void_p, # PGoauthBearerRequest * +) + + +# The token is established by connection.py while a connection is being +# created. ContextVar is preferable to Flask's session here because the +# libpq callback is process-global and may be invoked from C code. +_oauth_token: ContextVar[Optional[str]] = ContextVar( + "pgadmin_oauth_token", + default=None, +) + + +# Keep native objects alive for as long as libpq can call them. +_libpq = None +_hook = None +_cleanup_hook = None +_previous_hook = None + +# PGoauthBearerRequest pointer -> token buffer. +# +# libpq owns the request structure and calls cleanup() when it no longer +# needs request->token. +_token_buffers_lock = threading.Lock() +_token_buffers = {} + + +def _check_oauth_libpq_version(version: int) -> None: + if not 180000 <= version < 190000: + raise RuntimeError( + "This PostgreSQL OAuth adapter supports libpq 18.x; " + f"found libpq version {version}." + ) + + +def _delegate_to_previous(authdata_type, conn, data): + """Delegate auth-data requests not handled by pgAdmin.""" + if _previous_hook is None: + return 0 + + try: + return _previous_hook(authdata_type, conn, data) + except Exception: + # Never allow an exception to escape through a C callback. + logger.exception("Previous libpq auth-data hook failed.") + return -1 + + +_LIBPQ_FILENAME_RE = re.compile( + r"^libpq(?:[-_][A-Za-z0-9_-]+)?" + r"(?:\.so(?:\..*)?|(?:\.\d+)*\.dylib|\.dll)$", + re.IGNORECASE, +) + + +def _is_libpq_path(path: str) -> bool: + """Return whether the filename matches a recognized libpq name.""" + return bool( + _LIBPQ_FILENAME_RE.fullmatch(os.path.basename(path)) + ) + + +def _get_loaded_libpq_paths() -> list[str]: + """ + Return distinct recognized libpq paths from loaded-library listings. + + Use ctypes.util.dllist() when available. On Linux, also inspect + /proc/self/maps, combining results from both sources. + + Reject recognized libpq entries with deleted mappings or relative + paths. No library is loaded by this function. + + Results are candidates identified by filename; they do not prove + which library Psycopg uses or detect arbitrarily renamed libraries. + """ + paths = [] + + dllist = getattr(ctypes.util, "dllist", None) + + if callable(dllist): + try: + paths.extend(path for path in dllist() if path) + except OSError: + pass + + # Inspect Linux mappings even when dllist() returned entries. + if sys.platform.startswith("linux"): + try: + with open( + "/proc/self/maps", + "r", + encoding="utf-8", + ) as maps: + for line in maps: + parts = line.rstrip().split(None, 5) + + if len(parts) == 6: + paths.append(parts[5]) + + except OSError: + pass + + result = [] + seen = set() + + for path in paths: + deleted = path.endswith(" (deleted)") + candidate = path[:-10] if deleted else path + + if not _is_libpq_path(candidate): + continue + + if deleted: + raise RuntimeError( + "A loaded libpq library has a deleted mapping. " + "Restart pgAdmin before using PostgreSQL OAuth." + ) + + if not os.path.isabs(candidate): + raise RuntimeError( + "A loaded libpq candidate has a relative path: " + "'{0}'. Refusing to resolve it through another " + "library search.".format(candidate) + ) + + candidate = os.path.realpath(candidate) + + if candidate not in seen: + seen.add(candidate) + result.append(candidate) + + return result + + +def _get_libpq() -> ctypes.CDLL: + """ + Obtain a libpq handle and configure the OAuth auth-data hook API. + + Requires libpq 18.x. This function configures the ctypes function + signatures; install_oauth_hook() registers the callback. + + Psycopg implementations: + + * "python": + Reuse Psycopg's existing ctypes.CDLL handle directly, targeting + the same library instance used for its connections. This relies + on Psycopg's private _pq_ctypes module. + + * "c": + Supported on Linux only. Select the sole recognized libpq + library already loaded in the process and obtain its handle + without loading a new library. Reject missing or ambiguous + candidates. + + This assumes a normal dynamically linked Psycopg installation. + Library discovery is a conservative heuristic, not proof of + Psycopg's actual symbol bindings. Unrecognized library names + or unusual linking arrangements may invalidate this assumption. + + * "binary": + Unsupported by this adapter because it cannot reliably obtain + the bundled libpq instance used by Psycopg. + + Runtime versions must match, but matching PQlibVersion() values + are only a consistency check and do not establish library identity. + + This is a temporary integration pending usable native OAuth hook + support in the Psycopg version supported by pgAdmin. Replace it + with that API when available. + + Related Psycopg work: + https://github.com/psycopg/psycopg/pull/1190 + """ + global _libpq + + if _libpq is not None: + return _libpq + + from psycopg import pq + + implementation = pq.__impl__ + + if implementation == "binary": + raise RuntimeError( + "PostgreSQL OAuth bearer-token support cannot currently " + "install the authentication hook when Psycopg uses its " + "bundled 'binary' implementation. Use Psycopg's 'c' or " + "'python' implementation instead." + ) + + if implementation not in ("c", "python"): + raise RuntimeError( + "Unsupported Psycopg pq implementation '{0}' for PostgreSQL " + "OAuth bearer-token support.".format(implementation) + ) + + psycopg_version = pq.version() + _check_oauth_libpq_version(psycopg_version) + + libpq = None + libpq_identity = None + + if implementation == "c" and not sys.platform.startswith("linux"): + raise RuntimeError( + "This PostgreSQL OAuth adapter currently supports " + "Psycopg's 'c' implementation only on Linux. " + "Use the 'python' implementation with libpq 18.x." + ) + + if implementation == "python": + # Psycopg's pure-Python implementation itself uses ctypes. + # + # Reuse its existing CDLL object instead of rediscovering or + # reopening libpq. This guarantees that the hook is installed + # into the exact library through which Psycopg makes connections. + try: + from psycopg.pq import _pq_ctypes + except ImportError as exc: + raise RuntimeError( + "Could not access the libpq handle used by Psycopg's " + "'python' implementation." + ) from exc + + libpq = _pq_ctypes.pq + libpq_identity = getattr( + libpq, + "_name", + "", + ) + + elif implementation == "c": + # Importing psycopg.pq has initialized Psycopg's C implementation. + # + # Select the sole recognized libpq already loaded in this process. + # Under normal dynamic-linking assumptions, this is Psycopg's libpq. + # Reject missing or ambiguous candidates rather than guessing. + loaded_libpq = _get_loaded_libpq_paths() + + if not loaded_libpq: + raise RuntimeError( + "Psycopg is using its 'c' implementation, but the libpq " + "loaded by that implementation could not be identified. " + "Refusing to install the PostgreSQL OAuth hook into an " + "unverified libpq library." + ) + + if len(loaded_libpq) > 1: + raise RuntimeError( + "Psycopg is using its 'c' implementation, but multiple " + "libpq libraries are loaded in the current process. " + "Cannot safely determine which library Psycopg uses: " + "{0}".format(", ".join(loaded_libpq)) + ) + + libpq_identity = loaded_libpq[0] + + try: + no_load = getattr(os, "RTLD_NOLOAD", None) + if no_load is None: + raise RuntimeError( + "Cannot obtain an already-loaded libpq handle " + "on this platform." + ) + + libpq = ctypes.CDLL( + libpq_identity, + mode=os.RTLD_NOW | no_load, + ) + except OSError as exc: + raise RuntimeError( + "Could not obtain an existing handle to the selected " + "libpq at '{0}': {1}".format(libpq_identity, exc) + ) from exc + + # Validate the actual library behind the handle. + try: + pq_lib_version = libpq.PQlibVersion + except AttributeError as exc: + raise RuntimeError( + "The selected libpq does not expose PQlibVersion()." + ) from exc + + pq_lib_version.argtypes = [] + pq_lib_version.restype = ctypes.c_int + + version = pq_lib_version() + + # Check runtime-version consistency. + # For the C implementation, selecting the sole recognized loaded + # libpq is a conservative heuristic, not proof of Psycopg's binding. + if version != psycopg_version: + raise RuntimeError( + "The identified libpq reports version {0}, but Psycopg " + "reports version {1}.".format( + version, + psycopg_version, + ) + ) + + _check_oauth_libpq_version(version) + + try: + pq_set_auth_data_hook = libpq.PQsetAuthDataHook + pq_get_auth_data_hook = libpq.PQgetAuthDataHook + except AttributeError as exc: + raise RuntimeError( + "The selected libpq does not provide the " + "PostgreSQL 18 OAuth auth-data hook API." + ) from exc + + pq_set_auth_data_hook.argtypes = [_AUTH_HOOK] + pq_set_auth_data_hook.restype = None + + pq_get_auth_data_hook.argtypes = [] + pq_get_auth_data_hook.restype = ctypes.c_void_p + + _libpq = libpq + + logger.info( + "Using Psycopg %s libpq '%s' version %d for PostgreSQL " + "OAuth support.", + implementation, + libpq_identity, + version, + ) + + return _libpq + + +@_CLEANUP_HOOK +def _oauth_cleanup(conn, request): + """ + Release native memory associated with a libpq OAuth request. + + libpq calls this after it no longer needs request->token. + """ + try: + request_key = int(ctypes.cast(request, ctypes.c_void_p).value) + + with _token_buffers_lock: + _token_buffers.pop(request_key, None) + + logger.debug( + "Released PostgreSQL OAuth bearer-token " + "buffer for request %s.", request_key + ) + except Exception: + # Never allow an exception to escape a C callback. + pass + + +@_AUTH_HOOK +def _oauth_hook(authdata_type, conn, data): + """ + Supply the pgAdmin OAuth access token when libpq requests one. + """ + if authdata_type != PQAUTHDATA_OAUTH_BEARER_TOKEN: + return _delegate_to_previous(authdata_type, conn, data) + + try: + token = _oauth_token.get() + + if not token: + # This may be an unrelated OAuth connection made after the + # process-global pgAdmin hook was installed. + return _delegate_to_previous(authdata_type, conn, data) + + request = ctypes.cast(data, ctypes.POINTER( + PGoauthBearerRequest)).contents + + token_buffer = ctypes.create_string_buffer(token.encode("utf-8")) + + request_key = int(ctypes.cast(data, ctypes.c_void_p).value) + + # Keep the token memory alive until libpq invokes cleanup(). + with _token_buffers_lock: + _token_buffers[request_key] = token_buffer + + request.token = ctypes.cast(token_buffer, ctypes.c_char_p) + + # The token is supplied synchronously. No async callback is needed. + request.async_ = None + + # Tell libpq when it is safe to release the token buffer. + request.cleanup = ctypes.cast(_oauth_cleanup, ctypes.c_void_p) + + logger.debug( + "Supplying OAuth bearer token to " + "PostgreSQL libpq (token length=%d).", + len(token), + ) + + return 1 + + except Exception: + # Never allow an exception to cross the C callback boundary. + logger.exception( + "Exception while supplying OAuth bearer token to libpq.") + return -1 + + +@contextmanager +def oauth_token_context(token: str) -> Iterator[None]: + """ + Associate an OAuth access token with the current connection operation. + + The token is available to the libpq callback through ContextVar and is + automatically restored after the connection attempt. + """ + token_handle = _oauth_token.set(token) + + try: + yield + finally: + _oauth_token.reset(token_handle) + + +def install_oauth_hook() -> None: + """ + Install the process-wide pgAdmin OAuth bearer-token hook. + + Called lazily before the first PostgreSQL OAuth connection. + Repeated calls are safe; installation is serialized per process. + """ + global _hook + global _cleanup_hook + global _previous_hook + + if _hook is not None: + return + + with _install_lock: + if _hook is not None: + return + + libpq = _get_libpq() + + previous_hook_ptr = libpq.PQgetAuthDataHook() + _previous_hook = _AUTH_HOOK( + previous_hook_ptr) if previous_hook_ptr else None + + libpq.PQsetAuthDataHook(_oauth_hook) + + _cleanup_hook = _oauth_cleanup + _hook = _oauth_hook + + logger.info("Installed PostgreSQL libpq OAuth bearer-token hook.") + + +def _get_token_expiry(token_data: Mapping[str, Any]) -> Optional[float]: + """ + Return the access-token expiration timestamp. + + Authlib normally stores expires_at. Fall back to the JWT exp claim + because some providers omit expires_at from their token response. + """ + expires_at = token_data.get("expires_at") + + if expires_at is not None: + try: + return float(expires_at) + except (TypeError, ValueError): + logger.warning("Invalid expires_at value in pgAdmin OAuth token.") + + access_token = token_data.get("access_token") + + if not isinstance(access_token, str) or not access_token: + return None + + try: + parts = access_token.split(".") + + if len(parts) != 3: + # Opaque access token: expiration cannot be decoded locally. + return None + + payload = parts[1] + payload += "=" * (-len(payload) % 4) + + claims = json.loads(base64.urlsafe_b64decode(payload.encode("ascii"))) + + if not isinstance(claims, dict): + return None + + expires_at = claims.get("exp") + + if expires_at is not None: + return float(expires_at) + + except (ValueError, TypeError, UnicodeError, json.JSONDecodeError): + logger.warning( + "Could not decode expiration from pgAdmin OAuth access token.") + + return None + + +class OAuthTokenError(RuntimeError): + """Unable to obtain a PostgreSQL OAuth bearer token.""" + + +class OAuthTokenExchangeError(OAuthTokenError): + """Token exchange failed; the connection must not proceed.""" + + +class OAuthRequestContextRequired(OAuthTokenError): + """OAuth credentials require an authenticated request.""" + + +class OAuthRefreshUnsupportedError(OAuthTokenError): + """The configured client authentication cannot be used for refresh.""" + + +def _get_current_oauth_client() -> Any: + """Return the OAuth client used for the current pgAdmin login.""" + if not has_request_context(): + raise OAuthTokenError( + "OAuth authentication requires a pgAdmin login session." + ) + + provider_name = session.get("oauth2_provider") + + if not isinstance(provider_name, str) or not provider_name: + raise OAuthTokenError( + "The OAuth provider is not recorded in this pgAdmin " + "session. Sign in to pgAdmin again." + ) + + # Import lazily to avoid an import cycle during pgAdmin startup. + from pgadmin.authenticate import get_auth_sources + from pgadmin.utils.constants import OAUTH2 + + oauth_source = get_auth_sources(OAUTH2) + + if oauth_source is None: + raise OAuthTokenError( + "The pgAdmin OAuth authentication source is unavailable." + ) + + oauth_client = oauth_source.oauth2_clients.get(provider_name) + + if oauth_client is None: + raise OAuthTokenError( + "The OAuth provider recorded in this pgAdmin session " + "is not registered." + ) + + return oauth_client + + +def _refresh_pgadmin_oauth_token( + token_data: Mapping[str, Any], +) -> Optional[str]: + """ + Refresh the pgAdmin OAuth access token. + + Return the refreshed access token, or None when refreshing is not + possible. Raise OAuthRefreshUnsupportedError when workload-identity + client authentication is configured. + """ + refresh_token = token_data.get("refresh_token") + + if not isinstance(refresh_token, str) or not refresh_token: + logger.warning( + "The pgAdmin OAuth access token has expired, but no " + "refresh token is available." + ) + return None + + try: + from pgadmin.authenticate import get_auth_sources + from pgadmin.utils.constants import OAUTH2 + + oauth_client = _get_current_oauth_client() + + # Use this user's login provider, not process-global selection. + provider_name = session["oauth2_provider"] + oauth_source = get_auth_sources(OAUTH2) + provider = oauth_source.oauth2_config[provider_name] + + client_auth_method = provider.get( + "OAUTH2_CLIENT_AUTH_METHOD", + "client_secret", + ) + + if ( + isinstance(client_auth_method, str) and + client_auth_method.strip().lower() == "workload_identity" + ): + raise OAuthRefreshUnsupportedError( + "Automatic OAuth token refresh for workload-identity " + "login is not supported by this PostgreSQL integration. " + "Sign in to pgAdmin again, then retry the connection." + ) + + refreshed_token = oauth_client.fetch_access_token( + grant_type="refresh_token", + refresh_token=refresh_token, + ) + + if not isinstance(refreshed_token, dict): + logger.warning( + "The OAuth provider returned an invalid token " + "response while refreshing the pgAdmin access token." + ) + return None + + access_token = refreshed_token.get("access_token") + + if not isinstance(access_token, str) or not access_token: + logger.warning( + "The OAuth provider did not return a valid access " + "token while refreshing the pgAdmin access token." + ) + return None + + # Refresh responses may omit unchanged fields such as id_token, + # scope, token_type, or refresh_token. + updated_token = dict(token_data) + updated_token.pop("expires_at", None) + updated_token.pop("expires_in", None) + updated_token.update(refreshed_token) + + # Some providers do not return another refresh token. + # Keep the existing token in that case. + if not updated_token.get("refresh_token"): + updated_token["refresh_token"] = refresh_token + + session["oauth2_token"] = updated_token + session.modified = True + + logger.info("Refreshed pgAdmin OAuth access token.") + + return access_token + + except OAuthRefreshUnsupportedError: + # Preserve the actionable message for the connection error handler. + raise + except OAuthTokenError as exc: + logger.warning( + "Cannot refresh the pgAdmin OAuth access token: %s", + exc, + ) + return None + except Exception: + logger.warning( + "Failed to refresh the pgAdmin OAuth access token." + ) + return None + + +def _get_pgadmin_oauth_token() -> Optional[str]: + """ + Return the current pgAdmin OAuth access token. + + Refresh the token when it is expired or close to expiration. + Return None when no usable token is available. + """ + try: + token_data = session.get("oauth2_token") + except RuntimeError: + # No Flask request context. + return None + + if not isinstance(token_data, dict): + return None + + access_token = token_data.get("access_token") + + if not isinstance(access_token, str) or not access_token: + return None + + expires_at = _get_token_expiry(token_data) + + if expires_at is None: + return access_token + + now = time.time() + + # Refresh slightly early to avoid the token expiring during the + # PostgreSQL authentication handshake. + if now >= expires_at - _OAUTH_REFRESH_LEEWAY_SECONDS: + logger.info( + "pgAdmin OAuth access token is expired or close to " + "expiration: expires_at=%s now=%s.", + expires_at, + now, + ) + + return _refresh_pgadmin_oauth_token(token_data) + + return access_token + + +def _get_validated_oauth_client(oauth_issuer: Optional[str]): + """Return the login client, issuer-checked metadata and TLS setting.""" + if ( + not isinstance(oauth_issuer, str) or + not oauth_issuer.strip() or + "\0" in oauth_issuer + ): + raise OAuthTokenError( + "Set oauth_issuer to the issuer of the pgAdmin login " + "OIDC provider." + ) + + oauth_client = _get_current_oauth_client() + client_kwargs = oauth_client.client_kwargs or {} + verify_tls = client_kwargs.get("verify", True) + + if verify_tls is not True and not ( + isinstance(verify_tls, str) and verify_tls.strip() + ): + raise OAuthTokenError( + "PostgreSQL OAuth requires TLS verification for the login " + "OIDC provider. Enable verification or configure a trusted " + "CA bundle." + ) + + try: + metadata = oauth_client.load_server_metadata() + except Exception: + # Do not expose provider responses or exception text. + raise OAuthTokenError( + "Could not load the pgAdmin login OIDC provider metadata." + ) from None + + if not isinstance(metadata, dict): + raise OAuthTokenError( + "The pgAdmin login OIDC provider returned invalid metadata." + ) + + issuer = metadata.get("issuer") + if not isinstance(issuer, str) or not issuer.strip(): + raise OAuthTokenError( + "The pgAdmin login OIDC provider metadata contains no issuer." + ) + + # Issuer identifiers must match exactly. Do not normalize slashes, + # case or URL encoding, and do not use unverified JWT claims. + if oauth_issuer != issuer: + raise OAuthTokenError( + "The server's oauth_issuer does not match the pgAdmin " + "login OIDC provider issuer." + ) + + return oauth_client, metadata, verify_tls + + +def _get_token_exchange_auth(oauth_client, metadata): + """Authenticate exchange using the pgAdmin login client's credentials.""" + client_id = oauth_client.client_id + client_secret = oauth_client.client_secret + + if ( + not isinstance(client_id, str) or + not client_id.strip() or + "\0" in client_id or + not isinstance(client_secret, str) or + not client_secret.strip() or + "\0" in client_secret + ): + raise OAuthTokenExchangeError( + "PostgreSQL token exchange requires the pgAdmin login " + "OIDC client to have a client ID and client secret." + ) + + client_kwargs = oauth_client.client_kwargs or {} + auth_method = client_kwargs.get( + "token_endpoint_auth_method", + metadata.get("token_endpoint_auth_method"), + ) + if auth_method is None: + auth_method = "client_secret_basic" + + if auth_method not in ("client_secret_basic", "client_secret_post"): + raise OAuthTokenExchangeError( + "PostgreSQL token exchange supports client_secret_basic " + "or client_secret_post authentication. Configure a " + "supported method for the pgAdmin login OIDC client." + ) + + # Only construct the Requests authentication object here. No token + # request is made and no login scopes or token callbacks are copied. + with OAuth2Session( + client_id=client_id, + client_secret=client_secret, + token_endpoint_auth_method=auth_method, + ) as client: + return client.client_auth(auth_method) + + +def _exchange_oauth_access_token( + client_id: str, + oauth_scope: Optional[str], + oauth_issuer: Optional[str], +) -> str: + """ + Exchange the current pgAdmin access token for a cluster access token. + + client_id is the cluster's OAuth client identifier, obtained from + the server registration's oauth_client_id connection parameter. + + The token endpoint is obtained from the OAuth provider used for the + current pgAdmin login. An explicitly configured token URL takes + precedence over the token_endpoint discovered through the provider + metadata. + + The current pgAdmin token is refreshed first when necessary. + + Return the exchanged bearer token. Raise OAuthTokenExchangeError + on failure. + + The exchanged token is not stored in the Flask session or cached. + """ + if not has_request_context(): + raise OAuthTokenExchangeError( + "Token exchange requires a pgAdmin login session." + ) + + if ( + not isinstance(client_id, str) or + not client_id.strip() or + "\0" in client_id + ): + raise OAuthTokenExchangeError( + "Set the server connection parameter oauth_client_id " + "to the target cluster identifier." + ) + + try: + oauth_client, metadata, verify_tls = _get_validated_oauth_client( + oauth_issuer, + ) + except OAuthTokenError as exc: + raise OAuthTokenExchangeError(str(exc)) from None + + exchange_auth = _get_token_exchange_auth(oauth_client, metadata) + + endpoint = ( + oauth_client.access_token_url or + metadata.get("token_endpoint") + ) + + if not isinstance(endpoint, str) or not endpoint: + raise OAuthTokenExchangeError( + "The OAuth provider has no token endpoint configured " + "or advertised in its metadata." + ) + + try: + parsed_endpoint = urlsplit(endpoint) + valid_endpoint = ( + parsed_endpoint.scheme == "https" and + bool(parsed_endpoint.hostname) and + parsed_endpoint.username is None and + parsed_endpoint.password is None and + not parsed_endpoint.fragment + ) + except ValueError: + valid_endpoint = False + + if not valid_endpoint: + raise OAuthTokenExchangeError( + "The OAuth provider token endpoint must be an HTTPS URL " + "without embedded credentials or a fragment." + ) + + subject_token = _get_pgadmin_oauth_token() + + if ( + not isinstance(subject_token, str) or + not subject_token.strip() or + "\0" in subject_token + ): + raise OAuthTokenExchangeError( + "No current pgAdmin OAuth access token is available. " + "Sign in to pgAdmin again." + ) + + try: + request_data = { + "audience": client_id, + "grant_type": ( + "urn:ietf:params:oauth:grant-type:token-exchange" + ), + "subject_token": subject_token, + "subject_token_type": ( + "urn:ietf:params:oauth:token-type:access_token" + ), + } + + scope = (oauth_scope or "").strip() + if scope: + request_data["scope"] = scope + + response = requests.post( + endpoint, + data=request_data, + headers={ + "Accept": "application/json", + "Content-Type": "application/x-www-form-urlencoded", + }, + auth=exchange_auth, + verify=verify_tls, + timeout=(5, 30), + allow_redirects=False, + ) + except requests.exceptions.SSLError: + raise OAuthTokenExchangeError( + "TLS verification failed when contacting the OAuth " + "token endpoint." + ) from None + except requests.exceptions.Timeout: + raise OAuthTokenExchangeError( + "The request to the OAuth token endpoint timed out." + ) from None + except requests.exceptions.RequestException: + raise OAuthTokenExchangeError( + "Could not contact the OAuth token endpoint." + ) from None + + with response: + status_code = response.status_code + + try: + token_data = response.json() + except ValueError: + token_data = None + + has_oauth_error = ( + isinstance(token_data, dict) and + "error" in token_data + ) + + if status_code != 200 or has_oauth_error: + error_code = ( + token_data.get("error") + if isinstance(token_data, dict) + else None + ) + + # Include only a short error identifier. Never expose the + # provider response, descriptions, or token values. + if ( + not isinstance(error_code, str) or + re.fullmatch( + r"[A-Za-z0-9_-]{1,64}", + error_code, + ) is None + ): + error_code = "unknown_error" + + raise OAuthTokenExchangeError( + "OAuth token exchange failed " + "(HTTP {0}, error={1}).".format( + status_code, + error_code, + ) + ) + + if not isinstance(token_data, dict): + raise OAuthTokenExchangeError( + "The OAuth provider returned an invalid token exchange " + "response." + ) + + access_token = token_data.get("access_token") + token_type = token_data.get("token_type") + + if ( + not isinstance(access_token, str) or + not access_token.strip() or + "\0" in access_token + ): + raise OAuthTokenExchangeError( + "The OAuth token exchange response contains no valid " + "access token." + ) + + if ( + not isinstance(token_type, str) or + token_type.lower() != "bearer" + ): + raise OAuthTokenExchangeError( + "The OAuth token exchange response is not a bearer token." + ) + + return access_token + + +def get_postgres_oauth_token( + mode: str, + client_id: Optional[str], + oauth_scope: Optional[str], + oauth_issuer: Optional[str], +) -> str: + """Return a PostgreSQL OAuth token using the selected token mode.""" + if not has_request_context(): + logger.debug( + "PostgreSQL OAuth token retrieval outside a request context.", + stack_info=True, + ) + raise OAuthRequestContextRequired( + "PostgreSQL OAuth reconnect requires an authenticated " + "pgAdmin request. Reconnect through pgAdmin." + ) + + if mode == "direct": + _get_validated_oauth_client(oauth_issuer) + token = _get_pgadmin_oauth_token() + + if ( + not isinstance(token, str) or + not token.strip() or + "\0" in token + ): + raise OAuthTokenError( + "No current pgAdmin OAuth access token is available. " + "Sign in to pgAdmin again." + ) + + return token + + if mode == "exchange": + if ( + not isinstance(client_id, str) or + not client_id.strip() or + "\0" in client_id + ): + raise OAuthTokenError( + "Set the server connection parameter oauth_client_id " + "when using OAuth token exchange." + ) + + return _exchange_oauth_access_token( + client_id, oauth_scope, oauth_issuer, + ) + + raise OAuthTokenError( + 'Invalid pgAdmin OAuth token mode "{0}".'.format(mode) + ) diff --git a/web/pgadmin/utils/tests/test_pg_oauth2.py b/web/pgadmin/utils/tests/test_pg_oauth2.py new file mode 100644 index 00000000000..e414853d190 --- /dev/null +++ b/web/pgadmin/utils/tests/test_pg_oauth2.py @@ -0,0 +1,1984 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +"""Tests for pgAdmin PostgreSQL OAuth bearer-token support.""" + +import base64 +import ctypes +import unittest +from urllib.parse import parse_qs +from unittest.mock import MagicMock, Mock, patch, create_autospec +from pgadmin.utils.route import BaseTestGenerator + +from contextlib import ExitStack + +from flask import Flask, session +from pgadmin.utils import pg_oauth2 + +from authlib.integrations.flask_client.apps import FlaskOAuth2App + +ISSUER = 'https://keycloak.example.test/realms/postgres' + + +class _OAuthUnitTestMixin: + """Run each OAuth unit test as a separate regression scenario.""" + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + + cls.scenarios = [ + (name[5:], {'case_method': name}) + for name in sorted(dir(cls)) + if name.startswith('case_') and callable(getattr(cls, name)) + ] + + def setUp(self): + # These tests mock database connections. + unittest.TestCase.setUp(self) + + def runTest(self): + getattr(self, self.case_method)() + + +class TestOAuthTokenContext( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test propagation of tokens through ContextVar.""" + + def setUp(self): + self.token_handle = pg_oauth2._oauth_token.set(None) + + def tearDown(self): + pg_oauth2._oauth_token.reset(self.token_handle) + + def case_context_sets_and_restores_token(self): + self.assertIsNone(pg_oauth2._oauth_token.get()) + + with pg_oauth2.oauth_token_context('access-token'): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'access-token' + ) + + self.assertIsNone(pg_oauth2._oauth_token.get()) + + def case_context_restores_existing_token(self): + existing_handle = pg_oauth2._oauth_token.set( + 'existing-token' + ) + + try: + with pg_oauth2.oauth_token_context('temporary-token'): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'temporary-token' + ) + + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'existing-token' + ) + finally: + pg_oauth2._oauth_token.reset(existing_handle) + + def case_nested_contexts_restore_previous_token(self): + with pg_oauth2.oauth_token_context('outer-token'): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'outer-token' + ) + + with pg_oauth2.oauth_token_context('inner-token'): + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'inner-token' + ) + + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'outer-token' + ) + + self.assertIsNone(pg_oauth2._oauth_token.get()) + + def case_context_restores_token_when_exception_is_raised(self): + with self.assertRaisesRegex(RuntimeError, 'test error'): + with pg_oauth2.oauth_token_context('access-token'): + raise RuntimeError('test error') + + self.assertIsNone(pg_oauth2._oauth_token.get()) + + +class TestOAuthHook( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test the ctypes libpq OAuth callback.""" + + def setUp(self): + self.token_handle = pg_oauth2._oauth_token.set(None) + + self.previous_hook = pg_oauth2._previous_hook + pg_oauth2._previous_hook = None + + self.previous_token_buffers = dict( + pg_oauth2._token_buffers + ) + pg_oauth2._token_buffers.clear() + + def tearDown(self): + pg_oauth2._token_buffers.clear() + pg_oauth2._token_buffers.update( + self.previous_token_buffers + ) + + pg_oauth2._previous_hook = self.previous_hook + pg_oauth2._oauth_token.reset(self.token_handle) + + @staticmethod + def _make_request(): + request = pg_oauth2.PGoauthBearerRequest() + request_pointer = ctypes.pointer(request) + data = ctypes.cast( + request_pointer, + ctypes.c_void_p + ) + + return request, request_pointer, data + + def case_non_bearer_request_delegates_to_previous_hook(self): + previous_hook = Mock(return_value=73) + pg_oauth2._previous_hook = previous_hook + + result = pg_oauth2._oauth_hook( + 0, + None, + None + ) + + self.assertEqual(result, 73) + previous_hook.assert_called_once_with( + 0, + None, + None + ) + + def case_non_bearer_request_returns_zero_without_previous_hook(self): + result = pg_oauth2._oauth_hook( + 0, + None, + None + ) + + self.assertEqual(result, 0) + + def case_missing_token_delegates_to_previous_hook(self): + previous_hook = Mock(return_value=41) + pg_oauth2._previous_hook = previous_hook + + result = pg_oauth2._oauth_hook( + pg_oauth2.PQAUTHDATA_OAUTH_BEARER_TOKEN, + None, + None + ) + + self.assertEqual(result, 41) + previous_hook.assert_called_once_with( + pg_oauth2.PQAUTHDATA_OAUTH_BEARER_TOKEN, + None, + None + ) + + def case_missing_token_returns_zero_without_previous_hook(self): + result = pg_oauth2._oauth_hook( + pg_oauth2.PQAUTHDATA_OAUTH_BEARER_TOKEN, + None, + None + ) + + self.assertEqual(result, 0) + + def case_bearer_request_receives_token(self): + previous_hook = Mock(return_value=51) + pg_oauth2._previous_hook = previous_hook + + request, request_pointer, data = self._make_request() + + with pg_oauth2.oauth_token_context('test-access-token'): + result = pg_oauth2._oauth_hook( + pg_oauth2.PQAUTHDATA_OAUTH_BEARER_TOKEN, + None, + data + ) + + self.assertEqual(result, 1) + self.assertEqual(request.token, b'test-access-token') + self.assertIsNone(request.async_) + self.assertIsNotNone(request.cleanup) + + request_address = ctypes.addressof(request) + self.assertIn(request_address, pg_oauth2._token_buffers) + + previous_hook.assert_not_called() + + # Keep the ctypes pointer alive for the whole test. + self.assertIsNotNone(request_pointer) + + def case_cleanup_removes_retained_token_buffer(self): + request, request_pointer, data = self._make_request() + + with pg_oauth2.oauth_token_context( + 'test-access-token'): + result = pg_oauth2._oauth_hook( + pg_oauth2.PQAUTHDATA_OAUTH_BEARER_TOKEN, + None, + data + ) + + self.assertEqual(result, 1) + + request_address = ctypes.addressof(request) + + self.assertIn( + request_address, + pg_oauth2._token_buffers + ) + + pg_oauth2._oauth_cleanup(None, data) + + self.assertNotIn( + request_address, + pg_oauth2._token_buffers + ) + + # Keep request_pointer alive until cleanup has completed. + self.assertIsNotNone(request_pointer) + + def case_previous_hook_exception_returns_minus_one(self): + """ + Never allow an exception from the previous libpq hook to cross + the native callback boundary. + """ + previous_hook = Mock( + side_effect=RuntimeError('previous hook failed') + ) + pg_oauth2._previous_hook = previous_hook + + result = pg_oauth2._oauth_hook( + 0, + None, + None + ) + + self.assertEqual(result, -1) + previous_hook.assert_called_once_with( + 0, + None, + None + ) + + def case_bearer_hook_internal_failure_returns_minus_one(self): + """ + Never allow an exception raised while populating the bearer + request to cross the native callback boundary. + """ + previous_hook = Mock(return_value=51) + pg_oauth2._previous_hook = previous_hook + + request, request_pointer, data = self._make_request() + + # An access token must be a string. This deliberately supplies a + # malformed value so token.encode() raises inside the callback. + with pg_oauth2.oauth_token_context(123): + result = pg_oauth2._oauth_hook( + pg_oauth2.PQAUTHDATA_OAUTH_BEARER_TOKEN, + None, + data + ) + + self.assertEqual(result, -1) + self.assertIsNone(request.token) + self.assertNotIn( + ctypes.addressof(request), + pg_oauth2._token_buffers + ) + + # Once pgAdmin chooses to handle a bearer request, an internal + # failure must not invoke the previous hook afterward. + previous_hook.assert_not_called() + + # Keep the ctypes allocation alive for the complete callback. + self.assertIsNotNone(request_pointer) + + def case_none_context_clears_and_restores_existing_token(self): + with pg_oauth2.oauth_token_context('existing-token'): + with pg_oauth2.oauth_token_context(None): + self.assertIsNone(pg_oauth2._oauth_token.get()) + + self.assertEqual( + pg_oauth2._oauth_token.get(), + 'existing-token', + ) + + self.assertIsNone(pg_oauth2._oauth_token.get()) + + +class TestGetCurrentOAuthClient( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test lookup of the Authlib client for the current OAuth provider.""" + + def setUp(self): + self.app = Flask(__name__) + self.app.config.update( + SECRET_KEY='pg-oauth2-unit-test-secret', + TESTING=True, + ) + + @patch('pgadmin.authenticate.get_auth_sources') + def case_unknown_provider_raises_error(self, get_auth_sources): + oauth_source = MagicMock() + oauth_source.oauth2_clients = {} + get_auth_sources.return_value = oauth_source + + with self.app.test_request_context('/'): + session['oauth2_provider'] = 'unknown-provider' + + with self.assertRaises(pg_oauth2.OAuthTokenError): + pg_oauth2._get_current_oauth_client() + + @patch('pgadmin.authenticate.get_auth_sources') + def case_returns_registered_oauth_client(self, get_auth_sources): + oauth_client = MagicMock() + oauth_source = MagicMock() + oauth_source.oauth2_clients = { + 'keycloak': oauth_client, + } + get_auth_sources.return_value = oauth_source + + with self.app.test_request_context('/'): + session['oauth2_provider'] = 'keycloak' + + result = pg_oauth2._get_current_oauth_client() + + self.assertIs(result, oauth_client) + + +class TestGetPgAdminOAuthToken( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test pgAdmin access-token retrieval and refresh.""" + + PROVIDER_NAME = 'keycloak' + + def setUp(self): + self.app = Flask(__name__) + self.app.config.update( + SECRET_KEY='pg-oauth2-unit-test-secret', + TESTING=True, + ) + + self.oauth_source = Mock() + self.oauth_source.oauth2_config = { + self.PROVIDER_NAME: { + "OAUTH2_CLIENT_AUTH_METHOD": "client_secret", + }, + } + + source_patch = patch( + "pgadmin.authenticate.get_auth_sources", + return_value=self.oauth_source, + ) + source_patch.start() + self.addCleanup(source_patch.stop) + + def _make_oauth_client(self): + return create_autospec( + FlaskOAuth2App, + instance=True, + spec_set=True + ) + + def case_returns_none_without_request_context(self): + self.assertIsNone( + pg_oauth2._get_pgadmin_oauth_token() + ) + + def case_returns_none_without_oauth_token(self): + with self.app.test_request_context('/'): + self.assertIsNone( + pg_oauth2._get_pgadmin_oauth_token() + ) + + def case_returns_none_when_token_data_is_not_dictionary(self): + with self.app.test_request_context('/'): + session['oauth2_token'] = 'not-a-dictionary' + + self.assertIsNone( + pg_oauth2._get_pgadmin_oauth_token() + ) + + def case_returns_none_when_access_token_is_missing(self): + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'refresh_token': 'refresh-token', + } + + self.assertIsNone( + pg_oauth2._get_pgadmin_oauth_token() + ) + + def case_returns_current_unexpired_access_token(self): + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'current-access-token', + 'refresh_token': 'refresh-token', + 'expires_at': 2000, + } + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertEqual( + result, + 'current-access-token' + ) + + @patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client' + ) + def case_refreshes_expired_access_token(self, get_oauth_client): + oauth_client = self._make_oauth_client() + oauth_client.fetch_access_token.return_value = { + 'access_token': 'refreshed-access-token', + 'refresh_token': 'new-refresh-token', + } + + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'refresh_token': 'original-refresh-token', + 'expires_at': 999, + 'token_type': 'Bearer', + } + session['oauth2_provider'] = self.PROVIDER_NAME + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertEqual(result, 'refreshed-access-token') + + oauth_client.fetch_access_token.assert_called_once_with( + grant_type='refresh_token', + refresh_token='original-refresh-token' + ) + + @patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client' + ) + def case_refresh_failure_returns_none(self, get_oauth_client): + oauth_client = self._make_oauth_client() + oauth_client.fetch_access_token.side_effect = RuntimeError( + 'token endpoint unavailable' + ) + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'refresh_token': 'refresh-token', + 'expires_at': 999, + 'token_type': 'Bearer', + } + session['oauth2_provider'] = self.PROVIDER_NAME + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertIsNone(result) + oauth_client.fetch_access_token.assert_called_once_with( + grant_type='refresh_token', + refresh_token='refresh-token' + ) + + @patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client' + ) + def case_expired_token_without_refresh_token_returns_none( + self, get_oauth_client): + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'expires_at': 999, + 'token_type': 'Bearer', + } + session['oauth2_provider'] = self.PROVIDER_NAME + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertIsNone(result) + get_oauth_client.assert_not_called() + + @patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client' + ) + def case_refreshes_access_token_within_expiry_safety_window( + self, get_oauth_client): + """ + Refresh a token that has not expired yet but will expire within + the 30-second safety window. + """ + oauth_client = self._make_oauth_client() + oauth_client.fetch_access_token.return_value = { + 'access_token': 'refreshed-access-token', + 'refresh_token': 'refreshed-refresh-token', + 'expires_at': 3000, + 'token_type': 'Bearer', + } + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'nearly-expired-access-token', + 'refresh_token': 'original-refresh-token', + 'expires_at': 1020, + 'token_type': 'Bearer', + } + session['oauth2_provider'] = self.PROVIDER_NAME + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertEqual(result, 'refreshed-access-token') + self.assertEqual( + session['oauth2_token']['access_token'], + 'refreshed-access-token' + ) + + oauth_client.fetch_access_token.assert_called_once_with( + grant_type='refresh_token', + refresh_token='original-refresh-token' + ) + + def case_returns_access_token_without_expiration(self): + """ + Return an access token when the provider did not supply expires_at. + """ + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'access-token-without-expiration', + 'refresh_token': 'refresh-token', + 'token_type': 'Bearer', + } + + with patch( + 'pgadmin.utils.pg_oauth2.' + '_refresh_pgadmin_oauth_token' + ) as refresh_mock: + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertEqual( + result, + 'access-token-without-expiration' + ) + refresh_mock.assert_not_called() + + @patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client' + ) + def case_refresh_preserves_existing_refresh_token( + self, get_oauth_client): + """ + Preserve the old refresh token when the provider returns only a + new access token. + """ + oauth_client = self._make_oauth_client() + oauth_client.fetch_access_token.return_value = { + 'access_token': 'refreshed-access-token', + 'expires_at': 3000, + 'token_type': 'Bearer', + } + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'refresh_token': 'original-refresh-token', + 'expires_at': 900, + 'token_type': 'Bearer', + } + session['oauth2_provider'] = self.PROVIDER_NAME + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertEqual(result, 'refreshed-access-token') + self.assertEqual( + session['oauth2_token']['refresh_token'], + 'original-refresh-token' + ) + + oauth_client.fetch_access_token.assert_called_once_with( + grant_type='refresh_token', + refresh_token='original-refresh-token' + ) + + def case_expired_token_without_provider_returns_none(self): + """ + An expired token cannot be refreshed when the OAuth provider + identity is missing from the session. + """ + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'refresh_token': 'refresh-token', + 'expires_at': 900, + 'token_type': 'Bearer', + } + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertIsNone(result) + + @patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client', + side_effect=pg_oauth2.OAuthTokenError( + 'Unknown OAuth provider.' + ) + ) + def case_expired_token_returns_none_when_client_lookup_fails( + self, get_oauth_client): + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'refresh_token': 'refresh-token', + 'expires_at': 900, + 'token_type': 'Bearer', + } + session['oauth2_provider'] = 'unknown-provider' + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertIsNone(result) + get_oauth_client.assert_called_once_with() + + @patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client' + ) + def case_refresh_returns_none_for_non_dictionary_response( + self, get_oauth_client): + """ + Reject a malformed refresh response instead of attempting to use + or store it as an OAuth token. + """ + oauth_client = self._make_oauth_client() + oauth_client.fetch_access_token.return_value = ( + 'invalid-token-response' + ) + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'refresh_token': 'original-refresh-token', + 'expires_at': 900, + 'token_type': 'Bearer', + } + session['oauth2_provider'] = self.PROVIDER_NAME + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertIsNone(result) + self.assertEqual( + session['oauth2_token']['access_token'], + 'expired-access-token' + ) + + oauth_client.fetch_access_token.assert_called_once_with( + grant_type='refresh_token', + refresh_token='original-refresh-token' + ) + + @patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client' + ) + def case_refresh_returns_none_when_access_token_is_missing( + self, get_oauth_client): + """ + Reject a syntactically valid refresh response that does not contain + a new access token. + """ + oauth_client = self._make_oauth_client() + oauth_client.fetch_access_token.return_value = { + 'refresh_token': 'new-refresh-token', + 'expires_at': 3000, + 'token_type': 'Bearer', + } + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'refresh_token': 'original-refresh-token', + 'expires_at': 900, + 'token_type': 'Bearer', + } + session['oauth2_provider'] = self.PROVIDER_NAME + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000 + ): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertIsNone(result) + self.assertEqual( + session['oauth2_token']['access_token'], + 'expired-access-token' + ) + + oauth_client.fetch_access_token.assert_called_once_with( + grant_type='refresh_token', + refresh_token='original-refresh-token' + ) + + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + def case_refresh_does_not_reuse_previous_token_expiry( + self, get_oauth_client): + # The payload is {"exp": 3000}. Signature validation is not + # involved: the implementation only reads an expiration hint. + refreshed_access_token = ( + 'header.eyJleHAiOiAzMDAwfQ.signature' + ) + + oauth_client = self._make_oauth_client() + oauth_client.fetch_access_token.return_value = { + 'access_token': refreshed_access_token, + 'token_type': 'Bearer', + } + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + session['oauth2_provider'] = self.PROVIDER_NAME + session['oauth2_token'] = { + 'access_token': 'expired-access-token', + 'refresh_token': 'original-refresh-token', + 'expires_at': 999, + 'expires_in': 60, + 'token_type': 'Bearer', + } + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000): + first_token = pg_oauth2._get_pgadmin_oauth_token() + second_token = pg_oauth2._get_pgadmin_oauth_token() + + self.assertEqual(first_token, refreshed_access_token) + self.assertEqual(second_token, refreshed_access_token) + self.assertEqual( + pg_oauth2._get_token_expiry(session['oauth2_token']), + 3000, + ) + self.assertEqual( + session['oauth2_token']['refresh_token'], + 'original-refresh-token', + ) + self.assertNotIn('expires_in', session['oauth2_token']) + + oauth_client.fetch_access_token.assert_called_once_with( + grant_type='refresh_token', + refresh_token='original-refresh-token', + ) + + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + def case_workload_identity_refresh_is_explicitly_rejected( + self, get_oauth_client): + self.oauth_source.oauth2_config[self.PROVIDER_NAME][ + 'OAUTH2_CLIENT_AUTH_METHOD' + ] = 'workload_identity' + + oauth_client = self._make_oauth_client() + get_oauth_client.return_value = oauth_client + + # Cover both expiration and the early-refresh window. + for expires_at in (999, 1020): + with self.subTest(expires_at=expires_at): + original_token = { + 'access_token': 'existing-access-token', + 'refresh_token': 'existing-refresh-token', + 'expires_at': expires_at, + 'token_type': 'Bearer', + } + + with self.app.test_request_context('/'): + session['oauth2_provider'] = self.PROVIDER_NAME + session['oauth2_token'] = dict(original_token) + session.modified = False + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000): + with self.assertRaisesRegex( + pg_oauth2.OAuthRefreshUnsupportedError, + 'Automatic OAuth token refresh for ' + 'workload-identity login is not supported'): + pg_oauth2._get_pgadmin_oauth_token() + + self.assertEqual( + session['oauth2_token'], + original_token, + ) + self.assertFalse(session.modified) + + oauth_client.fetch_access_token.assert_not_called() + + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + def case_workload_identity_can_use_token_before_refresh_window( + self, get_oauth_client): + self.oauth_source.oauth2_config[self.PROVIDER_NAME][ + 'OAUTH2_CLIENT_AUTH_METHOD' + ] = 'workload_identity' + + original_token = { + 'access_token': 'current-access-token', + 'refresh_token': 'existing-refresh-token', + 'expires_at': 2000, + 'token_type': 'Bearer', + } + + with self.app.test_request_context('/'): + session['oauth2_provider'] = self.PROVIDER_NAME + session['oauth2_token'] = dict(original_token) + session.modified = False + + with patch( + 'pgadmin.utils.pg_oauth2.time.time', + return_value=1000): + result = pg_oauth2._get_pgadmin_oauth_token() + + self.assertEqual(result, 'current-access-token') + self.assertEqual(session['oauth2_token'], original_token) + self.assertFalse(session.modified) + + get_oauth_client.assert_not_called() + + +class TestExchangeOAuthAccessToken( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test OAuth access-token exchange for PostgreSQL.""" + + TOKEN_ENDPOINT = ( + 'https://keycloak.example.test/realms/postgres/' + 'protocol/openid-connect/token' + ) + + DISCOVERED_TOKEN_ENDPOINT = ( + 'https://keycloak.example.test/discovered/token' + ) + + def setUp(self): + self.app = Flask(__name__) + self.app.config.update( + SECRET_KEY='pg-oauth2-unit-test-secret', + TESTING=True, + ) + + @staticmethod + def _make_oauth_client( + access_token_url=None, + verify=True): + oauth_client = MagicMock() + oauth_client.client_id = 'pgadmin-client' + oauth_client.client_secret = 'pgadmin-client-secret' + oauth_client.access_token_url = access_token_url + oauth_client.client_kwargs = { + 'verify': verify, + } + oauth_client.load_server_metadata.return_value = { + 'issuer': ISSUER, + 'token_endpoint': ( + TestExchangeOAuthAccessToken.DISCOVERED_TOKEN_ENDPOINT + ), + } + + return oauth_client + + @staticmethod + def _make_response(status_code=200, token_data=None): + response = MagicMock() + response.status_code = status_code + response.json.return_value = token_data + response.__enter__.return_value = response + response.__exit__.return_value = False + + return response + + def case_exchange_requires_request_context(self): + with self.assertRaisesRegex( + pg_oauth2.OAuthTokenExchangeError, + 'pgAdmin login session'): + pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_returns_access_token( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + + oauth_client = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT + ) + get_oauth_client.return_value = oauth_client + + post.return_value = self._make_response( + token_data={ + 'access_token': 'exchanged-access-token', + 'token_type': 'Bearer', + } + ) + + with self.app.test_request_context('/'): + result = pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + self.assertEqual(result, 'exchanged-access-token') + + get_token.assert_called_once_with() + get_oauth_client.assert_called_once_with() + post.assert_called_once() + + args, kwargs = post.call_args + + self.assertEqual(args[0], self.TOKEN_ENDPOINT) + + self.assertEqual( + kwargs['data'], + { + 'audience': 'postgres-cluster', + 'grant_type': ( + 'urn:ietf:params:oauth:grant-type:token-exchange' + ), + 'subject_token': 'subject-access-token', + 'subject_token_type': ( + 'urn:ietf:params:oauth:token-type:access_token' + ), + 'scope': 'openid', + } + ) + + self.assertEqual( + kwargs['headers'], + { + 'Accept': 'application/json', + 'Content-Type': 'application/x-www-form-urlencoded', + } + ) + + prepared = pg_oauth2.requests.Request( + 'POST', + args[0], + data=kwargs['data'], + headers=kwargs['headers'], + auth=kwargs['auth'], + ).prepare() + + expected_credentials = base64.b64encode( + b'pgadmin-client:pgadmin-client-secret' + ).decode('ascii') + + self.assertEqual( + prepared.headers['Authorization'], + 'Basic ' + expected_credentials, + ) + self.assertTrue(kwargs['verify']) + self.assertEqual(kwargs['timeout'], (5, 30)) + self.assertFalse(kwargs['allow_redirects']) + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_client_secret_post_authenticates_pgadmin_for_database_audience( + self, get_token, get_oauth_client, post): + oauth_client = self._make_oauth_client(self.TOKEN_ENDPOINT) + oauth_client.client_kwargs['token_endpoint_auth_method'] = ( + 'client_secret_post' + ) + oauth_client.client_secret = 'secret+with&reserved=characters' + get_oauth_client.return_value = oauth_client + get_token.return_value = 'subject-access-token' + post.return_value = self._make_response(token_data={ + 'access_token': 'exchanged-access-token', + 'token_type': 'Bearer', + }) + + with self.app.test_request_context('/'): + result = pg_oauth2.get_postgres_oauth_token( + 'exchange', 'postgres-cluster', 'openid', ISSUER, + ) + + self.assertEqual(result, 'exchanged-access-token') + get_token.assert_called_once_with() + post.assert_called_once() + args, kwargs = post.call_args + + # Exercise the real Authlib authentication object. Inspecting + # request_data alone would miss the credentials it adds later. + prepared = pg_oauth2.requests.Request( + 'POST', args[0], data=kwargs['data'], + headers=kwargs['headers'], auth=kwargs['auth'], + ).prepare() + body = prepared.body + if isinstance(body, bytes): + body = body.decode('utf-8') + fields = parse_qs(body) + + self.assertEqual(prepared.url, self.TOKEN_ENDPOINT) + self.assertNotIn('Authorization', prepared.headers) + self.assertEqual(fields['client_id'], ['pgadmin-client']) + self.assertEqual( + fields['client_secret'], [oauth_client.client_secret], + ) + self.assertEqual(fields['audience'], ['postgres-cluster']) + self.assertEqual(fields['subject_token'], ['subject-access-token']) + self.assertEqual(fields['scope'], ['openid']) + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._refresh_pgadmin_oauth_token') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + def _assert_exchange_auth_rejected( + self, oauth_client, message, get_client, get_token, + refresh_token, post): + get_client.return_value = oauth_client + get_token.return_value = 'subject-access-token' + + with self.app.test_request_context('/'): + with self.assertRaisesRegex( + pg_oauth2.OAuthTokenExchangeError, message): + pg_oauth2.get_postgres_oauth_token( + 'exchange', 'postgres-cluster', 'openid', ISSUER, + ) + + get_token.assert_not_called() + refresh_token.assert_not_called() + post.assert_not_called() + + def case_exchange_rejects_missing_pgadmin_client_credentials(self): + for field in ('client_id', 'client_secret'): + for value in (None, '', ' '): + with self.subTest(field=field, value=value): + oauth_client = self._make_oauth_client(self.TOKEN_ENDPOINT) + setattr(oauth_client, field, value) + self._assert_exchange_auth_rejected( + oauth_client, 'client ID and client secret', + ) + + def case_exchange_rejects_unauthenticated_client(self): + oauth_client = self._make_oauth_client(self.TOKEN_ENDPOINT) + oauth_client.client_kwargs['token_endpoint_auth_method'] = 'none' + self._assert_exchange_auth_rejected( + oauth_client, + 'supports client_secret_basic or client_secret_post', + ) + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_uses_discovered_token_endpoint( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + + oauth_client = self._make_oauth_client( + access_token_url=None + ) + get_oauth_client.return_value = oauth_client + + post.return_value = self._make_response( + token_data={ + 'access_token': 'exchanged-access-token', + 'token_type': 'Bearer', + } + ) + + with self.app.test_request_context('/'): + result = pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + self.assertEqual(result, 'exchanged-access-token') + + oauth_client.load_server_metadata.assert_called_once_with() + post.assert_called_once() + + args, _ = post.call_args + self.assertEqual( + args[0], + self.DISCOVERED_TOKEN_ENDPOINT + ) + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_fails_when_provider_metadata_cannot_be_loaded( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + + oauth_client = self._make_oauth_client( + access_token_url=None + ) + oauth_client.load_server_metadata.side_effect = RuntimeError( + 'metadata unavailable' + ) + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + with self.assertRaises( + pg_oauth2.OAuthTokenExchangeError): + pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + oauth_client.load_server_metadata.assert_called_once_with() + post.assert_not_called() + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_rejects_disabled_tls_verification( + self, get_token, get_oauth_client, post): + oauth_client = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT, + verify=False, + ) + get_oauth_client.return_value = oauth_client + + with self.app.test_request_context('/'): + with self.assertRaises( + pg_oauth2.OAuthTokenExchangeError): + pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + get_oauth_client.assert_called_once_with() + get_token.assert_not_called() + oauth_client.load_server_metadata.assert_not_called() + post.assert_not_called() + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_translates_request_timeout( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + + get_oauth_client.return_value = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT + ) + + post.side_effect = pg_oauth2.requests.exceptions.Timeout( + 'request timed out' + ) + + with self.app.test_request_context('/'): + with self.assertRaisesRegex( + pg_oauth2.OAuthTokenExchangeError, + 'timed out'): + pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + post.assert_called_once() + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_rejects_oauth_error_response( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + + get_oauth_client.return_value = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT + ) + + post.return_value = self._make_response( + status_code=400, + token_data={ + 'error': 'invalid_grant', + } + ) + + with self.app.test_request_context('/'): + with self.assertRaisesRegex( + pg_oauth2.OAuthTokenExchangeError, + 'invalid_grant'): + pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + post.assert_called_once() + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_rejects_non_dictionary_response( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + + get_oauth_client.return_value = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT + ) + + post.return_value = self._make_response( + token_data='invalid-response' + ) + + with self.app.test_request_context('/'): + with self.assertRaises( + pg_oauth2.OAuthTokenExchangeError): + pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + post.assert_called_once() + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_rejects_invalid_access_token( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + + get_oauth_client.return_value = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT + ) + + invalid_tokens = ( + None, + '', + ' ', + 'invalid\0token', + ) + + with self.app.test_request_context('/'): + for access_token in invalid_tokens: + with self.subTest(access_token=access_token): + post.return_value = self._make_response( + token_data={ + 'access_token': access_token, + 'token_type': 'Bearer', + } + ) + + with self.assertRaises( + pg_oauth2.OAuthTokenExchangeError): + pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + self.assertEqual( + post.call_count, + len(invalid_tokens) + ) + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_rejects_non_bearer_token( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + + get_oauth_client.return_value = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT + ) + + post.return_value = self._make_response( + token_data={ + 'access_token': 'exchanged-access-token', + 'token_type': 'DPoP', + } + ) + + with self.app.test_request_context('/'): + with self.assertRaises( + pg_oauth2.OAuthTokenExchangeError): + pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + 'openid', + ISSUER, + ) + + post.assert_called_once() + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_omits_scope_when_not_configured( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + get_oauth_client.return_value = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT, + ) + post.return_value = self._make_response( + token_data={ + 'access_token': 'exchanged-access-token', + 'token_type': 'Bearer', + }, + ) + + for scope in (None, '', ' '): + with self.subTest(scope=scope): + post.reset_mock() + + with self.app.test_request_context('/'): + result = pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + scope, + ISSUER, + ) + + self.assertEqual(result, 'exchanged-access-token') + post.assert_called_once() + self.assertNotIn( + 'scope', + post.call_args.kwargs['data'], + ) + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._get_current_oauth_client') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def case_exchange_uses_configured_scopes( + self, get_token, get_oauth_client, post): + get_token.return_value = 'subject-access-token' + get_oauth_client.return_value = self._make_oauth_client( + access_token_url=self.TOKEN_ENDPOINT, + ) + post.return_value = self._make_response( + token_data={ + 'access_token': 'exchanged-access-token', + 'token_type': 'Bearer', + }, + ) + + cases = ( + ('openid', 'openid'), + ('openid database.read', 'openid database.read'), + (' openid ', 'openid'), + ) + + for configured_scope, expected_scope in cases: + with self.subTest(scope=configured_scope): + post.reset_mock() + + with self.app.test_request_context('/'): + result = pg_oauth2._exchange_oauth_access_token( + 'postgres-cluster', + configured_scope, + ISSUER, + ) + + self.assertEqual(result, 'exchanged-access-token') + post.assert_called_once() + self.assertEqual( + post.call_args.kwargs['data']['scope'], + expected_scope, + ) + + +class TestGetPostgresOAuthToken( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test selection of the PostgreSQL OAuth token source.""" + + def setUp(self): + self.app = Flask(__name__) + self.app.config.update( + SECRET_KEY='pg-oauth2-unit-test-secret', + TESTING=True, + ) + + request_context = self.app.test_request_context('/') + request_context.push() + self.addCleanup(request_context.pop) + + self.oauth_client = TestExchangeOAuthAccessToken._make_oauth_client( + TestExchangeOAuthAccessToken.TOKEN_ENDPOINT, + ) + + client_patch = patch( + 'pgadmin.utils.pg_oauth2._get_current_oauth_client', + return_value=self.oauth_client, + ) + client_patch.start() + self.addCleanup(client_patch.stop) + + @patch('pgadmin.utils.pg_oauth2.requests.post') + @patch('pgadmin.utils.pg_oauth2._refresh_pgadmin_oauth_token') + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + def _assert_issuer_rejected( + self, issuer, message, get_token, refresh_token, post): + get_token.return_value = 'subject-access-token' + post.return_value = TestExchangeOAuthAccessToken._make_response( + token_data={ + 'access_token': 'exchanged-access-token', + 'token_type': 'Bearer', + }, + ) + + # Leave issuer validation and the exchange implementation real. + for mode in ('direct', 'exchange'): + with self.subTest(mode=mode): + with self.assertRaisesRegex( + pg_oauth2.OAuthTokenError, message): + pg_oauth2.get_postgres_oauth_token( + mode, 'postgres-cluster', 'openid', issuer, + ) + + get_token.assert_not_called() + refresh_token.assert_not_called() + post.assert_not_called() + + def case_missing_configured_issuer_is_rejected_before_token_access(self): + for issuer in (None, '', ' '): + with self.subTest(issuer=issuer): + self._assert_issuer_rejected(issuer, 'Set oauth_issuer') + + def case_missing_metadata_issuer_is_rejected_before_token_access(self): + metadata = self.oauth_client.load_server_metadata.return_value + metadata.pop('issuer') + self._assert_issuer_rejected(ISSUER, 'metadata contains no issuer') + + for issuer in (None, '', ' '): + with self.subTest(metadata_issuer=issuer): + metadata['issuer'] = issuer + self._assert_issuer_rejected( + ISSUER, 'metadata contains no issuer', + ) + + def case_mismatched_issuer_is_rejected_before_token_access(self): + # A trailing slash is also a mismatch; issuer IDs are exact. + for issuer in ('https://other.example.test', ISSUER + '/'): + with self.subTest(issuer=issuer): + self._assert_issuer_rejected(issuer, 'does not match') + + def case_unavailable_metadata_is_rejected_before_token_access(self): + self.oauth_client.load_server_metadata.side_effect = RuntimeError( + 'OIDC metadata unavailable', + ) + self._assert_issuer_rejected(ISSUER, 'Could not load') + + @patch( + 'pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token', + return_value='pgadmin-access-token' + ) + def case_direct_mode_returns_pgadmin_access_token(self, get_token): + result = pg_oauth2.get_postgres_oauth_token( + 'direct', + 'postgres-cluster', + None, + ISSUER, + ) + + self.assertEqual(result, 'pgadmin-access-token') + get_token.assert_called_once_with() + + @patch( + 'pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token', + return_value=None + ) + def case_direct_mode_raises_when_pgadmin_token_is_missing( + self, get_token): + with self.assertRaisesRegex( + pg_oauth2.OAuthTokenError, + 'No current pgAdmin OAuth access token is available'): + pg_oauth2.get_postgres_oauth_token( + 'direct', + 'postgres-cluster', + None, + ISSUER, + ) + + get_token.assert_called_once_with() + + @patch( + 'pgadmin.utils.pg_oauth2._exchange_oauth_access_token', + return_value='exchanged-access-token' + ) + def case_exchange_mode_returns_exchanged_access_token(self, exchange): + result = pg_oauth2.get_postgres_oauth_token( + 'exchange', + 'postgres-cluster', + 'openid', + ISSUER, + ) + + self.assertEqual(result, 'exchanged-access-token') + exchange.assert_called_once_with('postgres-cluster', 'openid', ISSUER) + + @patch('pgadmin.utils.pg_oauth2._exchange_oauth_access_token') + def case_exchange_mode_requires_client_id(self, exchange): + invalid_client_ids = ( + None, + '', + ' ', + 'cluster\0name', + ) + + for client_id in invalid_client_ids: + with self.subTest(client_id=client_id): + with self.assertRaisesRegex( + pg_oauth2.OAuthTokenError, + 'oauth_client_id'): + pg_oauth2.get_postgres_oauth_token( + 'exchange', + client_id, + 'openid', + ISSUER, + ) + + exchange.assert_not_called() + + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + @patch('pgadmin.utils.pg_oauth2._exchange_oauth_access_token') + def case_invalid_mode_raises_error(self, exchange, get_token): + with self.assertRaisesRegex( + pg_oauth2.OAuthTokenError, + 'Invalid pgAdmin OAuth token mode'): + pg_oauth2.get_postgres_oauth_token( + 'invalid-mode', + 'postgres-cluster', + 'openid', + ISSUER, + ) + + get_token.assert_not_called() + exchange.assert_not_called() + + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + @patch('pgadmin.utils.pg_oauth2._exchange_oauth_access_token') + def case_exchange_failure_does_not_fall_back_to_direct( + self, exchange, get_token): + failure = pg_oauth2.OAuthTokenExchangeError( + 'Token exchange failed.' + ) + exchange.side_effect = failure + + # A usable login token must not become a fallback. + get_token.return_value = 'pgadmin-access-token' + + with self.app.test_request_context('/'): + with self.assertRaises( + pg_oauth2.OAuthTokenExchangeError) as raised: + pg_oauth2.get_postgres_oauth_token( + 'exchange', + 'postgres-cluster', + 'openid', + ISSUER, + ) + + self.assertIs(raised.exception, failure) + exchange.assert_called_once_with( + 'postgres-cluster', + 'openid', + ISSUER, + ) + get_token.assert_not_called() + + +class TestInstallOAuthHook( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test installation of the process-global libpq OAuth hook.""" + + def setUp(self): + self.original_libpq = pg_oauth2._libpq + self.original_hook = pg_oauth2._hook + self.original_cleanup_hook = pg_oauth2._cleanup_hook + self.original_previous_hook = pg_oauth2._previous_hook + + pg_oauth2._libpq = None + pg_oauth2._hook = None + pg_oauth2._cleanup_hook = None + pg_oauth2._previous_hook = None + + def tearDown(self): + pg_oauth2._libpq = self.original_libpq + pg_oauth2._hook = self.original_hook + pg_oauth2._cleanup_hook = self.original_cleanup_hook + pg_oauth2._previous_hook = self.original_previous_hook + + @staticmethod + def _make_libpq(previous_hook_pointer=0): + libpq = MagicMock(spec_set=[ + 'PQlibVersion', + 'PQgetAuthDataHook', + 'PQsetAuthDataHook', + ]) + + libpq.PQlibVersion.return_value = 180000 + + installed_hook_pointer = ctypes.cast( + pg_oauth2._oauth_hook, + ctypes.c_void_p + ).value + + # The first call obtains the previous hook. Some versions of the + # implementation make a second call to verify installation. + libpq.PQgetAuthDataHook.side_effect = [ + previous_hook_pointer, + installed_hook_pointer, + ] + + return libpq + + def case_installs_oauth_hook(self): + libpq = self._make_libpq() + + with patch( + 'pgadmin.utils.pg_oauth2._get_libpq', + return_value=libpq) as get_libpq: + pg_oauth2.install_oauth_hook() + + get_libpq.assert_called_once_with() + libpq.PQsetAuthDataHook.assert_called_once_with( + pg_oauth2._oauth_hook + ) + + self.assertIs( + pg_oauth2._hook, + pg_oauth2._oauth_hook + ) + self.assertIs( + pg_oauth2._cleanup_hook, + pg_oauth2._oauth_cleanup + ) + self.assertIsNone(pg_oauth2._previous_hook) + + def case_installation_is_idempotent(self): + libpq = self._make_libpq() + + with patch( + 'pgadmin.utils.pg_oauth2._get_libpq', + return_value=libpq) as get_libpq: + pg_oauth2.install_oauth_hook() + pg_oauth2.install_oauth_hook() + + # The second call should return before loading libpq or replacing + # the process-global hook again. + get_libpq.assert_called_once_with() + libpq.PQsetAuthDataHook.assert_called_once_with( + pg_oauth2._oauth_hook + ) + + def case_preserves_previous_auth_data_hook(self): + @pg_oauth2._AUTH_HOOK + def previous_hook(authdata_type, conn, data): + return 67 + + previous_hook_pointer = ctypes.cast( + previous_hook, + ctypes.c_void_p + ).value + + libpq = self._make_libpq(previous_hook_pointer) + + with patch( + 'pgadmin.utils.pg_oauth2._get_libpq', + return_value=libpq): + pg_oauth2.install_oauth_hook() + + self.assertIsNotNone(pg_oauth2._previous_hook) + self.assertEqual( + pg_oauth2._previous_hook(0, None, None), + 67 + ) + + libpq.PQsetAuthDataHook.assert_called_once_with( + pg_oauth2._oauth_hook + ) + + # Keep the original ctypes callback alive until after the + # reconstructed callback pointer has been invoked. + self.assertIsNotNone(previous_hook) + + def case_failed_installation_can_be_retried(self): + libpq = self._make_libpq() + + # Override the fixture's sequence: each installation attempt + # starts with no previous hook. + libpq.PQgetAuthDataHook.side_effect = None + libpq.PQgetAuthDataHook.return_value = 0 + + libpq.PQsetAuthDataHook.side_effect = [ + RuntimeError('Hook installation failed.'), + None, + ] + + with patch( + 'pgadmin.utils.pg_oauth2._get_libpq', + return_value=libpq) as get_libpq: + with self.assertRaisesRegex( + RuntimeError, + 'Hook installation failed'): + pg_oauth2.install_oauth_hook() + + self.assertIsNone(pg_oauth2._hook) + + pg_oauth2.install_oauth_hook() + + self.assertEqual(get_libpq.call_count, 2) + self.assertEqual(libpq.PQsetAuthDataHook.call_count, 2) + + for hook_call in libpq.PQsetAuthDataHook.call_args_list: + self.assertIs( + hook_call.args[0], + pg_oauth2._oauth_hook, + ) + + self.assertIs(pg_oauth2._hook, pg_oauth2._oauth_hook) + self.assertIs( + pg_oauth2._cleanup_hook, + pg_oauth2._oauth_cleanup, + ) + + +class TestOAuthRequestContext( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test rejection of token retrieval outside a request.""" + + @patch('pgadmin.utils.pg_oauth2._get_pgadmin_oauth_token') + @patch('pgadmin.utils.pg_oauth2._exchange_oauth_access_token') + def case_token_modes_require_request_context( + self, exchange, get_token): + for mode in ('direct', 'exchange'): + with self.subTest(mode=mode): + with self.assertRaisesRegex( + pg_oauth2.OAuthRequestContextRequired, + 'authenticated pgAdmin request'): + pg_oauth2.get_postgres_oauth_token( + mode, + 'postgres-cluster', + 'openid', + ISSUER, + ) + + get_token.assert_not_called() + exchange.assert_not_called() + + +class TestLibpqSelection( + _OAuthUnitTestMixin, BaseTestGenerator): + """Test the restrictions used when selecting Psycopg's libpq.""" + + def setUp(self): + from psycopg import pq + + self.pq = pq + self.stack = ExitStack() + self.addCleanup(self.stack.close) + + self.libpq = MagicMock() + self.libpq.PQlibVersion.return_value = 180000 + self.library_path = '/test/libpq.so.5' + + self.stack.enter_context( + patch.object(pg_oauth2, '_libpq', None) + ) + self.stack.enter_context( + patch.object(self.pq, '__impl__', 'c') + ) + self.stack.enter_context( + patch.object(self.pq, 'version', return_value=180000) + ) + self.stack.enter_context( + patch.object(pg_oauth2.sys, 'platform', 'linux') + ) + + # Provide deterministic flags even when running these mocked + # Linux loader tests on Windows or macOS. + self.stack.enter_context( + patch.object( + pg_oauth2.os, 'RTLD_NOW', 2, create=True, + ) + ) + self.stack.enter_context( + patch.object( + pg_oauth2.os, 'RTLD_NOLOAD', 4, create=True, + ) + ) + + self.discover = self.stack.enter_context( + patch.object( + pg_oauth2, + '_get_loaded_libpq_paths', + return_value=[self.library_path], + ) + ) + self.load_library = self.stack.enter_context( + patch.object( + pg_oauth2.ctypes, + 'CDLL', + return_value=self.libpq, + ) + ) + + def case_python_implementation_reuses_existing_handle(self): + ctypes_module = MagicMock() + ctypes_module.pq = self.libpq + + with ExitStack() as stack: + stack.enter_context( + patch.object(self.pq, '__impl__', 'python') + ) + stack.enter_context( + patch.dict( + 'sys.modules', + {'psycopg.pq._pq_ctypes': ctypes_module}, + ) + ) + stack.enter_context( + patch.object( + self.pq, + '_pq_ctypes', + ctypes_module, + create=True, + ) + ) + + result = pg_oauth2._get_libpq() + + self.assertIs(result, self.libpq) + self.discover.assert_not_called() + self.load_library.assert_not_called() + + def case_c_implementation_uses_existing_loaded_library(self): + result = pg_oauth2._get_libpq() + + self.assertIs(result, self.libpq) + self.discover.assert_called_once_with() + self.load_library.assert_called_once_with( + self.library_path, + mode=pg_oauth2.os.RTLD_NOW | pg_oauth2.os.RTLD_NOLOAD, + ) + + def case_c_implementation_rejects_missing_or_ambiguous_library(self): + cases = ( + [], + ['/test/libpq.so.5', '/other/libpq.so.5'], + ) + + for candidates in cases: + with self.subTest(candidates=candidates): + self.discover.return_value = candidates + + with self.assertRaises(RuntimeError): + pg_oauth2._get_libpq() + + self.load_library.assert_not_called() + self.assertIsNone(pg_oauth2._libpq) + + def case_unsupported_implementations_fail_before_library_acquisition(self): + cases = ( + ('binary', 'linux'), + ('c', 'win32'), + ('c', 'darwin'), + ) + + for implementation, platform in cases: + with self.subTest( + implementation=implementation, + platform=platform): + with ExitStack() as stack: + stack.enter_context( + patch.object( + self.pq, '__impl__', implementation, + ) + ) + stack.enter_context( + patch.object( + pg_oauth2.sys, 'platform', platform, + ) + ) + + with self.assertRaises(RuntimeError): + pg_oauth2._get_libpq() + + self.discover.assert_not_called() + self.load_library.assert_not_called() + self.assertIsNone(pg_oauth2._libpq) + + def case_runtime_version_mismatch_is_rejected(self): + # Both versions are supported individually, but they disagree. + self.libpq.PQlibVersion.return_value = 180001 + + with self.assertRaisesRegex( + RuntimeError, + 'Psycopg reports version'): + pg_oauth2._get_libpq() + + self.assertIsNone(pg_oauth2._libpq) + self.libpq.PQsetAuthDataHook.assert_not_called() + + def case_only_libpq_18_versions_are_accepted(self): + for version in (180000, 180001, 189999): + with self.subTest(version=version): + pg_oauth2._check_oauth_libpq_version(version) + + for version in (170000, 179999, 190000, 190001): + with self.subTest(version=version): + with self.assertRaisesRegex(RuntimeError, '18.x'): + pg_oauth2._check_oauth_libpq_version(version)