mirror of
https://github.com/pgadmin-org/pgadmin4.git
synced 2026-08-18 17:04:53 -05:00
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.
This commit is contained in:
Ashesh Vashi
@@ -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):
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
##########################################################################
|
||||
#
|
||||
# pgAdmin 4 - PostgreSQL Tools
|
||||
#
|
||||
# Copyright (C) 2013 - 2026, The pgAdmin Development Team
|
||||
# This software is released under the PostgreSQL Licence
|
||||
#
|
||||
##########################################################################
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user