[EventHub] add ssl_context kwarg to clients

sdk/eventhub/azure-eventhub/CHANGELOG.md

@@ -1,18 +1,15 @@
 # Release History
 
-## 5.12.3 (Unreleased)
-
-### Features Added
-
-### Breaking Changes
+## 5.12.3 (2024-10-09)
 
 ### Bugs Fixed
 
-### Other Changes
+- Fixed a bug where creating the SSL context in the async clients was making a blocking call outside of the constructor.([#37246](https://github.com/Azure/azure-sdk-for-python/issues/37246))
 
 ## 5.12.2 (2024-10-02)
 
 ### Bugs Fixed
+
 - Implemented backpressure for  async consumer  to address a memory leak issue. ([#36398](https://github.com/Azure/azure-sdk-for-python/issues/36398))
 
 ## 5.12.1 (2024-06-11)

sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py

@@ -2,12 +2,14 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
+import warnings
 from typing import Optional, Dict, Any, Union, TYPE_CHECKING
 from urllib.parse import urlparse
 
 from azure.core.pipeline.policies import RetryMode
 from ._constants import TransportType, DEFAULT_AMQPS_PORT, DEFAULT_AMQP_WSS_PORT
 if TYPE_CHECKING:
+    from ssl import SSLContext
     from ._transport._base import AmqpTransport
     from .aio._transport._base_async import AmqpTransportAsync
 
@@ -36,6 +38,7 @@ def __init__(
             custom_endpoint_address: Optional[str] = None,
             connection_verify: Optional[str] = None,
             use_tls: bool = True,
+            ssl_context: Optional["SSLContext"] = None,
             **kwargs: Any
         ):
         self.user_agent = user_agent
@@ -57,7 +60,13 @@ def __init__(
         self.receive_timeout = receive_timeout
         self.send_timeout = send_timeout
         self.custom_endpoint_address = custom_endpoint_address
-        self.connection_verify = connection_verify
+        self.connection_verify: Optional[Union[str, "SSLContext"]]
+        if ssl_context:
+            if connection_verify:
+                warnings.warn("ssl_context is specified, connection_verify will be ignored.")
+            self.connection_verify = ssl_context
+        else:
+            self.connection_verify = connection_verify
         self.custom_endpoint_hostname = None
         self.hostname = hostname
         self.use_tls = use_tls

sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py

@@ -15,6 +15,7 @@
 
 
 if TYPE_CHECKING:
+    from ssl import SSLContext
     from ._eventprocessor.partition_context import PartitionContext
     from ._common import EventData
     from ._client_base import CredentialTypes
@@ -124,6 +125,9 @@ class EventHubConsumerClient(
      authenticate the identity of the connection endpoint.
      Default is None in which case `certifi.where()` will be used.
     :paramtype connection_verify: str or None
+    :keyword ssl_context: The SSLContext object to use in the underlying Pure Python AMQP transport. If specified,
+     connection_verify will be ignored.
+    :paramtype ssl_context: ssl.SSLContext or None
     :keyword uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is
      False and the Pure Python AMQP library will be used as the underlying transport.
     :paramtype uamqp_transport: bool
@@ -244,6 +248,7 @@ def from_connection_string(
         load_balancing_strategy: Union[str, LoadBalancingStrategy] = LoadBalancingStrategy.GREEDY,
         custom_endpoint_address: Optional[str] = None,
         connection_verify: Optional[str] = None,
+        ssl_context: Optional["SSLContext"] = None,
         uamqp_transport: bool = False,
         **kwargs: Any
     ) -> "EventHubConsumerClient":
@@ -319,6 +324,9 @@ def from_connection_string(
          authenticate the identity of the connection endpoint.
          Default is None in which case `certifi.where()` will be used.
         :paramtype connection_verify: str or None
+        :keyword ssl_context: The SSLContext object to use in the underlying Pure Python AMQP transport. If specified,
+         connection_verify will be ignored.
+        :paramtype ssl_context: ssl.SSLContext or None
         :keyword uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is
          False and the Pure Python AMQP library will be used as the underlying transport.
         :paramtype uamqp_transport: bool
@@ -356,6 +364,7 @@ def from_connection_string(
             load_balancing_strategy=load_balancing_strategy,
             custom_endpoint_address=custom_endpoint_address,
             connection_verify=connection_verify,
+            ssl_context=ssl_context,
             uamqp_transport=uamqp_transport,
             **kwargs,
         )

sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py

@@ -30,6 +30,7 @@
 from .exceptions import ConnectError, EventHubError
 
 if TYPE_CHECKING:
+    from ssl import SSLContext
     from ._client_base import CredentialTypes
 
 SendEventTypes = List[Union[EventData, AmqpAnnotatedMessage]]
@@ -130,6 +131,9 @@ class EventHubProducerClient(
      authenticate the identity of the connection endpoint.
      Default is None in which case `certifi.where()` will be used.
     :paramtype connection_verify: Optional[str]
+    :keyword ssl_context: The SSLContext object to use in the underlying Pure Python AMQP transport. If specified,
+     connection_verify will be ignored.
+    :paramtype ssl_context: ssl.SSLContext or None
     :keyword uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is
      False and the Pure Python AMQP library will be used as the underlying transport.
     :paramtype uamqp_transport: bool
@@ -466,6 +470,7 @@ def from_connection_string(
         transport_type: Optional["TransportType"] = TransportType.Amqp,
         custom_endpoint_address: Optional[str] = None,
         connection_verify: Optional[str] = None,
+        ssl_context: Optional["SSLContext"] = None,
         uamqp_transport: bool = False,
         **kwargs: Any
     ) -> "EventHubProducerClient":
@@ -553,6 +558,9 @@ def from_connection_string(
          authenticate the identity of the connection endpoint.
          Default is None in which case `certifi.where()` will be used.
         :paramtype connection_verify: Optional[str]
+        :keyword ssl_context: The SSLContext object to use in the underlying Pure Python AMQP transport. If specified,
+         connection_verify will be ignored.
+        :paramtype ssl_context: ssl.SSLContext or None
         :keyword uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is
          False and the Pure Python AMQP library will be used as the underlying transport.
         :paramtype uamqp_transport: bool
@@ -589,6 +597,7 @@ def from_connection_string(
             transport_type=transport_type,
             custom_endpoint_address=custom_endpoint_address,
             connection_verify=connection_verify,
+            ssl_context=ssl_context,
             uamqp_transport=uamqp_transport,
             **kwargs
         )

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py

@@ -46,8 +46,6 @@
 import logging
 from threading import Lock
 
-import certifi
-
 from ._platform import KNOWN_TCP_OPTS, SOL_TCP
 from ._encode import encode_frame
 from ._decode import decode_frame, decode_empty_frame
@@ -171,7 +169,7 @@ def __init__(
         socket_timeout=SOCKET_TIMEOUT,
         socket_settings=None,
         raise_on_initial_eintr=True,
-        use_tls: bool =True,
+        use_tls: bool = True,
         **kwargs
     ):
         self._quick_recv = None
@@ -516,19 +514,15 @@ def _setup_transport(self):
             self.sock = self._wrap_socket(self.sock, **self.sslopts)
         self._quick_recv = self.sock.recv
 
-    def _wrap_socket(self, sock, context=None, **sslopts):
-        if context:
-            return self._wrap_context(sock, sslopts, **context)
-        return self._wrap_socket_sni(sock, **sslopts)
-
-    def _wrap_context(
-        self, sock, sslopts, check_hostname=None, **ctx_options
+    def _wrap_socket(
+        self,
+        sock,
+        **sslopts
     ):
-        ctx = ssl.create_default_context(**ctx_options)
-        ctx.verify_mode = ssl.CERT_REQUIRED
-        ctx.load_verify_locations(cafile=certifi.where())
-        ctx.check_hostname = check_hostname
-        return ctx.wrap_socket(sock, **sslopts)
+        if "context" in sslopts:
+            context = sslopts.pop("context")
+            return context.wrap_socket(sock, **sslopts)
+        return self._wrap_socket_sni(sock, **sslopts)
 
     def _wrap_socket_sni(
         self,

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py

@@ -12,7 +12,6 @@
 from functools import partial
 from typing import Any, Callable, Coroutine, List, Dict, Optional, Tuple, Union, overload, cast
 from typing_extensions import Literal
-import certifi
 
 from ..outcomes import Accepted, Modified, Received, Rejected, Released
 from ._connection_async import Connection
@@ -131,9 +130,9 @@ class AMQPClientAsync(AMQPClientSync):
      If port is not specified in the `custom_endpoint_address`, by default port 443 will be used.
     :paramtype custom_endpoint_address: str
     :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to
-     authenticate the identity of the connection endpoint.
+     authenticate the identity of the connection endpoint OR an instance of ssl.SSLContext to be used.
      Default is None in which case `certifi.where()` will be used.
-    :paramtype connection_verify: str
+    :paramtype connection_verify: str or ssl.SSLContext or None
     :keyword float socket_timeout: The maximum time in seconds that the underlying socket in the transport should
      wait when reading or writing data before timing out. The default value is 0.2 (for transport type Amqp),
      and 1 for transport type AmqpOverWebsocket.
@@ -246,7 +245,7 @@ async def open_async(self, connection=None):
             self._connection = Connection(
                 "amqps://" + self._hostname if self._use_tls else "amqp://" + self._hostname,
                 sasl_credential=self._auth.sasl,
-                ssl_opts={'ca_certs': self._connection_verify or certifi.where()},
+                ssl_opts=self._connection_verify,
                 container_id=self._name,
                 max_frame_size=self._max_frame_size,
                 channel_max=self._channel_max,
@@ -495,9 +494,9 @@ class SendClientAsync(SendClientSync, AMQPClientAsync):
      If port is not specified in the `custom_endpoint_address`, by default port 443 will be used.
     :paramtype custom_endpoint_address: str
     :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to
-     authenticate the identity of the connection endpoint.
+     authenticate the identity of the connection endpoint OR an instance of ssl.SSLContext to be used.
      Default is None in which case `certifi.where()` will be used.
-    :paramtype connection_verify: str
+    :paramtype connection_verify: str or ssl.SSLContext or None
     """
 
     async def _client_ready_async(self):
@@ -717,9 +716,9 @@ class ReceiveClientAsync(ReceiveClientSync, AMQPClientAsync):
      If port is not specified in the `custom_endpoint_address`, by default port 443 will be used.
     :paramtype custom_endpoint_address: str
     :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to
-     authenticate the identity of the connection endpoint.
+     authenticate the identity of the connection endpoint OR an instance of ssl.SSLContext to be used.
      Default is None in which case `certifi.where()` will be used.
-    :paramtype connection_verify: str
+    :paramtype connection_verify: str or ssl.SSLContext or None
     """
 
     async def _client_ready_async(self):

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py

@@ -41,10 +41,6 @@
 from io import BytesIO
 import logging
 
-
-
-import certifi
-
 from .._platform import KNOWN_TCP_OPTS, SOL_TCP
 from .._encode import encode_frame
 from .._decode import decode_frame, decode_empty_frame
@@ -182,7 +178,7 @@ def _build_ssl_opts(self, sslopts):
             return sslopts
         try:
             if "context" in sslopts:
-                return self._build_ssl_context(**sslopts.pop("context"))
+                return sslopts["context"]
             ssl_version = sslopts.get("ssl_version")
             if ssl_version is None:
                 ssl_version = ssl.PROTOCOL_TLS_CLIENT
@@ -224,15 +220,6 @@ def _build_ssl_opts(self, sslopts):
                 "SSL configuration must be a dictionary, or the value True."
             ) from None
 
-    def _build_ssl_context(
-        self, check_hostname=None, **ctx_options
-    ):
-        ctx = ssl.create_default_context(**ctx_options)
-        ctx.verify_mode = ssl.CERT_REQUIRED
-        ctx.load_verify_locations(cafile=certifi.where())
-        ctx.check_hostname = check_hostname
-        return ctx
-
 
 class AsyncTransport(
     AsyncTransportMixin

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py

@@ -145,9 +145,9 @@ class AMQPClient(
      If port is not specified in the `custom_endpoint_address`, by default port 443 will be used.
     :paramtype custom_endpoint_address: str
     :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to
-     authenticate the identity of the connection endpoint.
+     authenticate the identity of the connection endpoint OR an instance of ssl.SSLContext to be used.
      Default is None in which case `certifi.where()` will be used.
-    :paramtype connection_verify: str
+    :paramtype connection_verify: str or ssl.SSLContext or None
     :keyword float socket_timeout: The maximum time in seconds that the underlying socket in the transport should
      wait when reading or writing data before timing out. The default value is 0.2 (for transport type Amqp),
      and 1 for transport type AmqpOverWebsocket.
@@ -212,7 +212,12 @@ def __init__(self, hostname, **kwargs):
 
         # Custom Endpoint
         self._custom_endpoint_address = kwargs.get("custom_endpoint_address")
-        self._connection_verify = kwargs.get("connection_verify")
+        connection_verify = kwargs.get("connection_verify")
+        if connection_verify and not isinstance(connection_verify, str):    # ssl.SSLContext
+            self._connection_verify = {"context": connection_verify}
+        else: # str or None
+            self._connection_verify = {"ca_certs": connection_verify or certifi.where()}
+
 
         # Emulator
         self._use_tls: bool = kwargs.get("use_tls", True)
@@ -317,7 +322,7 @@ def open(self, connection=None):
             self._connection = Connection(
                 "amqps://" + self._hostname if self._use_tls else "amqp://" + self._hostname,
                 sasl_credential=self._auth.sasl,
-                ssl_opts={"ca_certs": self._connection_verify or certifi.where()},
+                ssl_opts=self._connection_verify,
                 container_id=self._name,
                 max_frame_size=self._max_frame_size,
                 channel_max=self._channel_max,
@@ -567,9 +572,9 @@ class SendClient(AMQPClient):
      If port is not specified in the `custom_endpoint_address`, by default port 443 will be used.
     :paramtype custom_endpoint_address: str
     :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to
-     authenticate the identity of the connection endpoint.
+     authenticate the identity of the connection endpoint OR an instance of ssl.SSLContext to be used.
      Default is None in which case `certifi.where()` will be used.
-    :paramtype connection_verify: str
+    :paramtype connection_verify: str or ssl.SSLContext or None
     """
 
     def __init__(self, hostname, target, **kwargs):
@@ -806,9 +811,9 @@ class ReceiveClient(AMQPClient): # pylint:disable=too-many-instance-attributes
      If port is not specified in the `custom_endpoint_address`, by default port 443 will be used.
     :paramtype custom_endpoint_address: str
     :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to
-     authenticate the identity of the connection endpoint.
+     authenticate the identity of the connection endpoint OR an instance of ssl.SSLContext to be used.
      Default is None in which case `certifi.where()` will be used.
-    :paramtype connection_verify: str
+    :paramtype connection_verify: str or ssl.SSLContext or None
     """
 
     def __init__(self, hostname, source, **kwargs):

sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py

@@ -579,7 +579,6 @@ def create_token_auth(
             timeout=config.auth_timeout,
             custom_endpoint_hostname=config.custom_endpoint_hostname,
             port=config.connection_port,
-            verify=config.connection_verify,
         )
         # if update_token:
         #    token_auth.update_token()  # TODO: why don't we need to update in pyamqp?

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py

@@ -28,6 +28,7 @@
 
 
 if TYPE_CHECKING:
+    from ssl import SSLContext
     from ._client_base_async import CredentialTypes
     from ._eventprocessor.partition_context import PartitionContext
     from ._eventprocessor.checkpoint_store import CheckpointStore
@@ -136,6 +137,9 @@ class EventHubConsumerClient(
      authenticate the identity of the connection endpoint.
      Default is None in which case `certifi.where()` will be used.
     :paramtype connection_verify: Optional[str]
+    :keyword ssl_context: The SSLContext object to use in the underlying Pure Python AMQP transport. If specified,
+     connection_verify will be ignored.
+    :paramtype ssl_context: ssl.SSLContext or None
     :keyword uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is
      False and the Pure Python AMQP library will be used as the underlying transport.
     :paramtype uamqp_transport: bool
@@ -258,6 +262,7 @@ def from_connection_string(
         load_balancing_strategy: Union[str, LoadBalancingStrategy] = LoadBalancingStrategy.GREEDY,
         custom_endpoint_address: Optional[str] = None,
         connection_verify: Optional[str] = None,
+        ssl_context: Optional["SSLContext"] = None,
         uamqp_transport: bool = False,
         **kwargs: Any
     ) -> "EventHubConsumerClient":
@@ -333,6 +338,9 @@ def from_connection_string(
          authenticate the identity of the connection endpoint.
          Default is None in which case `certifi.where()` will be used.
         :paramtype connection_verify: Optional[str]
+        :keyword ssl_context: The SSLContext object to use in the underlying Pure Python AMQP transport. If specified,
+         connection_verify will be ignored.
+        :paramtype ssl_context: ssl.SSLContext or None
         :keyword uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is
          False and the Pure Python AMQP library will be used as the underlying transport.
         :paramtype uamqp_transport: bool
@@ -369,6 +377,7 @@ def from_connection_string(
             load_balancing_strategy=load_balancing_strategy,
             custom_endpoint_address=custom_endpoint_address,
             connection_verify=connection_verify,
+            ssl_context=ssl_context,
             uamqp_transport=uamqp_transport,
             **kwargs,
         )