[EH Emulator] WIP support for emulator (Azure#34103)

* add use_tls bool

* adding on

* async consumer

* add for localhsot

* add to producer

* prod client

* getting closer

* reocgnize emulator slug if true

* udate

* sync websocket

* remove print

* async

* update

* d't use custom endpont

* pass through

* updates

* comment

Co-authored-by: Kashif Khan <[email protected]>

* updates

* comment

* pass thrpugh

* update

* error

* async

* changes

* fix

* sync send/recieve

* sync

* nits

* nit

* add async

* remove

* test

* test

* update typing

* pylint

* update if

* nit

* add ipv6 test

* set true at top level

* remove true

* use config

* my changes

* move everything to pyamqp

* remove

* nit

* space

* remove print

* remove print + typing

* merge error

* nits

* Update sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py

Co-authored-by: Kashif Khan <[email protected]>

* websocket port

* apply to sync

* websocket port

* Update sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py

Co-authored-by: Kashif Khan <[email protected]>

---------

Co-authored-by: Kashif Khan <[email protected]>
Co-authored-by: swathipil <[email protected]>

Author: 3 people

sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py

@@ -9,6 +9,7 @@
 import time
 import functools
 import collections
+import re
 from typing import Any, Dict, Tuple, List, Optional, TYPE_CHECKING, cast, Union
 from datetime import timedelta
 from urllib.parse import urlparse
@@ -68,20 +69,23 @@
 _LOGGER = logging.getLogger(__name__)
 _Address = collections.namedtuple("_Address", "hostname path")
 
+def _is_local_endpoint(endpoint: str) -> bool:
+    return re.match("^(127\\.[\\d.]+|[0:]+1|localhost)", endpoint.lower()) is not None
 
 def _parse_conn_str(
         conn_str: str,  # pylint:disable=unused-argument
         *,
         eventhub_name: Optional[str] = None,
         check_case: bool = False,
         **kwargs: Any
-    ) -> Tuple[str, Optional[str], Optional[str], str, Optional[str], Optional[int]]:
+    ) -> Tuple[str, Optional[str], Optional[str], str, Optional[str], Optional[int], bool]:
     endpoint = None
     shared_access_key_name = None
     shared_access_key = None
     entity_path: Optional[str] = None
     shared_access_signature: Optional[str] = None
     shared_access_signature_expiry = None
+    use_emulator: Optional[str] = None
     conn_settings = core_parse_connection_string(
         conn_str, case_sensitive_keys=check_case
     )
