Authenticate Relay with Managed Identity token (#103)

In production, replace locally-stored SAS token in websockets URL with with Azure-generated Managed Identity JWT token in HTTP header.

In non-production environments, use a SAS token if set in the environment, otherwise default to DefaultAzureCredential.

Author: carlosmartinez

.env.development

@@ -1,6 +1,7 @@
 # Azure Relay Configuration
 AZURE_RELAY_NAMESPACE=manbrs-gateway-dev.servicebus.windows.net
 AZURE_RELAY_HYBRID_CONNECTION=name-of-your-choice-relay-test-hc
+# Optional: set these to use SAS token auth locally instead of managed identity / az login
 AZURE_RELAY_KEY_NAME=RootManageSharedAccessKey
 AZURE_RELAY_SHARED_ACCESS_KEY=YOUR_SHARED_ACCESS_KEY_HERE
 

docs/adr/ADR-005_Managed_Identity_For_Azure_Relay.md

@@ -0,0 +1,47 @@
+# ADR-005: Use Managed Identity for Azure Relay authentication
+
+Date: 2026-04-28
+
+Status: Accepted
+
+## Context
+
+The gateway connects to Azure Relay Hybrid Connections to receive worklist actions from Manage Breast Screening. A connection must be authenticated to Azure Relay.
+
+The initial implementation used Shared Access Signature (SAS) tokens. These are HMAC-SHA256 signatures computed from a shared secret key, embedded in the WebSocket connection URL as a query parameter (`sb-hc-token`). This required:
+
+- A shared access key to be provisioned and stored as an environment variable (`AZURE_RELAY_SHARED_ACCESS_KEY`)
+- The key name to be configured separately (`AZURE_RELAY_KEY_NAME`)
+- Manual key rotation when keys needed to change
+
+As the gateway runs inside the hospital network but is provisioned via Azure Arc, it can be assigned a managed identity through Arc-enabled infrastructure. Storing a long-lived shared secret in the environment is therefore unnecessary operational overhead and a potential security risk.
+
+That said, setting up a working Relay connection locally is already complex. Mandating managed identity for all environments would add further friction for developers, who would need Azure CLI credentials with a Relay Listener role assignment before they could run the service.
+
+## Decision
+
+In **production** (`ENVIRONMENT=prod`), the gateway uses `ManagedIdentityCredential` exclusively. The SAS token path is unavailable regardless of what environment variables are set. The gateway's managed identity must be assigned the **Azure Relay Listener** role on the hybrid connection resource in Azure.
+
+In **non-production** environments, the auth method is determined by whether `AZURE_RELAY_SHARED_ACCESS_KEY` is set:
+
+- If set, a SAS token is generated and embedded in the WebSocket URL (`sb-hc-token`), preserving the simpler local development setup.
+- If absent, `DefaultAzureCredential` is used, which works with Azure CLI credentials (`az login`) for developers who have the Listener role assigned to their identity.
+
+The token is passed as an `Authorization: Bearer` HTTP header on the WebSocket upgrade request for managed identity paths. Azure Relay validates it against Azure AD.
+
+A startup credential check (`verify_credentials()`) runs before the listen loop. In production this will raise `ClientAuthenticationError` immediately if the managed identity is not correctly configured. In non-production with a SAS key it logs the auth method and continues.
+
+`ManagedIdentityCredential` is preferred over `DefaultAzureCredential` in production because it is predictable — it only attempts the IMDS endpoint and fails clearly, rather than traversing a credential chain that could succeed unexpectedly via another mechanism.
+
+## Consequences
+
+### Positive Consequences
+
+- **No secrets in production:** No shared key to store, rotate, or accidentally leak in deployed environments
+- **Fail-fast on misconfiguration:** Startup validation raises `ClientAuthenticationError` immediately rather than failing silently in the reconnect loop
+- **Preserved local developer experience:** SAS tokens continue to work locally when `AZURE_RELAY_SHARED_ACCESS_KEY` is set
+- **Consistent with platform direction:** Aligns with how the gateway already authenticates to the DICOM API
+
+### Negative Consequences
+
+- **Azure infrastructure dependency in production:** The managed identity and its role assignment must exist before the service can start

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
   "requests>=2.33.1",
   "python-dotenv>=1.2.1",
   "azure-monitor-opentelemetry>=1.8.7",
+  "azure-identity>=1.23.0,<2.0.0",
 ]
 
 [dependency-groups]

