[webpubsub] fix jwt token + reverse proxy endpoint (Azure#22587)

Author: iscai-msft

sdk/webpubsub/azure-messaging-webpubsubservice/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- fix authentication who use a `reverse_proxy_endpoint` with either `WebPubSubServiceClient.from_connection_string` or pass in an `AzureKeyCredential` for authentication #22587
+
 ### Other Changes
 
 ## 1.0.0 (2021-11-10)

sdk/webpubsub/azure-messaging-webpubsubservice/azure/messaging/webpubsubservice/_patch.py

@@ -140,8 +140,14 @@ class JwtCredentialPolicy(SansIOHTTPPolicy):
 
     NAME_CLAIM_TYPE = "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/name"
 
-    def __init__(self, credential, user=None):
-        # type: (AzureKeyCredential, typing.Optional[str]) -> None
+    def __init__(
+        self,
+        credential: AzureKeyCredential,
+        *,
+        user: typing.Optional[str] = None,
+        origin_endpoint: typing.Optional[str] = None,
+        reverse_proxy_endpoint: typing.Optional[str] = None,
+    ) -> None:
         """Create a new instance of the policy associated with the given credential.
 
         :param credential: The azure.core.credentials.AzureKeyCredential instance to use
@@ -151,6 +157,8 @@ def __init__(self, credential, user=None):
         """
         self._credential = credential
         self._user = user
+        self._original_url = origin_endpoint
+        self._reverse_proxy_endpoint = reverse_proxy_endpoint
 
     def on_request(self, request):
         # type: (PipelineRequest) -> typing.Union[None, typing.Awaitable[None]]
@@ -159,8 +167,11 @@ def on_request(self, request):
         :param request: Request to be modified before sent from next policy.
         :type request: ~azure.core.pipeline.PipelineRequest
         """
+        url = request.http_request.url
+        if self._reverse_proxy_endpoint:
+            url = url.replace(self._reverse_proxy_endpoint, self._original_url, 1)
         request.http_request.headers["Authorization"] = "Bearer " + self._encode(
-            request.http_request.url
+            url
         )
         return super(JwtCredentialPolicy, self).on_request(request)
 
@@ -266,7 +277,12 @@ def _configure(
         self.authentication_policy = kwargs.get('authentication_policy')
         if self.credential and not self.authentication_policy:
             if isinstance(self.credential, AzureKeyCredential):
-                self.authentication_policy = JwtCredentialPolicy(self.credential, kwargs.get('user'))
+                self.authentication_policy = JwtCredentialPolicy(
+                    self.credential,
+                    user=kwargs.get("user"),
+                    origin_endpoint=kwargs.get("origin_endpoint"),
+                    reverse_proxy_endpoint=kwargs.get("reverse_proxy_endpoint")
+                )
             else:
                 self.authentication_policy = policies.BearerTokenCredentialPolicy(self.credential, *self.credential_scopes, **kwargs)
 

sdk/webpubsub/azure-messaging-webpubsubservice/azure/messaging/webpubsubservice/aio/_patch.py

@@ -112,7 +112,12 @@ def _configure(
         self.authentication_policy = kwargs.get('authentication_policy')
         if self.credential and not self.authentication_policy:
             if isinstance(self.credential, AzureKeyCredential):
-                self.authentication_policy = JwtCredentialPolicy(self.credential, kwargs.get('user'))
+                self.authentication_policy = JwtCredentialPolicy(
+                    self.credential,
+                    user=kwargs.get("user"),
+                    origin_endpoint=kwargs.get("origin_endpoint"),
+                    reverse_proxy_endpoint=kwargs.get("reverse_proxy_endpoint")
+                )
             else:
                 self.authentication_policy = policies.AsyncBearerTokenCredentialPolicy(self.credential, *self.credential_scopes, **kwargs)
 

sdk/webpubsub/azure-messaging-webpubsubservice/tests/recordings/test_reverse_proxy.test_reverse_proxy_call.yaml

@@ -0,0 +1,36 @@
+interactions:
+- request:
+    body: '{"Hello": "reverse_proxy_endpoint!"}'
+    headers:
+      Accept:
+      - application/json, text/json
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '36'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-messaging-webpubsubservice/1.0.1 Python/3.9.7 (macOS-12.1-x86_64-i386-64bit)
+    method: POST
+    uri: https://myservice.webpubsub.azure.com/api/hubs/hub/:send?api-version=2021-10-01
+  response:
+    body:
+      string: ''
+    headers:
+      api-supported-versions:
+      - '2021-10-01'
+      connection:
+      - keep-alive
+      content-length:
+      - '0'
+      date:
+      - Thu, 20 Jan 2022 21:53:17 GMT
+      strict-transport-security:
+      - max-age=15724800; includeSubDomains
+    status:
+      code: 202
+      message: Accepted
+version: 1

sdk/webpubsub/azure-messaging-webpubsubservice/tests/recordings/test_reverse_proxy_async.test_reverse_proxy_call.yaml

@@ -0,0 +1,28 @@
+interactions:
+- request:
+    body: '{"Hello": "reverse_proxy_endpoint!"}'
+    headers:
+      Accept:
+      - application/json, text/json
+      Content-Length:
+      - '36'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-messaging-webpubsubservice/1.0.1 Python/3.9.7 (macOS-12.1-x86_64-i386-64bit)
+    method: POST
+    uri: https://myservice.webpubsub.azure.com/api/hubs/hub/:send?api-version=2021-10-01
+  response:
+    body:
+      string: ''
+    headers:
+      api-supported-versions: '2021-10-01'
+      connection: keep-alive
+      content-length: '0'
+      date: Thu, 20 Jan 2022 21:55:24 GMT
+      strict-transport-security: max-age=15724800; includeSubDomains
+    status:
+      code: 202
+      message: Accepted
+    url: https://webpubsub-yyc.webpubsub.azure.com/api/hubs/hub/:send?api-version=2021-10-01
+version: 1

sdk/webpubsub/azure-messaging-webpubsubservice/tests/test_reverse_proxy.py

@@ -9,31 +9,45 @@
 from azure.messaging.webpubsubservice._operations._operations import build_send_to_all_request
 from azure.core.credentials import AzureKeyCredential
 from devtools_testutils.fake_credential import FakeTokenCredential
+from testcase import WebpubsubTest, WebpubsubPowerShellPreparer
+from azure.identity import DefaultAzureCredential
 
-def test_reverse_proxy_endpoint_redirection_azure_key_credential():
-    def _callback(pipeline_request):
-        assert pipeline_request.http_request.url.startswith("https://apim.contoso.com/")
-        raise ValueError("Success!")
-    wps_endpoint = "https://wps.contoso.com/"
-    apim_endpoint = "https://apim.contoso.com/"
-    credential = AzureKeyCredential("abcdabcdabcdabcdabcdabcdabcdabcd")
-    client = WebPubSubServiceClient(wps_endpoint, "Hub", credential, reverse_proxy_endpoint=apim_endpoint)
-    request = build_send_to_all_request('Hub', content='test_webpubsub_send_request', content_type='text/plain')
+class WebpubsubReverseProxyTest(WebpubsubTest):
 
-    with pytest.raises(ValueError) as ex:
-        client.send_request(request, raw_request_hook=_callback)
-    assert "Success!" in str(ex.value)
+    def test_reverse_proxy_endpoint_redirection_azure_key_credential(self):
+        def _callback(pipeline_request):
+            assert pipeline_request.http_request.url.startswith("https://apim.contoso.com/")
+            raise ValueError("Success!")
+        wps_endpoint = "https://wps.contoso.com/"
+        apim_endpoint = "https://apim.contoso.com/"
+        credential = AzureKeyCredential("abcdabcdabcdabcdabcdabcdabcdabcd")
+        client = WebPubSubServiceClient(wps_endpoint, "Hub", credential, reverse_proxy_endpoint=apim_endpoint)
+        request = build_send_to_all_request('Hub', content='test_webpubsub_send_request', content_type='text/plain')
 
-def test_reverse_proxy_endpoint_redirection_identity():
-    def _callback(pipeline_request):
-        assert pipeline_request.http_request.url.startswith("https://apim.contoso.com/")
-        raise ValueError("Success!")
-    wps_endpoint = "https://wps.contoso.com/"
-    apim_endpoint = "https://apim.contoso.com/"
-    credential = FakeTokenCredential()
-    client = WebPubSubServiceClient(wps_endpoint, "Hub", credential, reverse_proxy_endpoint=apim_endpoint)
-    request = build_send_to_all_request('Hub', content='test_webpubsub_send_request', content_type='text/plain')
+        with pytest.raises(ValueError) as ex:
+            client.send_request(request, raw_request_hook=_callback)
+        assert "Success!" in str(ex.value)
 
-    with pytest.raises(ValueError) as ex:
-        client.send_request(request, raw_request_hook=_callback)
-    assert "Success!" in str(ex.value)
+    def test_reverse_proxy_endpoint_redirection_identity(self):
+        def _callback(pipeline_request):
+            assert pipeline_request.http_request.url.startswith("https://apim.contoso.com/")
+            raise ValueError("Success!")
+        wps_endpoint = "https://wps.contoso.com/"
+        apim_endpoint = "https://apim.contoso.com/"
+        credential = FakeTokenCredential()
+        client = WebPubSubServiceClient(wps_endpoint, "Hub", credential, reverse_proxy_endpoint=apim_endpoint)
+        request = build_send_to_all_request('Hub', content='test_webpubsub_send_request', content_type='text/plain')
+
+        with pytest.raises(ValueError) as ex:
+            client.send_request(request, raw_request_hook=_callback)
+        assert "Success!" in str(ex.value)
+
+    @WebpubsubPowerShellPreparer()
+    def test_reverse_proxy_call(self, webpubsub_connection_string, webpubsub_reverse_proxy_endpoint):
+        client = self.create_client(
+            connection_string=webpubsub_connection_string,
+            hub='hub',
+            logging_enable=True,
+            reverse_proxy_endpoint=webpubsub_reverse_proxy_endpoint
+        )
+        client.send_to_all({'Hello': 'reverse_proxy_endpoint!'})

sdk/webpubsub/azure-messaging-webpubsubservice/tests/test_reverse_proxy_async.py

@@ -10,30 +10,47 @@
 from azure.core.credentials import AzureKeyCredential
 from devtools_testutils.fake_async_credential import AsyncFakeCredential
 
-@pytest.mark.asyncio
-async def test_reverse_proxy_endpoint_redirection():
-    def _callback(pipeline_request):
-        assert pipeline_request.http_request.url.startswith("https://apim.contoso.com/")
-        raise ValueError("Success!")
-    wps_endpoint = "https://wps.contoso.com/"
-    apim_endpoint = "https://apim.contoso.com/"
-    credential = AzureKeyCredential("abcdabcdabcdabcdabcdabcdabcdabcd")
-    request = build_send_to_all_request("Hub", content='test_webpubsub_send_request', content_type='text/plain')
-    async with WebPubSubServiceClient(wps_endpoint, "Hub", credential, reverse_proxy_endpoint=apim_endpoint) as client:
-        with pytest.raises(ValueError) as ex:
-            await client.send_request(request, raw_request_hook=_callback)
-        assert "Success!" in str(ex.value)
+from testcase import WebpubsubPowerShellPreparer
+from testcase_async import WebpubsubAsyncTest
+
+class WebpubsubReverseProxyTestAsync(WebpubsubAsyncTest):
+
+    @pytest.mark.asyncio
+    async def test_reverse_proxy_endpoint_redirection(self):
+        def _callback(pipeline_request):
+            assert pipeline_request.http_request.url.startswith("https://apim.contoso.com/")
+            raise ValueError("Success!")
+        wps_endpoint = "https://wps.contoso.com/"
+        apim_endpoint = "https://apim.contoso.com/"
+        credential = AzureKeyCredential("abcdabcdabcdabcdabcdabcdabcdabcd")
+        request = build_send_to_all_request("Hub", content='test_webpubsub_send_request', content_type='text/plain')
+        async with WebPubSubServiceClient(wps_endpoint, "Hub", credential, reverse_proxy_endpoint=apim_endpoint) as client:
+            with pytest.raises(ValueError) as ex:
+                await client.send_request(request, raw_request_hook=_callback)
+            assert "Success!" in str(ex.value)
+
+    @pytest.mark.asyncio
+    async def test_reverse_proxy_endpoint_redirection_identity(self):
+        def _callback(pipeline_request):
+            assert pipeline_request.http_request.url.startswith("https://apim.contoso.com/")
+            raise ValueError("Success!")
+        wps_endpoint = "https://wps.contoso.com/"
+        apim_endpoint = "https://apim.contoso.com/"
+        credential = AsyncFakeCredential()
+        request = build_send_to_all_request('Hub', content='test_webpubsub_send_request', content_type='text/plain')
+        async with WebPubSubServiceClient(wps_endpoint, "Hub", credential, reverse_proxy_endpoint=apim_endpoint) as client:
+            with pytest.raises(ValueError) as ex:
+                await client.send_request(request, raw_request_hook=_callback)
+            assert "Success!" in str(ex.value)
+
+    @pytest.mark.asyncio
+    @WebpubsubPowerShellPreparer()
+    async def test_reverse_proxy_call(self, webpubsub_connection_string, webpubsub_reverse_proxy_endpoint):
+        client = self.create_client(
+            connection_string=webpubsub_connection_string,
+            hub='hub',
+            logging_enable=True,
+            reverse_proxy_endpoint=webpubsub_reverse_proxy_endpoint
+        )
+        await client.send_to_all({'Hello': 'reverse_proxy_endpoint!'})
 
-@pytest.mark.asyncio
-async def test_reverse_proxy_endpoint_redirection_identity():
-    def _callback(pipeline_request):
-        assert pipeline_request.http_request.url.startswith("https://apim.contoso.com/")
-        raise ValueError("Success!")
-    wps_endpoint = "https://wps.contoso.com/"
-    apim_endpoint = "https://apim.contoso.com/"
-    credential = AsyncFakeCredential()
-    request = build_send_to_all_request('Hub', content='test_webpubsub_send_request', content_type='text/plain')
-    async with WebPubSubServiceClient(wps_endpoint, "Hub", credential, reverse_proxy_endpoint=apim_endpoint) as client:
-        with pytest.raises(ValueError) as ex:
-            await client.send_request(request, raw_request_hook=_callback)
-        assert "Success!" in str(ex.value)

sdk/webpubsub/azure-messaging-webpubsubservice/tests/test_smoke.py

@@ -5,7 +5,6 @@
 # license information.
 # -------------------------------------------------------------------------
 import pytest
-from requests.packages.urllib3.connection import connection
 from testcase import WebpubsubTest, WebpubsubPowerShellPreparer
 from azure.messaging.webpubsubservice._operations._operations import build_send_to_all_request
 from azure.core.exceptions import ServiceRequestError

sdk/webpubsub/azure-messaging-webpubsubservice/tests/testcase.py

@@ -15,7 +15,7 @@ def __init__(self, method_name, **kwargs):
 
     def create_client(self, endpoint=None, hub=None, reverse_proxy_endpoint=None, **kwargs):
         if kwargs.get("connection_string"):
-            return WebPubSubServiceClient.from_connection_string(kwargs.pop("connection_string"), hub)
+            return WebPubSubServiceClient.from_connection_string(kwargs.pop("connection_string"), hub, **kwargs)
         credential = self.get_credential(WebPubSubServiceClient)
         return self.create_client_from_credential(
             WebPubSubServiceClient,

sdk/webpubsub/azure-messaging-webpubsubservice/tests/testcase_async.py

@@ -14,7 +14,7 @@ def __init__(self, method_name, **kwargs):
 
     def create_client(self, endpoint=None, hub=None, reverse_proxy_endpoint=None, **kwargs):
         if kwargs.get("connection_string"):
-            return WebPubSubServiceClient.from_connection_string(kwargs.pop("connection_string"), hub)
+            return WebPubSubServiceClient.from_connection_string(kwargs.pop("connection_string"), hub, **kwargs)
         credential = self.get_credential(WebPubSubServiceClient, is_async=True)
         return self.create_client_from_credential(
             WebPubSubServiceClient,