Update the WebSocketTransport to detect upgrade to TLS (#443)

Update the WebSocketTransport to detect the attempt to upgrade the connection to TLS.

If the connection is already using TLS then continue. Otherwise raise a ProtocolError.

Fixes #442

nats/aio/transport.py

@@ -10,6 +10,8 @@
 except ImportError:
     aiohttp = None  # type: ignore[assignment]
 
+from nats.errors import ProtocolError
+
 
 class Transport(abc.ABC):
 
@@ -197,6 +199,7 @@ def __init__(self):
         self._client: aiohttp.ClientSession = aiohttp.ClientSession()
         self._pending = asyncio.Queue()
         self._close_task = asyncio.Future()
+        self._using_tls: bool | None = None
 
     async def connect(
         self, uri: ParseResult, buffer_size: int, connect_timeout: int
@@ -205,6 +208,7 @@ async def connect(
         self._ws = await self._client.ws_connect(
             uri.geturl(), timeout=connect_timeout
         )
+        self._using_tls = False
 
     async def connect_tls(
         self,
@@ -213,11 +217,17 @@ async def connect_tls(
         buffer_size: int,
         connect_timeout: int,
     ):
+        if self._ws and not self._ws.closed:
+            if self._using_tls:
+                return
+            raise ProtocolError("ws: cannot upgrade to TLS")
+
         self._ws = await self._client.ws_connect(
             uri if isinstance(uri, str) else uri.geturl(),
             ssl=ssl_context,
             timeout=connect_timeout
         )
+        self._using_tls = True
 
     def write(self, payload):
         self._pending.put_nowait(payload)

tests/utils.py

@@ -477,7 +477,9 @@ def setUp(self):
         )
 
         server = NATSD(
-            port=4222, config_file=get_config_file("conf/ws_tls.conf")
+            port=4222,
+            tls=True,
+            config_file=get_config_file("conf/ws_tls.conf")
         )
         self.server_pool.append(server)
         for natsd in self.server_pool: