diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py index cc84b642fba4..9d8c445bd231 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py @@ -185,7 +185,7 @@ async def _disconnect(self, *args) -> None: if self.state == ConnectionState.END: return await self._set_state(ConnectionState.END) - self._transport.close() + await self._transport.close() def _can_read(self): # type: () -> bool diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py index cd4e8d49514a..b53971ddb203 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py @@ -37,9 +37,12 @@ import socket import ssl import struct -from ssl import SSLError +from ssl import SSLContext, SSLError +from contextlib import contextmanager from io import BytesIO import logging +from threading import Lock +from typing import Optional import certifi @@ -57,6 +60,7 @@ set_cloexec, AMQP_PORT, TIMEOUT_INTERVAL, + WebSocketTransport, ) @@ -145,37 +149,6 @@ async def send_frame(self, channel, frame, **kwargs): # _LOGGER.info("OCH%d -> %r", channel, frame) -class AsyncTransport(AsyncTransportMixin): - """Common superclass for TCP and SSL transports.""" - - def __init__( - self, - host, - port=AMQP_PORT, - connect_timeout=None, - read_timeout=None, - write_timeout=None, - ssl=False, - socket_settings=None, - raise_on_initial_eintr=True, - **kwargs - ): - self.connected = False - self.sock = None - self.reader = None - self.writer = None - self.raise_on_initial_eintr = raise_on_initial_eintr - self._read_buffer = BytesIO() - self.host, self.port = to_host_port(host, port) - - self.connect_timeout = connect_timeout - self.read_timeout = read_timeout - self.write_timeout = write_timeout - self.socket_settings = socket_settings - self.loop = get_running_loop() - self.socket_lock = asyncio.Lock() - self.sslopts = self._build_ssl_opts(ssl) - def _build_ssl_opts(self, sslopts): if sslopts in [True, False, None, {}]: return sslopts @@ -214,6 +187,40 @@ def _build_ssl_context(self, sslopts, check_hostname=None, **ctx_options): ctx.check_hostname = check_hostname return ctx + +class AsyncTransport(AsyncTransportMixin): + """Common superclass for TCP and SSL transports.""" + + def __init__( + self, + host, + port=AMQP_PORT, + connect_timeout=None, + read_timeout=None, + write_timeout=None, + ssl=False, + socket_settings=None, + raise_on_initial_eintr=True, + **kwargs, + ): + self.connected = False + self.sock = None + self.reader = None + self.writer = None + self.raise_on_initial_eintr = raise_on_initial_eintr + self._read_buffer = BytesIO() + self.host, self.port = to_host_port(host, port) + + self.connect_timeout = connect_timeout + self.read_timeout = read_timeout + self.write_timeout = write_timeout + self.socket_settings = socket_settings + self.loop = get_running_loop() + self.socket_lock = asyncio.Lock() + self.sslopts = self._build_ssl_opts(ssl) + + + async def connect(self): try: # are we already connected? @@ -263,8 +270,7 @@ async def _connect(self, host, port, timeout): # if getaddrinfo succeeded before for another address # family, reraise the previous socket.error since it's more # relevant to users - raise e if e is not None else socket.error("failed to resolve broker hostname") - continue # pragma: no cover + raise (e if e is not None else socket.error("failed to resolve broker hostname")) # now that we have address(es) for the hostname, connect to broker for i, res in enumerate(entries): @@ -378,7 +384,7 @@ async def _write(self, s): """Write a string out to the SSL socket fully.""" self.writer.write(s) - def close(self): + async def close(self): if self.writer is not None: if self.sslopts: # see issue: https://github.com/encode/httpx/issues/914 @@ -426,74 +432,81 @@ async def negotiate(self): class WebSocketTransportAsync(AsyncTransportMixin): def __init__(self, host, port=WEBSOCKET_PORT, connect_timeout=None, ssl=None, **kwargs): self._read_buffer = BytesIO() - self.loop = get_running_loop() self.socket_lock = asyncio.Lock() - self.sslopts = ssl if isinstance(ssl, dict) else {} + self.sslopts = self._build_ssl_opts(ssl) if isinstance(ssl, dict) else None self._connect_timeout = connect_timeout or TIMEOUT_INTERVAL self._custom_endpoint = kwargs.get("custom_endpoint") self.host = host self.ws = None + self.session = None self._http_proxy = kwargs.get("http_proxy", None) - async def connect(self): - http_proxy_host, http_proxy_port, http_proxy_auth = None, None, None + username, password = None, None + http_proxy_host, http_proxy_port = None, None + http_proxy_auth = None + if self._http_proxy: http_proxy_host = self._http_proxy["proxy_hostname"] http_proxy_port = self._http_proxy["proxy_port"] + if http_proxy_host and http_proxy_port: + http_proxy_host = f"{http_proxy_host}:{http_proxy_port}" username = self._http_proxy.get("username", None) password = self._http_proxy.get("password", None) - if username or password: - http_proxy_auth = (username, password) + try: - from websocket import create_connection + from aiohttp import ClientSession - self.ws = create_connection( + if username or password: + from aiohttp import BasicAuth + http_proxy_auth = BasicAuth(login=username, password=password) + + self.session = ClientSession() + self.ws = await self.session.ws_connect( url="wss://{}".format(self._custom_endpoint or self.host), - subprotocols=[AMQP_WS_SUBPROTOCOL], timeout=self._connect_timeout, - skip_utf8_validation=True, - sslopt=self.sslopts, - http_proxy_host=http_proxy_host, - http_proxy_port=http_proxy_port, - http_proxy_auth=http_proxy_auth, + protocols=[AMQP_WS_SUBPROTOCOL], + autoclose=False, + proxy=http_proxy_host, + proxy_auth=http_proxy_auth, + ssl=self.sslopts, ) + except ImportError: - raise ValueError("Please install websocket-client library to use websocket transport.") + raise ValueError("Please install aiohttp library to use websocket transport.") async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-arguments """Read exactly n bytes from the peer.""" - from websocket import WebSocketTimeoutException length = 0 view = buffer or memoryview(bytearray(n)) nbytes = self._read_buffer.readinto(view) length += nbytes n -= nbytes + try: while n: - data = await self.loop.run_in_executor(None, self.ws.recv) - + data = await self.ws.receive_bytes() if len(data) <= n: view[length : length + len(data)] = data n -= len(data) else: view[length : length + n] = data[0:n] self._read_buffer = BytesIO(data[n:]) - n = 0 - + n = 0 return view - except WebSocketTimeoutException: + except (asyncio.TimeoutError) as wex: raise TimeoutError() - def close(self): + async def close(self): """Do any preliminary work in shutting down the connection.""" - self.ws.close() + await self.ws.close() + await self.session.close() self.connected = False async def write(self, s): - """Completely write a string to the peer. + """Completely write a string (byte array) to the peer. ABNF, OPCODE_BINARY = 0x2 See http://tools.ietf.org/html/rfc5234 http://tools.ietf.org/html/rfc6455#section-5.2 - """ - await self.loop.run_in_executor(None, self.ws.send_binary, s) + """ + await self.ws.send_bytes(s)