diff --git a/services/core/auth/src/nmp/core/auth/app/embedded_pdp/engine.py b/services/core/auth/src/nmp/core/auth/app/embedded_pdp/engine.py index b7c42f64f6..c26195d630 100644 --- a/services/core/auth/src/nmp/core/auth/app/embedded_pdp/engine.py +++ b/services/core/auth/src/nmp/core/auth/app/embedded_pdp/engine.py @@ -23,13 +23,35 @@ class PolicyEngineError(Exception): """Error during policy evaluation.""" +_engine: Optional[Engine] = None +_engine_lock = threading.Lock() + + +def _get_engine() -> Engine: + """Return the process-wide wasmtime Engine, creating it once (thread-safe, double-checked locking). + + Engine setup is the heavyweight step in wasmtime — JIT code generation and process-wide + trap/signal-handling registration — and wasmtime's own guidance is one Engine per process with + many cheap Stores created from it. Creating a fresh Engine per OPAPolicy instance (i.e. per + AuthService lifecycle, i.e. per test client) churns through hundreds of Engines over an xdist + worker's lifetime; that churn, not any single test, is what was crashing workers with no + traceback ("node down: Not properly terminated"). + """ + global _engine + if _engine is None: + with _engine_lock: + if _engine is None: + config = Config() + config.consume_fuel = True + _engine = Engine(config) + return _engine + + class OPAPolicy: """Wrapper for OPA WASM policy evaluation.""" def __init__(self, wasm_path: str, *, fuel_limit: int = 200_000_000, memory_limit_mb: int = 32): - config = Config() - config.consume_fuel = True - engine = Engine(config) + engine = _get_engine() self.fuel_limit = fuel_limit self.store = Store(engine) diff --git a/services/core/auth/tests/integration/test_scoped_access_keys.py b/services/core/auth/tests/integration/test_scoped_access_keys.py index de3c26ab1a..34964610ad 100644 --- a/services/core/auth/tests/integration/test_scoped_access_keys.py +++ b/services/core/auth/tests/integration/test_scoped_access_keys.py @@ -5,7 +5,6 @@ from pathlib import Path from unittest.mock import patch -import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from fastapi.testclient import TestClient @@ -16,11 +15,6 @@ from nmp.core.auth.config import AuthServiceConfig from nmp.testing.client import create_test_client -# Run on a dedicated worker so the inline create_test_client call doesn't share -# the module-level wasmtime singleton with the module-scoped test_client fixture -# used by sibling integration tests (which would violate wasmtime thread affinity). -pytestmark = pytest.mark.xdist_group("auth_scoped_access_keys") - ACCESS_KEYS_PATH = "/apis/auth/v2/access-keys" IAM_ROLE_BINDINGS_PATH = "/apis/auth/v2/iam/role-bindings" WORKSPACES_PATH = "/apis/entities/v2/workspaces" @@ -63,7 +57,7 @@ def _auth_configs(private_key_file: str) -> tuple[AuthConfig, AuthServiceConfig] ) service_config = AuthServiceConfig( **shared_config.model_dump(), - policy_data_refresh_interval=0.2, + policy_data_refresh_interval=0.05, bundle_cache_seconds=0, admin_email="admin@example.com", ) @@ -111,10 +105,6 @@ def test_scoped_access_key_created_by_auth_service_authenticates_platform_reques access_key = access_key_body["token"] access_key_jti = access_key_body["jti"] - list_keys = client.get(ACCESS_KEYS_PATH, headers=user_headers) - assert list_keys.status_code == 200, list_keys.text - assert [key["jti"] for key in list_keys.json()["data"]] == [access_key_jti] - role_binding = client.post( IAM_ROLE_BINDINGS_PATH, json={"principal": group, "role": "Viewer", "workspace": workspace}, @@ -132,9 +122,11 @@ async def validate_with_local_jwks(config: AuthConfig, token: str) -> TokenClaim f"{WORKSPACES_PATH}/{workspace}", headers={"Authorization": f"Bearer {access_key}"}, ) - assert response.status_code == 200, response.text - assert response.json()["name"] == workspace + assert response.status_code == 200, response.text + assert response.json()["name"] == workspace + + with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks): authenticate_response = client.get( "/apis/auth/authenticate", headers={"Authorization": f"Bearer {access_key}"}, @@ -143,33 +135,37 @@ async def validate_with_local_jwks(config: AuthConfig, token: str) -> TokenClaim "/apis/auth/authenticate", headers={"Authorization": f"Bearer {_tamper_jwt(access_key)}"}, ) - assert authenticate_response.status_code == 200, authenticate_response.text - assert authenticate_response.json()["principal"] == user - assert authenticate_response.json()["token_kind"] == "access_key" - assert invalid_authenticate_response.status_code == 401, invalid_authenticate_response.text + assert authenticate_response.status_code == 200, authenticate_response.text + assert authenticate_response.json()["principal"] == user + assert authenticate_response.json()["token_kind"] == "access_key" + assert invalid_authenticate_response.status_code == 401, invalid_authenticate_response.text + + with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks): invalid_workspace_response = client.get( f"{WORKSPACES_PATH}/{workspace}", headers={"Authorization": f"Bearer {_tamper_jwt(access_key)}"}, ) - assert invalid_workspace_response.status_code == 401, invalid_workspace_response.text - revoke_key = client.delete(f"{ACCESS_KEYS_PATH}/{access_key_jti}", headers=user_headers) - assert revoke_key.status_code == 200, revoke_key.text - revoked_keys = client.get(ACCESS_KEYS_PATH, headers=user_headers).json()["data"] - assert len(revoked_keys) == 1 - assert revoked_keys[0]["status"] == "REVOKED" + assert invalid_workspace_response.status_code == 401, invalid_workspace_response.text + + revoke_key = client.delete(f"{ACCESS_KEYS_PATH}/{access_key_jti}", headers=user_headers) + assert revoke_key.status_code == 200, revoke_key.text + revoked_keys = client.get(ACCESS_KEYS_PATH, headers=user_headers).json()["data"] + assert len(revoked_keys) == 1 + assert revoked_keys[0]["status"] == "REVOKED" + with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks): revoked_authenticate_response = client.get( "/apis/auth/authenticate", headers={"Authorization": f"Bearer {access_key}"}, ) - assert revoked_authenticate_response.status_code == 401, revoked_authenticate_response.text - revoked_workspace_response = client.get( f"{WORKSPACES_PATH}/{workspace}", headers={"Authorization": f"Bearer {access_key}"}, ) - assert revoked_workspace_response.status_code == 401, revoked_workspace_response.text + + assert revoked_authenticate_response.status_code == 401, revoked_authenticate_response.text + assert revoked_workspace_response.status_code == 401, revoked_workspace_response.text finally: client.delete(f"{WORKSPACES_PATH}/{workspace}", headers=SERVICE_HEADERS)