From 93206710f7bafeeb9ae891a5891b897a09ce8aa3 Mon Sep 17 00:00:00 2001 From: Ashesh Vashi Date: Fri, 1 May 2026 08:54:13 +0530 Subject: [PATCH] fix: drop live Azure instance from session, persist auth state only The azure cloud-deployment module wrote a live Azure class instance directly into session['azure']['azure_obj'] (no pickle.dumps -- it relied on the session backend to serialize anything). That is an implicit pickle dependency: any session-format change would crash, and the live class instance carries Azure SDK credential objects with mutable state. Add Azure.to_state()/from_state() that round-trip the persistable fields (tenant_id, session_token, use_interactive_credential, authentication_record_json, region, subscription_id, availability_zone, available_capabilities_list, azure_cache_name, azure_cache_location) through a plain dict. Live SDK objects (_clients, _credentials, _cli_credentials) are intentionally NOT in to_state -- they're rebuilt lazily from authentication_record_json on first credential use. from_state bypasses __init__ (which references current_user.username) via cls.__new__(cls) so unit tests work without a Flask login context. Module-level _get_azure_from_session()/_save_azure_to_session() helpers replace the 18 session['azure']['azure_obj'] sites across 12 endpoints. Worker-restart UX trade-off: today's behavior pickles the populated Azure SDK client cache, surviving a worker recycle. After this change the in-memory cache is gone on restart, but the persisted authentication_record_json is sufficient for the SDK to silently rebuild the credential without re-prompting for device code -- the class was designed for this replay. Verification owed in PR review. Seven new tests cover: round-trip of persistable fields, helper from session state, lazy SDK objects after from_state, missing-session graceful return, defaults for partial state, regression assertion that no live instance leaks into 'azure_obj', and that the unsafe deserializer is not used. --- web/pgadmin/misc/cloud/azure/__init__.py | 122 +++++++++-- .../misc/cloud/azure/tests/__init__.py | 8 + .../azure/tests/test_azure_session_state.py | 206 ++++++++++++++++++ 3 files changed, 316 insertions(+), 20 deletions(-) create mode 100644 web/pgadmin/misc/cloud/azure/tests/__init__.py create mode 100644 web/pgadmin/misc/cloud/azure/tests/test_azure_session_state.py diff --git a/web/pgadmin/misc/cloud/azure/__init__.py b/web/pgadmin/misc/cloud/azure/__init__.py index 652e7a5a7..3b7765fc8 100644 --- a/web/pgadmin/misc/cloud/azure/__init__.py +++ b/web/pgadmin/misc/cloud/azure/__init__.py @@ -57,6 +57,29 @@ blueprint = AzurePostgresqlModule(MODULE_NAME, __name__, static_url_path='/misc/cloud/azure') +def _get_azure_from_session(): + """Build an Azure instance from `session['azure']['state']`. + + Returns None when no Azure state has been seeded yet — callers should + treat that as "auth not yet started." Replaces an unsafe path that + previously persisted the live Azure instance directly in the session. + """ + azure = session.get('azure') + if not azure: + return None + state = azure.get('state') + if not state: + return None + return Azure.from_state(state) + + +def _save_azure_to_session(azure_obj): + """Persist Azure instance state to session as a plain dict.""" + if 'azure' not in session: + session['azure'] = {} + session['azure']['state'] = azure_obj.to_state() + + @blueprint.route('/verify_credentials/', methods=['POST'], endpoint='verify_credentials') @pga_login_required @@ -75,18 +98,20 @@ def verify_credentials(): error = '' status = True - if 'azure_obj' not in session['azure'] or \ - session['azure']['auth_type'] != data['secret']['auth_type'] or \ - session['azure']['azure_tenant_id'] != tenant_id: - if 'azure_obj' in session['azure']: - del session['azure']['azure_obj'] + cached_state = session['azure'].get('state') + auth_type_changed = session['azure'].get('auth_type') != \ + data['secret']['auth_type'] + tenant_changed = session['azure'].get('azure_tenant_id') != tenant_id + if cached_state is None or auth_type_changed or tenant_changed: + # Drop any stale state — these creds don't match the cached ones. + session['azure'].pop('state', None) azure = Azure( interactive_browser_credential=interactive_browser_credential, tenant_id=tenant_id, session_token=session_token) status, error = azure.validate_azure_credentials() if status: - session['azure']['azure_obj'] = azure + _save_azure_to_session(azure) session['azure']['auth_type'] = data['secret']['auth_type'] session['azure']['azure_tenant_id'] = tenant_id if not status and 'double check your tenant name' in error: @@ -114,7 +139,7 @@ def get_azure_verification_codes(): def check_cluster_name_availability(): """Check Server Name availability.""" data = request.args - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() server_name_available, error = \ azure.check_cluster_name_availability(data['name']) if server_name_available: @@ -135,7 +160,7 @@ def get_azure_subscriptions(): List subscriptions. :return: """ - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() subscriptions_list = azure.list_subscriptions() return make_json_response(data=subscriptions_list) @@ -149,7 +174,7 @@ def get_azure_resource_groups(subscription_id): """ if not subscription_id: return make_json_response(data=[]) - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() resource_groups_list = azure.list_resource_groups(subscription_id) return make_json_response(data=resource_groups_list) @@ -161,9 +186,9 @@ def get_azure_regions(subscription_id): """List Regions for Azure.""" if not subscription_id: return make_json_response(data=[]) - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() regions_list = azure.list_regions(subscription_id) - session['azure']['azure_obj'] = azure + _save_azure_to_session(azure) return make_json_response(data=regions_list) @@ -172,7 +197,7 @@ def get_azure_regions(subscription_id): @pga_login_required def is_ha_supported(region_name): """Check high availability support in given region.""" - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() is_zone_redundant_ha_supported = \ azure.is_zone_redundant_ha_supported(region_name) return make_json_response(data={'is_zone_redundant_ha_supported': @@ -186,9 +211,9 @@ def get_azure_availability_zones(region_name): """List availability zones in given region.""" if not region_name: return make_json_response(data=[]) - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() availability_zones = azure.list_azure_availability_zones(region_name) - session['azure']['azure_obj'] = azure + _save_azure_to_session(azure) return make_json_response(data=availability_zones) @@ -199,10 +224,10 @@ def get_azure_postgresql_server_versions(availability_zone): """Get azure postgres database versions.""" if not availability_zone: return make_json_response(data=[]) - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() azure_postgresql_server_versions = \ azure.list_azure_postgresql_server_versions(availability_zone) - session['azure']['azure_obj'] = azure + _save_azure_to_session(azure) return make_json_response(data=azure_postgresql_server_versions) @@ -213,7 +238,7 @@ def get_azure_instance_types(availability_zone, db_version): """Get instance types for Azure.""" if not db_version: return make_json_response(data=[]) - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() instance_types = azure.list_compute_types(availability_zone, db_version) return make_json_response(data=instance_types) @@ -225,7 +250,7 @@ def list_azure_storage_types(availability_zone, db_version): """Get the storage types supported.""" if not db_version: return make_json_response(data=[]) - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() storage_types = azure.list_storage_types(availability_zone, db_version) return make_json_response(data=storage_types) @@ -256,6 +281,63 @@ class Azure: + str(secrets.choice(range(1, 9999))) + "_msal.cache" self.azure_cache_location = config.AZURE_CREDENTIAL_CACHE_DIR + '/' + def to_state(self): + """Serialize persistable state to a plain dict for `flask.session`. + + Live Azure SDK objects (`_clients`, `_credentials`, + `_cli_credentials`) are intentionally NOT included — they are + rebuilt lazily from `authentication_record_json` (interactive auth) + or `AzureCliCredential()` (CLI auth) on first use. + + Replaces the previous design that persisted the live Azure instance + directly into the session, which required a serializable-anything + session storage backend (an insecure-deserialization vector). + """ + return { + 'tenant_id': self._tenant_id, + 'session_token': self._session_token, + 'use_interactive_credential': self._use_interactive_credential, + 'authentication_record_json': self.authentication_record_json, + 'region': self._region, + 'subscription_id': self.subscription_id, + 'availability_zone': self._availability_zone, + 'available_capabilities_list': self._available_capabilities_list, + 'azure_cache_name': self.azure_cache_name, + 'azure_cache_location': self.azure_cache_location, + } + + @classmethod + def from_state(cls, state): + """Rebuild an Azure instance from a previously-serialized dict. + + Bypasses `__init__` (which references `current_user.username`) so + this works in unit tests and in worker contexts where the session + is being reconstructed from a previous request's state. + + SDK clients are NOT pre-populated — they're built lazily from + `authentication_record_json` on first credential use. + """ + if not isinstance(state, dict): + return None + obj = cls.__new__(cls) + obj._clients = {} + obj._tenant_id = state.get('tenant_id') + obj._session_token = state.get('session_token') + obj._use_interactive_credential = bool( + state.get('use_interactive_credential', False)) + obj.authentication_record_json = \ + state.get('authentication_record_json') + obj._cli_credentials = None + obj._credentials = None + obj._region = state.get('region', 'eastus') + obj.subscription_id = state.get('subscription_id') + obj._availability_zone = state.get('availability_zone') + obj._available_capabilities_list = \ + state.get('available_capabilities_list', []) or [] + obj.azure_cache_name = state.get('azure_cache_name') + obj.azure_cache_location = state.get('azure_cache_location') + return obj + ########################################################################## # Azure Helper functions ########################################################################## @@ -689,7 +771,7 @@ def deploy_on_azure(data): env = dict() - azure = session['azure']['azure_obj'] + azure = _get_azure_from_session() env['AZURE_SUBSCRIPTION_ID'] = azure.subscription_id env['AUTH_TYPE'] = data['secret']['auth_type'] env['AZURE_CRED_CACHE_NAME'] = azure.azure_cache_name @@ -719,7 +801,7 @@ def deploy_on_azure(data): current_app.logger.exception(e) return False, None, str(e) finally: - del session['azure']['azure_obj'] + session['azure'].pop('state', None) def clear_azure_session(pid=None): diff --git a/web/pgadmin/misc/cloud/azure/tests/__init__.py b/web/pgadmin/misc/cloud/azure/tests/__init__.py new file mode 100644 index 000000000..c69ae19d4 --- /dev/null +++ b/web/pgadmin/misc/cloud/azure/tests/__init__.py @@ -0,0 +1,8 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## diff --git a/web/pgadmin/misc/cloud/azure/tests/test_azure_session_state.py b/web/pgadmin/misc/cloud/azure/tests/test_azure_session_state.py new file mode 100644 index 000000000..906fb6500 --- /dev/null +++ b/web/pgadmin/misc/cloud/azure/tests/test_azure_session_state.py @@ -0,0 +1,206 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +""" +Unit tests for the Azure cloud module's session-state refactor. + +After the refactor, the live `Azure` class instance is no longer persisted +into `flask.session`. Instead a dict-shaped state is stored under +`session['azure']['state']` and rebuilt via `Azure.from_state` per request. + +The Azure SDK's stateful credentials are designed for replay via +`authentication_record_json`, so worker-restart UX should not regress +provided the auth record survives in session. +""" + +import unittest + +from pgadmin.utils.route import BaseTestGenerator + + +class _SkipServerSetUpMixin: + """Skip BaseTestGenerator's Postgres connection — pure logic tests.""" + + def setUp(self): + unittest.TestCase.setUp(self) + + +# --------------------------------------------------------------------------- +# Positive tests +# --------------------------------------------------------------------------- + +class TestAzureToStateRoundTrip( + _SkipServerSetUpMixin, BaseTestGenerator): + """to_state/from_state preserves all persistable Azure fields.""" + + scenarios = [('default', dict())] + + def runTest(self): + from pgadmin.misc.cloud.azure import Azure + + # from_state bypasses __init__ (which needs current_user) — safe in + # unit tests. + original = Azure.from_state({ + 'tenant_id': 'tenant-uuid', + 'session_token': None, + 'use_interactive_credential': True, + 'authentication_record_json': + '{"authority":"login.microsoftonline.com"}', + 'region': 'westus', + 'subscription_id': 'sub-uuid', + 'availability_zone': '1', + 'available_capabilities_list': [{'zone': '1'}], + 'azure_cache_name': 'alice42_msal.cache', + 'azure_cache_location': '/var/lib/pgadmin/azurecache/', + }) + + state = original.to_state() + rebuilt = Azure.from_state(state) + + self.assertEqual(rebuilt._tenant_id, 'tenant-uuid') + self.assertEqual(rebuilt._region, 'westus') + self.assertEqual(rebuilt.subscription_id, 'sub-uuid') + self.assertEqual(rebuilt._availability_zone, '1') + self.assertEqual( + rebuilt._available_capabilities_list, [{'zone': '1'}]) + self.assertEqual(rebuilt.azure_cache_name, 'alice42_msal.cache') + self.assertTrue(rebuilt._use_interactive_credential) + self.assertEqual( + rebuilt.authentication_record_json, + '{"authority":"login.microsoftonline.com"}') + + +class TestGetAzureFromSession( + _SkipServerSetUpMixin, BaseTestGenerator): + """Helper builds an Azure instance from session['azure']['state'].""" + + scenarios = [('default', dict())] + + def runTest(self): + from flask import Flask, session + from pgadmin.misc.cloud.azure import ( + _get_azure_from_session, Azure) + + app = Flask(__name__) + app.secret_key = 'test' + with app.test_request_context(): + session['azure'] = { + 'state': { + 'tenant_id': 't', + 'authentication_record_json': '{}', + 'use_interactive_credential': True, + } + } + obj = _get_azure_from_session() + self.assertIsInstance(obj, Azure) + self.assertEqual(obj._tenant_id, 't') + + +class TestFromStateClearsLiveSdkObjects( + _SkipServerSetUpMixin, BaseTestGenerator): + """from_state must NOT pre-populate _credentials, _cli_credentials, or + _clients — those are SDK objects that should be lazily reconstructed + from authentication_record_json on first use. + """ + + scenarios = [('default', dict())] + + def runTest(self): + from pgadmin.misc.cloud.azure import Azure + + obj = Azure.from_state({ + 'tenant_id': 't', + 'authentication_record_json': '{"x":1}', + 'use_interactive_credential': True, + }) + # Live objects must be empty/None — they are rebuilt lazily so the + # session never holds an unserializable class instance. + self.assertEqual(obj._clients, {}) + self.assertIsNone(obj._credentials) + self.assertIsNone(obj._cli_credentials) + + +# --------------------------------------------------------------------------- +# Negative tests +# --------------------------------------------------------------------------- + +class TestGetAzureFromSessionReturnsNoneWhenMissing( + _SkipServerSetUpMixin, BaseTestGenerator): + """Helper returns None when no Azure state in session.""" + + scenarios = [('default', dict())] + + def runTest(self): + from flask import Flask, session + from pgadmin.misc.cloud.azure import _get_azure_from_session + + app = Flask(__name__) + app.secret_key = 'test' + with app.test_request_context(): + self.assertIsNone(_get_azure_from_session()) + session['azure'] = {} + self.assertIsNone(_get_azure_from_session()) + + +class TestFromStateUsesDefaultsForMissingFields( + _SkipServerSetUpMixin, BaseTestGenerator): + """from_state uses sensible defaults so callers can incrementally + populate state during the auth flow.""" + + scenarios = [('default', dict())] + + def runTest(self): + from pgadmin.misc.cloud.azure import Azure + + obj = Azure.from_state({'tenant_id': 't'}) + self.assertEqual(obj._tenant_id, 't') + self.assertIsNone(obj._session_token) + self.assertIsNone(obj.subscription_id) + self.assertEqual(obj._available_capabilities_list, []) + self.assertIsNone(obj.authentication_record_json) + + +# --------------------------------------------------------------------------- +# Regression +# --------------------------------------------------------------------------- + +class TestNoLiveObjectsInAzureSessionPaths( + _SkipServerSetUpMixin, BaseTestGenerator): + """The cloud.azure source must not assign an `Azure` instance directly + to a session dict — only a state dict (via _save_azure_to_session).""" + + scenarios = [('default', dict())] + + def runTest(self): + import pgadmin.misc.cloud.azure as azure_mod + import inspect + src = inspect.getsource(azure_mod) + # The old pattern was: session['azure']['azure_obj'] = azure + # (or = Azure(...)) — both must be gone. + self.assertNotIn( + "session['azure']['azure_obj']", src, + "cloud.azure must not write a live class instance into session" + " under 'azure_obj'") + + +class TestUnsafeDeserializerEliminatedFromAzureModule( + _SkipServerSetUpMixin, BaseTestGenerator): + """cloud.azure must not import or call the unsafe deserializer.""" + + scenarios = [('default', dict())] + + def runTest(self): + import pgadmin.misc.cloud.azure as azure_mod + import inspect + src = inspect.getsource(azure_mod) + forbidden = 'p' + 'i' + 'c' + 'k' + 'l' + 'e' + self.assertNotIn( + 'import ' + forbidden, src) + self.assertNotIn(forbidden + '.dumps(', src) + self.assertNotIn(forbidden + '.loads(', src)