Added websocket client handler, and websocket-client module

README.md

@@ -153,6 +153,20 @@ from fastapi_websocket_rpc.logger import logging_config, LoggingModes
 logging_config.set_mode(LoggingModes.UVICORN)
 ```
 
+## HTTP(S) Proxy
+By default, fastapi-websocket-rpc uses websockets module as websocket client handler. This does not support HTTP(S) Proxy, see https://github.com/python-websockets/websockets/issues/364 . If proxy is important, others websocket client implementation can be used, such as websocket-client module (https://websocket-client.readthedocs.io). Here is how to use it, by defining websocket_client_handler_cls parameter:
+
+```python
+import asyncio
+from fastapi_websocket_rpc import RpcMethodsBase, WebSocketRpcClient, WebSocketClientHandler
+
+async def run_client(uri):
+    async with WebSocketRpcClient(uri, RpcMethodsBase(), websocket_client_handler_cls = WebSocketClientHandler) as client:
+```
+
+Just set standard environment variables (lowercase and uppercase works): http_proxy, https_proxy, and no_proxy before running python script.
+
+
 ## Pull requests - welcome!
 - Please include tests for new features 
 

fastapi_websocket_rpc/__init__.py

@@ -1,5 +1,7 @@
 from .rpc_methods import RpcMethodsBase, RpcUtilityMethods
 from .websocket_rpc_client import WebSocketRpcClient
+from .websocket_rpc_client import WebSocketClientHandler
+from .websocket_rpc_client import WebSocketsClientHandler
 from .websocket_rpc_endpoint import WebsocketRPCEndpoint
 from .rpc_channel import RpcChannel
 from .schemas import WebSocketFrameType

fastapi_websocket_rpc/simplewebsocket.py

@@ -9,6 +9,10 @@ class SimpleWebSocket(ABC):
     Abstract base class for all websocket related wrappers.
     """
 
+    @abstractmethod
+    def connect(self, uri: str, **connect_kwargs):
+        pass
+
     @abstractmethod
     def send(self, msg):
         pass
@@ -21,11 +25,22 @@ def recv(self):
     def close(self, code: int = 1000):
         pass
 
+    @abstractmethod
+    def handle_exception(self, exception: Exception):
+        pass
+
+    @abstractmethod
+    def isConnectionClosedException(self, exception: Exception) -> bool:
+        pass
+
 
 class JsonSerializingWebSocket(SimpleWebSocket):
     def __init__(self, websocket: SimpleWebSocket):
         self._websocket = websocket
 
+    async def connect(self, uri: str, **connect_kwargs):
+        await self._websocket.connect(uri, **connect_kwargs)
+
     def _serialize(self, msg):
         return pydantic_serialize(msg)
 
@@ -42,3 +57,12 @@ async def recv(self):
 
     async def close(self, code: int = 1000):
         await self._websocket.close(code)
+
+    async def handle_exception(self, exception: Exception):
+        await self._websocket.handle_exception(exception)
+
+    async def isConnectionClosedException(self, exception: Exception) -> bool:
+        return await self._websocket.isConnectionClosedException(exception)
+
+
+

fastapi_websocket_rpc/websocket_rpc_client.py

@@ -5,6 +5,7 @@
 import tenacity
 from tenacity.retry import retry_if_exception
 
+import websocket
 import websockets
 from websockets.exceptions import InvalidStatusCode, WebSocketException, ConnectionClosedError, ConnectionClosedOK
 
@@ -16,6 +17,145 @@
 logger = get_logger("RPC_CLIENT")
 
 
+class WebSocketClientHandler(SimpleWebSocket):
+    """
+    Handler that use https://websocket-client.readthedocs.io/en/latest module.
+    This implementation supports HTTP proxy, though HTTP_PROXY and HTTPS_PROXY environment variable.
+    This is not documented, but in code, see https://github.com/websocket-client/websocket-client/blob/master/websocket/_url.py#L163
+    The module is not written as coroutine: https://websocket-client.readthedocs.io/en/latest/threading.html#asyncio-library-usage, so 
+    as a workaround, the send/recv are called in "run_in_executor" method.
+    TODO: remove this implementation after https://github.com/python-websockets/websockets/issues/364 is fixed and use WebSocketsClientHandler instead.
+
+    Note: the connect timeout, if not specified, is the default socket connect timeout, which could be around 2min, so a bit longer than WebSocketsClientHandler.
+    """
+    def __init__(self):
+        self._websocket = None
+
+    """
+    Args:
+        **kwargs: Additional args passed to connect
+            https://websocket-client.readthedocs.io/en/latest/examples.html#connection-options
+            https://websocket-client.readthedocs.io/en/latest/core.html#websocket._core.WebSocket.connect
+    """
+    async def connect(self, uri: str, **connect_kwargs):
+        self._websocket = await asyncio.get_event_loop().run_in_executor(None, websocket.create_connection, uri, **connect_kwargs)
+
+    async def send(self, msg):
+        if self._websocket is None:
+            # connect must be called before.
+            logging.error("Websocket connect() must be called before.")
+        await asyncio.get_event_loop().run_in_executor(None, self._websocket.send, msg)
+
+    async def recv(self):
+        if self._websocket is None:
+            # connect must be called before.
+            logging.error("Websocket connect() must be called before.")
+        msg = await asyncio.get_event_loop().run_in_executor(None, self._websocket.recv)
+        return msg
+
+    async def close(self, code: int = 1000):
+        if self._websocket is not None:
+            # Case opened, we have something to close.
+            await self._websocket.close(code)
+
+    async def handle_exception(self, exception: Exception):
+        try:
+            raise exception
+        # See https://websocket-client.readthedocs.io/en/latest/exceptions.html
+        except websocket._exceptions.WebSocketAddressException:
+            logger.info("websocket address info cannot be found")
+            raise
+        except websocket._exceptions.WebSocketBadStatusException:
+            logger.info("bad handshake status code")
+            raise
+        except websocket._exceptions.WebSocketConnectionClosedException:
+            logger.info("remote host closed the connection or some network error happened")
+            raise
+        except websocket._exceptions.WebSocketPayloadException:
+            logger.info(
+                f"WebSocket payload is invalid")
+            raise
+        except websocket._exceptions.WebSocketProtocolException:
+            logger.info(f"WebSocket protocol is invalid")
+            raise
+        except websocket._exceptions.WebSocketProxyException:
+            logger.info(f"proxy error occurred")
+            raise
+        except OSError as err:
+            logger.info("RPC Connection failed - %s", err)
+            raise
+        except Exception as err:
+            logger.exception("RPC Error")
+            raise
+
+    async def isConnectionClosedException(self, exception: Exception) -> bool:
+        # websocket.WebSocketConnectionClosedException means remote host closed the connection or some network error happened
+        return isinstance(exception, websocket.WebSocketConnectionClosedException)
+
+
+class WebSocketsClientHandler(SimpleWebSocket):
+    """
+    Handler that use https://websockets.readthedocs.io/en/stable module.
+    This implementation does not support HTTP proxy (see https://github.com/python-websockets/websockets/issues/364).
+    """
+    def __init__(self):
+        self._websocket = None
+
+    """
+    Args:
+        **kwargs: Additional args passed to connect
+            https://websockets.readthedocs.io/en/stable/reference/asyncio/client.html#opening-a-connection
+    """
+    async def connect(self, uri: str, **connect_kwargs):
+        self._websocket = await websockets.connect(uri, **connect_kwargs)
+
+    async def send(self, msg):
+        if self._websocket is None:
+            # connect must be called before.
+            logging.error("Websocket connect() must be called before.")
+        await self._websocket.send(msg)
+
+    async def recv(self):
+        if self._websocket is None:
+            # connect must be called before.
+            logging.error("Websocket connect() must be called before.")
+        msg = await self._websocket.recv()
+        return msg
+
+    async def close(self, code: int = 1000):
+        if self._websocket is not None:
+            # Case opened, we have something to close.
+            await self._websocket.close(code)
+
+    async def handle_exception(self, exception: Exception):
+        try:
+            raise exception
+        except ConnectionRefusedError:
+            logger.info("RPC connection was refused by server")
+            raise
+        except ConnectionClosedError:
+            logger.info("RPC connection lost")
+            raise
+        except ConnectionClosedOK:
+            logger.info("RPC connection closed")
+            raise
+        except InvalidStatusCode as err:
+            logger.info(
+                f"RPC Websocket failed - with invalid status code {err.status_code}")
+            raise
+        except WebSocketException as err:
+            logger.info(f"RPC Websocket failed - with {err}")
+            raise
+        except OSError as err:
+            logger.info("RPC Connection failed - %s", err)
+            raise
+        except Exception as err:
+            logger.exception("RPC Error")
+            raise
+
+    async def isConnectionClosedException(self, exception: Exception) -> bool:
+        return isinstance(exception, websockets.exceptions.ConnectionClosed)
+
 def isNotInvalidStatusCode(value):
     return not isinstance(value, InvalidStatusCode)
 
@@ -59,6 +199,7 @@ def __init__(self, uri: str, methods: RpcMethodsBase = None,
                  on_disconnect: List[OnDisconnectCallback] = None,
                  keep_alive: float = 0,
                  serializing_socket_cls: Type[SimpleWebSocket] = JsonSerializingWebSocket,
+                 websocket_client_handler_cls: Type[SimpleWebSocket] = WebSocketsClientHandler,
                  **kwargs):
         """
         Args:
@@ -71,8 +212,7 @@ def __init__(self, uri: str, methods: RpcMethodsBase = None,
             on_disconnect (List[Coroutine]): callbacks on connection termination (each callback is called with the channel)
             keep_alive(float): interval in seconds to send a keep-alive ping, Defaults to 0, which means keep alive is disabled.
 
-            **kwargs: Additional args passed to connect (@see class Connect at websockets/client.py)
-                      https://websockets.readthedocs.io/en/stable/api.html#websockets.client.connect
+            **kwargs: Additional args passed to connect, depends on websocket_client_handler_cls
 
 
             usage:
@@ -105,15 +245,18 @@ def __init__(self, uri: str, methods: RpcMethodsBase = None,
         self._on_connect = on_connect
         # serialization
         self._serializing_socket_cls = serializing_socket_cls
+        # websocket client implementation
+        self._websocket_client_handler_cls = websocket_client_handler_cls
 
     async def __connect__(self):
         try:
             try:
                 logger.info(f"Trying server - {self.uri}")
                 # Start connection
-                raw_ws = await websockets.connect(self.uri, **self.connect_kwargs)
+                raw_ws = self._websocket_client_handler_cls()
                 # Wrap socket in our serialization class
                 self.ws = self._serializing_socket_cls(raw_ws)
+                await self.ws.connect(self.uri, **self.connect_kwargs)
                 # Init an RPC channel to work on-top of the connection
                 self.channel = RpcChannel(
                     self.methods, self.ws, default_response_timeout=self.default_response_timeout)
@@ -137,28 +280,13 @@ async def __connect__(self):
                     await self.channel.close()
                 self.cancel_tasks()
                 raise
-        except ConnectionRefusedError:
-            logger.info("RPC connection was refused by server")
-            raise
-        except ConnectionClosedError:
-            logger.info("RPC connection lost")
-            raise
-        except ConnectionClosedOK:
-            logger.info("RPC connection closed")
-            raise
-        except InvalidStatusCode as err:
-            logger.info(
-                f"RPC Websocket failed - with invalid status code {err.status_code}")
-            raise
-        except WebSocketException as err:
-            logger.info(f"RPC Websocket failed - with {err}")
-            raise
-        except OSError as err:
-            logger.info("RPC Connection failed - %s", err)
-            raise
         except Exception as err:
-            logger.exception("RPC Error")
-            raise
+            if self.ws is not None:
+                # Exception could be websocket client specific.
+                await self.ws.handle_exception(err)
+            else:
+                logger.exception("RPC Error")
+                raise
 
     async def __aenter__(self):
         if self.retry_config is False:
@@ -205,12 +333,13 @@ async def reader(self):
         # task was canceled
         except asyncio.CancelledError:
             pass
-        except websockets.exceptions.ConnectionClosed:
-            logger.info("Connection was terminated.")
-            await self.close()
-        except:
-            logger.exception("RPC Reader task failed")
-            raise
+        except Exception as err:
+            if self.ws.isConnectionClosedException(err):
+                logger.info("Connection was terminated.")
+                await self.close()
+            else:
+                logger.exception("RPC Reader task failed")
+                raise
 
     async def _keep_alive(self):
         try:

requirements.txt

@@ -3,4 +3,5 @@ packaging>=20.4
 pydantic>=1.9.1
 uvicorn>=0.17.6,<1
 websockets>=10.3,<11
+websocket-client>=1.6.4
 tenacity>=8.0.1,<9