src/relay_listener.py

@@ -14,11 +14,12 @@
 import time
 import urllib.parse
 
+from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
 from dotenv import load_dotenv
 from websockets.asyncio.client import connect
 from websockets.exceptions import ConnectionClosedError
-from websockets.frames import CloseCode
 
+from environment import Environment
 from services.mwl.create_worklist_item import CreateWorklistItem
 from services.storage import MWLStorage
 from telemetry import configure_telemetry
@@ -28,10 +29,14 @@
 logger = logging.getLogger(__name__)
 
 DB_PATH = os.getenv("MWL_DB_PATH", "/var/lib/pacs/worklist.db")
-EXPIRED_TOKEN = "ExpiredToken"
+AZURE_RELAY_SCOPE = "https://relay.azure.com/.default"
 SAS_TOKEN_EXPIRY_SECONDS = 3600
 
 
+class CredentialNotAvailableError(RuntimeError):
+    pass
+
+
 class RelayListener:
     """
     Socket Listener for Azure Relay.
@@ -40,9 +45,11 @@ class RelayListener:
     Environment variables:
     AZURE_RELAY_NAMESPACE: Azure Relay namespace (default: relay-test.servicebus.windows.net)
     AZURE_RELAY_HYBRID_CONNECTION: Azure Relay hybrid connection name (default: relay-test-hc)
-    AZURE_RELAY_KEY_NAME: Azure Relay shared access key name (default: RootManageSharedAccessKey)
-    AZURE_RELAY_SHARED_ACCESS_KEY: Azure Relay shared access key (default: none)
     MWL_DB_PATH: Path to the MWL SQLite database file (default: /var/lib/pacs/worklist.db)
+
+    Non-production only (SAS token fallback):
+    AZURE_RELAY_KEY_NAME: Shared access policy name (default: RootManageSharedAccessKey)
+    AZURE_RELAY_SHARED_ACCESS_KEY: Shared access key value
     """
 
     def __init__(self, storage: MWLStorage):
@@ -91,7 +98,11 @@ def process_action(self, payload: dict):
 
     def _connect(self):
         """Connect to Azure Relay."""
-        return connect(self.relay_uri.connection_url(), compression=None)
+        return connect(
+            self.relay_uri.connection_url(),
+            compression=None,
+            additional_headers=self.relay_uri.auth_headers(),
+        )
 
 
 class RelayURI:
@@ -100,9 +111,35 @@ def __init__(self):
         self.hybrid_connection_name = os.getenv("AZURE_RELAY_HYBRID_CONNECTION", "relay-test-hc")
         self.key_name = os.getenv("AZURE_RELAY_KEY_NAME", "RootManageSharedAccessKey")
         self.shared_access_key = os.getenv("AZURE_RELAY_SHARED_ACCESS_KEY", "")
+        self._env = Environment()
+        self._credential = None if self._use_sas() else self._build_credential()
+
+    def _use_sas(self) -> bool:
+        return not self._env.production and bool(self.shared_access_key)
 