@@ -95,6 +99,7 @@ def _parse_conn_str(
             # only sas check is non case sensitive for both conn str properties and internal use
             if key.lower() == "sharedaccesssignature":
                 shared_access_signature = value
+        use_emulator = conn_settings.get("UseDevelopmentEmulator")
 
     if not check_case:
         endpoint = conn_settings.get("endpoint") or conn_settings.get("hostname")
@@ -104,6 +109,7 @@ def _parse_conn_str(
         shared_access_key = conn_settings.get("sharedaccesskey")
         entity_path = conn_settings.get("entitypath")
         shared_access_signature = conn_settings.get("sharedaccesssignature")
+        use_emulator = conn_settings.get("usedevelopmentemulator")
 
     if shared_access_signature:
         try:
@@ -124,11 +130,21 @@ def _parse_conn_str(
     # check that endpoint is valid
     if not endpoint:
         raise ValueError("Connection string is either blank or malformed.")
+
     parsed = urlparse(endpoint)
     if not parsed.netloc:
         raise ValueError("Invalid Endpoint on the Connection String.")
     host = cast(str, parsed.netloc.strip())
 
+    emulator = use_emulator=="true"
+    if emulator and not _is_local_endpoint(host):
+        raise ValueError(
+            "Invalid endpoint on the connection string. "
+            "For development connection strings, should be in the format: "
+            "Endpoint=sb://localhost;SharedAccessKeyName=<KeyName>;SharedAccessKey=<KeyValue>;"
+            "UseDevelopmentEmulator=true;"
+        )
+
     if any([shared_access_key, shared_access_key_name]) and not all(
         [shared_access_key, shared_access_key_name]
     ):
@@ -154,6 +170,7 @@ def _parse_conn_str(
         entity,
         str(shared_access_signature) if shared_access_signature else None,
         shared_access_signature_expiry,
+        emulator
     )
 
 
@@ -327,11 +344,15 @@ def __init__(
 
     @staticmethod
     def _from_connection_string(conn_str: str, **kwargs: Any) -> Dict[str, Any]:
-        host, policy, key, entity, token, token_expiry = _parse_conn_str(
+        host, policy, key, entity, token, token_expiry, emulator = _parse_conn_str(
             conn_str, **kwargs
         )
+
         kwargs["fully_qualified_namespace"] = host
         kwargs["eventhub_name"] = entity
+        # Check if emulator is in use, unset tls if it is
+        if emulator:
+            kwargs["use_tls"] = False
         if token and token_expiry:
             kwargs["credential"] = EventHubSASTokenCredential(token, token_expiry)
         elif policy and key:

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

@@ -35,6 +35,7 @@ def __init__(
             send_timeout: int = 60,
             custom_endpoint_address: Optional[str] = None,
             connection_verify: Optional[str] = None,
+            use_tls: bool = True,
             **kwargs: Any
         ):
         self.user_agent = user_agent
@@ -59,6 +60,7 @@ def __init__(
         self.connection_verify = connection_verify
         self.custom_endpoint_hostname = None
         self.hostname = hostname
+        self.use_tls = use_tls
 
         if self.http_proxy or self.transport_type.value == TransportType.AmqpOverWebsocket.value:
             self.transport_type = TransportType.AmqpOverWebsocket

sdk/eventhub/azure-eventhub/azure/eventhub/_connection_string_parser.py

@@ -91,7 +91,7 @@ def parse_connection_string(conn_str: str) -> "EventHubConnectionStringPropertie
     :return: A properties bag containing the parsed connection string.
     :rtype: ~azure.eventhub.EventHubConnectionStringProperties
     """
-    fully_qualified_namespace, policy, key, entity, signature = _parse_conn_str(
+    fully_qualified_namespace, policy, key, entity, signature, emulator = _parse_conn_str(
         conn_str, check_case=True
     )[:-1]
     endpoint = "sb://" + fully_qualified_namespace + "/"
@@ -102,5 +102,6 @@ def parse_connection_string(conn_str: str) -> "EventHubConnectionStringPropertie
         "shared_access_signature": signature,
         "shared_access_key_name": policy,
         "shared_access_key": key,
+        "emulator": emulator,
     }
     return EventHubConnectionStringProperties(**props)

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

@@ -126,6 +126,10 @@ def __init__( # pylint:disable=too-many-locals,too-many-statements
             self._port = PORT
         self.state: Optional[ConnectionState] = None
 
+        # Set the port for AmqpOverWebsocket
+        if transport_type.value == TransportType.AmqpOverWebsocket.value:
+            self._port = WEBSOCKET_PORT
+
         # Custom Endpoint
         custom_endpoint_address = kwargs.get("custom_endpoint_address")
         custom_endpoint = None
@@ -157,6 +161,7 @@ def __init__( # pylint:disable=too-many-locals,too-many-statements
             self._transport = sasl_transport(
                 host=endpoint,
                 credential=kwargs["sasl_credential"],
+                port=self._port,
                 custom_endpoint=custom_endpoint,
                 socket_timeout=self._socket_timeout,
                 network_trace_params=self._network_trace_params,

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

@@ -171,6 +171,7 @@ def __init__(
         socket_timeout=SOCKET_TIMEOUT,
         socket_settings=None,
         raise_on_initial_eintr=True,
+        use_tls: bool =True,
         **kwargs
     ):
         self._quick_recv = None
@@ -186,6 +187,8 @@ def __init__(
         self.socket_settings = socket_settings
         self.socket_lock = Lock()
 
+        self._use_tls = use_tls
+
     def connect(self):
         try:
             # are we already connected?
@@ -509,7 +512,8 @@ def __init__(
 
     def _setup_transport(self):
         """Wrap the socket in an SSL object."""
-        self.sock = self._wrap_socket(self.sock, **self.sslopts)
+        if self._use_tls:
+            self.sock = self._wrap_socket(self.sock, **self.sslopts)
         self._quick_recv = self.sock.recv
 
     def _wrap_socket(self, sock, context=None, **sslopts):
@@ -599,7 +603,8 @@ def _shutdown_transport(self):
         """Unwrap a SSL socket, so we can call shutdown()."""
         if self.sock is not None:
             try:
-                self.sock = self.sock.unwrap()
+                if self._use_tls:
+                    self.sock = self.sock.unwrap()
             except OSError:
                 pass
 
@@ -737,11 +742,12 @@ def connect(self):
             ) from None
         try:
             self.sock = create_connection(
-                url="wss://{}".format(self._custom_endpoint or self._host),
+                url="wss://{}".format(self._custom_endpoint or self._host) if self._use_tls
+                    else "ws://{}".format(self._custom_endpoint or self._host),
                 subprotocols=[AMQP_WS_SUBPROTOCOL],
                 timeout=self.socket_timeout,    # timeout for read/write operations
                 skip_utf8_validation=True,
-                sslopt=self.sslopts,
+                sslopt=self.sslopts if self._use_tls else None,
                 http_proxy_host=http_proxy_host,
                 http_proxy_port=http_proxy_port,
                 http_proxy_auth=http_proxy_auth,

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

@@ -244,7 +244,7 @@ async def open_async(self, connection=None):
             self._external_connection = True
         if not self._connection:
             self._connection = Connection(
-                "amqps://" + self._hostname,
+                "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()},
                 container_id=self._name,
@@ -257,6 +257,7 @@ async def open_async(self, connection=None):
                 http_proxy=self._http_proxy,
                 custom_endpoint_address=self._custom_endpoint_address,
                 socket_timeout=self._socket_timeout,
+                use_tls=self._use_tls,
             )
             await self._connection.open()
         if not self._session:

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

@@ -107,6 +107,10 @@ def __init__(# pylint:disable=too-many-locals,too-many-statements
             self._port = PORT
         self.state: Optional[ConnectionState] = None
 
+        # Set the port for AmqpOverWebsocket
+        if transport_type.value == TransportType.AmqpOverWebsocket.value:
+            self._port = WEBSOCKET_PORT
+
         # Custom Endpoint
         custom_endpoint_address = kwargs.get("custom_endpoint_address")
         custom_endpoint = None
@@ -141,6 +145,7 @@ def __init__(# pylint:disable=too-many-locals,too-many-statements
             self._transport: Union[SASLTransport, SASLWithWebSocket, AsyncTransport] = sasl_transport(
                 host=endpoint,
                 credential=kwargs["sasl_credential"],
+                port=self._port,
                 custom_endpoint=custom_endpoint,
                 socket_timeout=self._socket_timeout,
                 network_trace_params=self._network_trace_params,

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

@@ -247,6 +247,7 @@ def __init__(
         ssl_opts=False,
         socket_settings=None,
         raise_on_initial_eintr=True,
+        use_tls: bool = True,
         **kwargs,  # pylint: disable=unused-argument
     ):
         self.connected = False
@@ -260,6 +261,7 @@ def __init__(
         self.socket_lock = asyncio.Lock()
         self.sslopts = ssl_opts
         self.network_trace_params = kwargs.get('network_trace_params')
+        self._use_tls = use_tls
 
     async def connect(self):
         try:
@@ -280,10 +282,10 @@ async def connect(self):
             self.reader, self.writer = await asyncio.open_connection(
                 host=self.host,
                 port=self.port,
-                ssl=self.sslopts,
+                ssl=self.sslopts if self._use_tls else None,
                 family=socket.AF_UNSPEC,
                 proto=SOL_TCP,
-                server_hostname=self.host if self.sslopts else None,
+                server_hostname=self.host if self._use_tls else None,
             )
             self.connected = True
             sock = self.writer.transport.get_extra_info("socket")
@@ -436,6 +438,7 @@ def __init__(
         port=WEBSOCKET_PORT,
         socket_timeout=CONNECT_TIMEOUT,
         ssl_opts=None,
+        use_tls: bool =True,
         **kwargs
     ):
         self._read_buffer = BytesIO()
@@ -449,6 +452,7 @@ def __init__(
         self._http_proxy = kwargs.get("http_proxy", None)
         self.connected = False
         self.network_trace_params = kwargs.get('network_trace_params')
+        self._use_tls = use_tls
 
     async def connect(self):
         self.sslopts = self._build_ssl_opts(self.sslopts)
@@ -478,9 +482,9 @@ async def connect(self):
 
         self.session = ClientSession()
         if self._custom_endpoint:
-            url = f"wss://{self._custom_endpoint}"
+            url = f"wss://{self._custom_endpoint}" if self._use_tls else f"ws://{self._custom_endpoint}"
         else:
-            url = f"wss://{self.host}"
+            url = f"wss://{self.host}" if self._use_tls else f"ws://{self.host}"
             parsed_url = urlsplit(url)
             url = f"{parsed_url.scheme}://{parsed_url.netloc}:{self.port}{parsed_url.path}"
 
@@ -500,7 +504,7 @@ async def connect(self):
                 autoclose=False,
                 proxy=http_proxy_host,
                 proxy_auth=http_proxy_auth,
-                ssl=self.sslopts,
+                ssl=self.sslopts if self._use_tls else None,
                 heartbeat=DEFAULT_WEBSOCKET_HEARTBEAT_SECONDS,
             )
         except ClientConnectorError as exc:

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

@@ -214,6 +214,9 @@ def __init__(self, hostname, **kwargs):
         self._custom_endpoint_address = kwargs.get("custom_endpoint_address")
         self._connection_verify = kwargs.get("connection_verify")
 
+        # Emulator
+        self._use_tls: bool = kwargs.get("use_tls", True)
+
     def __enter__(self):
         """Run Client in a context manager.
 
@@ -312,7 +315,7 @@ def open(self, connection=None):
             self._external_connection = True
         elif not self._connection:
             self._connection = Connection(
-                "amqps://" + self._hostname,
+                "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()},
                 container_id=self._name,
@@ -325,6 +328,7 @@ def open(self, connection=None):
                 http_proxy=self._http_proxy,
                 custom_endpoint_address=self._custom_endpoint_address,
                 socket_timeout=self._socket_timeout,
+                use_tls=self._use_tls,
             )
             self._connection.open()
         if not self._session:

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

@@ -351,6 +351,7 @@ def create_send_client(# pylint: disable=unused-argument
             link_properties=link_properties,
             properties=properties,
             client_name=client_name,
+            use_tls=config.use_tls,
             **kwargs,
         )
 
@@ -507,6 +508,7 @@ def create_receive_client(
             keep_alive_interval=keep_alive_interval,
             streaming_receive=streaming_receive,
             timeout=timeout,
+            use_tls=config.use_tls,
             **kwargs,
         )
 
@@ -604,6 +606,7 @@ def create_mgmt_client(
             http_proxy=config.http_proxy,
             custom_endpoint_address=config.custom_endpoint_address,
             connection_verify=config.connection_verify,
+            use_tls=config.use_tls,
         )
 
     @staticmethod