Skip to content

Commit 8836ca8

Browse files
Fix managed identity authentication for Relay connection (#109)
- Correct OAuth scope from relay.azure.com to relay.azure.net - Pass Bearer token as sb-hc-token URL parameter (matching SAS behaviour), rather than via Authorization header
1 parent f5e86e4 commit 8836ca8

2 files changed

Lines changed: 14 additions & 28 deletions

File tree

src/relay_listener.py

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
logger = logging.getLogger(__name__)
3030

3131
DB_PATH = os.getenv("MWL_DB_PATH", "/var/lib/pacs/worklist.db")
32-
AZURE_RELAY_SCOPE = "https://relay.azure.com/.default"
32+
AZURE_RELAY_SCOPE = "https://relay.azure.net/.default"
3333
SAS_TOKEN_EXPIRY_SECONDS = 3600
3434

3535

@@ -101,7 +101,6 @@ def _connect(self):
101101
return connect(
102102
self.relay_uri.connection_url(),
103103
compression=None,
104-
additional_headers=self.relay_uri.auth_headers(),
105104
)
106105

107106

@@ -126,18 +125,16 @@ def connection_url(self) -> str:
126125
base = f"wss://{self.relay_namespace}/$hc/{self.hybrid_connection_name}?sb-hc-action=listen"
127126
if self._use_sas():
128127
token = self._create_sas_token()
129-
return f"{base}&sb-hc-token={urllib.parse.quote_plus(token)}"
130-
return base
128+
else:
129+
token = self._create_bearer_token()
130+
return f"{base}&sb-hc-token={urllib.parse.quote_plus(token)}"
131131

132-
def auth_headers(self) -> dict:
133-
if self._use_sas():
134-
return {}
132+
def _create_bearer_token(self) -> str:
135133
if self._credential is None:
136134
raise CredentialNotAvailableError(
137135
"No credential available — _credential should never be None when not using SAS"
138136
)
139-
token = self._credential.get_token(AZURE_RELAY_SCOPE).token
140-
return {"Authorization": f"Bearer {token}"}
137+
return f"Bearer {self._credential.get_token(AZURE_RELAY_SCOPE).token}"
141138

142139
def _create_sas_token(self, expiry_seconds: int = SAS_TOKEN_EXPIRY_SECONDS) -> str:
143140
uri = f"http://{self.relay_namespace}/{self.hybrid_connection_name}"

tests/test_relay_listener.py

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -119,14 +119,11 @@ def setup(self, monkeypatch):
119119
monkeypatch.delenv("AZURE_RELAY_SHARED_ACCESS_KEY", raising=False)
120120
yield
121121

122-
def test_connection_url(self):
122+
def test_connection_url(self, mock_azure_credential):
123123
subject = RelayURI()
124-
assert subject.connection_url() == "wss://test-namespace/$hc/test-connection?sb-hc-action=listen"
125-
126-
def test_auth_headers(self, mock_azure_credential):
127-
subject = RelayURI()
128-
assert subject.auth_headers() == {"Authorization": "Bearer test-token"}
129-
mock_azure_credential.get_token.assert_called_once_with("https://relay.azure.com/.default")
124+
url = subject.connection_url()
125+
assert url.startswith("wss://test-namespace/$hc/test-connection?sb-hc-action=listen")
126+
assert "sb-hc-token=Bearer+test-token" in url
130127

131128
def test_uses_default_azure_credential(self, mock_azure_credential):
132129
with patch("relay_listener.DefaultAzureCredential") as mock_dac:
@@ -152,10 +149,6 @@ def test_connection_url_includes_sas_token(self):
152149
assert url.startswith("wss://test-namespace/$hc/test-connection?sb-hc-action=listen")
153150
assert "sb-hc-token=SharedAccessSignature" in url
154151

155-
def test_auth_headers_are_empty(self):
156-
subject = RelayURI()
157-
assert subject.auth_headers() == {}
158-
159152
def test_no_credential_is_created(self):
160153
with patch("relay_listener.DefaultAzureCredential") as mock_dac:
161154
with patch("relay_listener.ManagedIdentityCredential") as mock_mic:
@@ -180,15 +173,11 @@ def test_uses_managed_identity_credential(self, mock_azure_credential):
180173
subject = RelayURI()
181174
assert subject._credential is mock_mic.return_value
182175

183-
def test_sas_key_is_ignored(self, monkeypatch):
176+
def test_sas_key_is_ignored(self, mock_azure_credential, monkeypatch):
184177
monkeypatch.setenv("AZURE_RELAY_SHARED_ACCESS_KEY", "some-key")
185178
subject = RelayURI()
186179
assert not subject._use_sas()
187-
assert "sb-hc-token" not in subject.connection_url()
188-
189-
def test_auth_headers_use_bearer_token(self, mock_azure_credential):
190-
subject = RelayURI()
191-
assert subject.auth_headers() == {"Authorization": "Bearer test-token"}
180+
assert "sb-hc-token=Bearer+test-token" in subject.connection_url()
192181

193182

194183
class TestVerifyCredentials:
@@ -201,12 +190,12 @@ def test_logs_sas_when_key_present(self, monkeypatch):
201190
def test_verifies_default_azure_credential_when_no_key(self, mock_azure_credential, monkeypatch):
202191
monkeypatch.delenv("AZURE_RELAY_SHARED_ACCESS_KEY", raising=False)
203192
verify_credentials()
204-
mock_azure_credential.get_token.assert_called_with("https://relay.azure.com/.default")
193+
mock_azure_credential.get_token.assert_called_with("https://relay.azure.net/.default")
205194

206195
def test_verifies_managed_identity_in_production(self, mock_azure_credential, monkeypatch):
207196
monkeypatch.setenv("ENVIRONMENT", "prod")
208197
verify_credentials()
209-
mock_azure_credential.get_token.assert_called_with("https://relay.azure.com/.default")
198+
mock_azure_credential.get_token.assert_called_with("https://relay.azure.net/.default")
210199

211200
def test_raises_client_authentication_error_on_credential_failure(self, monkeypatch):
212201
monkeypatch.delenv("AZURE_RELAY_SHARED_ACCESS_KEY", raising=False)

0 commit comments

Comments
 (0)