-    def create_sas_token(self, expiry_seconds: int = SAS_TOKEN_EXPIRY_SECONDS) -> str:
-        """Create SAS token for Azure Relay authentication."""
+    def _build_credential(self):
+        if self._env.production:
+            return ManagedIdentityCredential()
+        return DefaultAzureCredential()
+
+    def connection_url(self) -> str:
+        base = f"wss://{self.relay_namespace}/$hc/{self.hybrid_connection_name}?sb-hc-action=listen"
+        if self._use_sas():
+            token = self._create_sas_token()
+            return f"{base}&sb-hc-token={urllib.parse.quote_plus(token)}"
+        return base
+
+    def auth_headers(self) -> dict:
+        if self._use_sas():
+            return {}
+        if self._credential is None:
+            raise CredentialNotAvailableError(
+                "No credential available — _credential should never be None when not using SAS"
+            )
+        token = self._credential.get_token(AZURE_RELAY_SCOPE).token
+        return {"Authorization": f"Bearer {token}"}
+
+    def _create_sas_token(self, expiry_seconds: int = SAS_TOKEN_EXPIRY_SECONDS) -> str:
         uri = f"http://{self.relay_namespace}/{self.hybrid_connection_name}"
         encoded_uri = urllib.parse.quote_plus(uri)
         expiry = str(int(time.time() + expiry_seconds))
@@ -115,12 +152,25 @@ def create_sas_token(self, expiry_seconds: int = SAS_TOKEN_EXPIRY_SECONDS) -> st
             f"&se={expiry}&skn={self.key_name}"
         )
 
-    def connection_url(self) -> str:
-        token = self.create_sas_token()
-        return (
-            f"wss://{self.relay_namespace}/$hc/{self.hybrid_connection_name}"
-            f"?sb-hc-action=listen&sb-hc-token={urllib.parse.quote_plus(token)}"
-        )
+
+def verify_credentials():
+    """
+    Verify relay credentials are available at startup.
+
+    In production, raises ClientAuthenticationError if managed identity is not configured.
+    In non-production with a SAS key present, logs the auth method and returns immediately.
+    """
+    uri = RelayURI()
+    if uri._use_sas():
+        logger.info("Using SAS token authentication for Azure Relay.")
+    else:
+        if uri._credential is None:
+            raise CredentialNotAvailableError(
+                "No credential available — _credential should never be None when not using SAS"
+            )
+        uri._credential.get_token(AZURE_RELAY_SCOPE)
+        credential_type = "ManagedIdentityCredential" if uri._env.production else "DefaultAzureCredential"
+        logger.info(f"Azure Relay credentials verified ({credential_type}).")
 
 
 async def main():
@@ -131,6 +181,7 @@ async def main():
     configure_telemetry(service_name="relay-listener")
 
     logger.info("Socket Listener Starting...")
+    verify_credentials()
     storage = MWLStorage(db_path=DB_PATH)
 
     while True:
@@ -142,13 +193,9 @@ async def main():
         except ConnectionClosedError as e:
             code = e.rcvd.code if e.rcvd else "N/A"
             reason = e.rcvd.reason if e.rcvd else "N/A"
-
-            if code == CloseCode.INTERNAL_ERROR.value and EXPIRED_TOKEN in reason:
-                logger.info("SAS token expired, refreshing...")
-            else:
-                logger.warning(f"Connection closed with code {code}: {reason}")
-                logger.warning("Retrying in 5 seconds...")
-                await asyncio.sleep(5)
+            logger.warning(f"Connection closed with code {code}: {reason}")
+            logger.warning("Retrying in 5 seconds...")
+            await asyncio.sleep(5)
         except Exception as e:
             logger.warning(f"Connection error: {e}")
             logger.warning("Retrying in 5 seconds...")

tests/conftest.py

@@ -2,7 +2,7 @@
 import sys
 from contextlib import contextmanager
 from pathlib import Path
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import numpy as np
 import pytest
@@ -15,6 +15,17 @@
 sys.path.append(f"{Path(__file__).parent.parent}/src")
 
 
+@pytest.fixture(autouse=True)
+def mock_azure_credential():
+    mock = MagicMock()
+    mock.get_token.return_value.token = "test-token"
+    with (
+        patch("relay_listener.DefaultAzureCredential", return_value=mock),
+        patch("relay_listener.ManagedIdentityCredential", return_value=mock),
+    ):
+        yield mock
+
+
 @pytest.fixture
 def tmp_dir():
     return f"{Path(__file__).parent}/tmp"