wslink 1.12.3__py3-none-any.whl → 2.0.0__py3-none-any.whl

wslink/backends/aiohttp/__init__.py

@@ -15,6 +15,8 @@ import aiohttp.web as aiohttp_web
 MAX_MSG_SIZE = int(os.environ.get("WSLINK_MAX_MSG_SIZE", 4194304))
 HEART_BEAT = int(os.environ.get("WSLINK_HEART_BEAT", 30))  # 30 seconds
 
+logger = logging.getLogger(__name__)
+
 # -----------------------------------------------------------------------------
 # HTTP helpers
 # -----------------------------------------------------------------------------
@@ -98,20 +100,20 @@ class WebAppServer(AbstractWebApp):
             self.app, handle_signals=self.handle_signals
         )
 
-        logging.info("awaiting runner setup")
+        logger.info("awaiting runner setup")
         await self._runner.setup()
 
         self._site = aiohttp_web.TCPSite(
             self._runner, self.host, self.port, ssl_context=self.ssl_context
         )
 
-        logging.info("awaiting site startup")
+        logger.info("awaiting site startup")
         await self._site.start()
 
         if port_callback is not None:
             port_callback(self.get_port())
 
-        logging.info(f"Print WSLINK_READY_MSG")
+        logger.info("Print WSLINK_READY_MSG")
         STARTUP_MSG = os.environ.get("WSLINK_READY_MSG", "wslink: Starting factory")
         if STARTUP_MSG:
             # Emit an expected log message so launcher.py knows we've started up.
@@ -119,10 +121,10 @@ class WebAppServer(AbstractWebApp):
             # We've seen some issues with stdout buffering - be conservative.
             sys.stdout.flush()
 
-        logging.info(f"Schedule auto shutdown with timout {self.timeout}")
+        logger.info(f"Schedule auto shutdown with timout {self.timeout}")
         self.shutdown_schedule()
 
-        logging.info("awaiting running future")
+        logger.info("awaiting running future")
         await self.completion
 
     async def stop(self):
@@ -133,12 +135,12 @@ class WebAppServer(AbstractWebApp):
         # Neither site.stop() nor runner.cleanup() actually stop the server
         # as documented, but at least runner.cleanup() results in the
         # "on_shutdown" signal getting sent.
-        logging.info("Performing runner.cleanup()")
+        logger.info("Performing runner.cleanup()")
         await self.runner.cleanup()
 
         # So to actually stop the server, the workaround is just to resolve
         # the future we awaited in the start method.
-        logging.info("Stopping server")
+        logger.info("Stopping server")
         self.completion.set_result(True)
 
 
@@ -163,7 +165,7 @@ class ReverseWebAppServer(AbstractWebApp):
 
 def create_webserver(server_config):
     if "logging_level" in server_config and server_config["logging_level"]:
-        logging.basicConfig(level=server_config["logging_level"])
+        logging.getLogger("wslink").setLevel(server_config["logging_level"])
 
     # Shortcut for reverse connection
     if "reverse_url" in server_config:
@@ -184,10 +186,10 @@ def is_binary(msg):
 
 class AioHttpWsHandler(WslinkHandler):
     async def disconnectClients(self):
-        logging.info("Closing client connections:")
+        logger.info("Closing client connections:")
         keys = list(self.connections.keys())
         for client_id in keys:
-            logging.info("  {0}".format(client_id))
+            logger.info("  {0}".format(client_id))
             ws = self.connections[client_id]
             await ws.close(
                 code=aiohttp.WSCloseCode.GOING_AWAY, message="Server shutdown"
@@ -202,7 +204,7 @@ class AioHttpWsHandler(WslinkHandler):
         )
         self.connections[client_id] = current_ws
 
-        logging.info("client {0} connected".format(client_id))
+        logger.info("client {0} connected".format(client_id))
 
         self.web_app.shutdown_cancel()
 
@@ -217,31 +219,31 @@ class AioHttpWsHandler(WslinkHandler):
             del self.connections[client_id]
             self.authentified_client_ids.discard(client_id)
 
-            logging.info("client {0} disconnected".format(client_id))
+            logger.info("client {0} disconnected".format(client_id))
 
             if not self.connections:
-                logging.info("No more connections, scheduling shutdown")
+                logger.info("No more connections, scheduling shutdown")
                 self.web_app.shutdown_schedule()
 
         return current_ws
 
     async def reverse_connect_to(self, url):
-        logging.debug("reverse_connect_to: running with url %s", url)
+        logger.debug("reverse_connect_to: running with url %s", url)
         client_id = self.reverse_connection_client_id
         async with aiohttp.ClientSession() as session:
-            logging.debug("reverse_connect_to: client session started")
+            logger.debug("reverse_connect_to: client session started")
             async with session.ws_connect(url) as current_ws:
-                logging.debug("reverse_connect_to: ws started")
+                logger.debug("reverse_connect_to: ws started")
                 self.connections[client_id] = current_ws
-                logging.debug("reverse_connect_to: onConnect")
+                logger.debug("reverse_connect_to: onConnect")
                 await self.onConnect(url, client_id)
 
                 async for msg in current_ws:
                     if not current_ws.closed:
                         await self.onMessage(is_binary(msg), msg, client_id)
 
-                logging.debug("reverse_connect_to: onClose")
+                logger.debug("reverse_connect_to: onClose")
                 await self.onClose(client_id)
                 del self.connections[client_id]
 
-        logging.debug("reverse_connect_to: exited")
+        logger.debug("reverse_connect_to: exited")

wslink/backends/aiohttp/launcher.py

@@ -23,6 +23,8 @@ from wslink.launcher import (
     STATUS_NOT_FOUND,
 )
 
+logger = logging.getLogger(__name__)
+
 # ===========================================================================
 # Launcher ENV configuration
 # ===========================================================================
@@ -50,7 +52,7 @@ class LauncherResource(object):
     def __del__(self):
         try:
             # causes an exception when server is killed with Ctrl-C
-            logging.warning("Server factory shutting down. Stopping all processes")
+            logger.warning("Server factory shutting down. Stopping all processes")
         except:
             pass
 
@@ -141,17 +143,17 @@ class LauncherResource(object):
 
         if not id:
             message = "id not provided in GET request"
-            logging.error(message)
+            logger.error(message)
             return aiohttp_web.json_response(
                 {"error": message}, status=STATUS_BAD_REQUEST
             )
 
-        logging.info("GET request received for id: %s" % id)
+        logger.info("GET request received for id: %s" % id)
 
         session = self.session_manager.getSession(id)
         if not session:
             message = "No session with id: %s" % id
-            logging.error(message)
+            logger.error(message)
             return aiohttp_web.json_response(
                 {"error": message}, status=STATUS_BAD_REQUEST
             )
@@ -170,17 +172,17 @@ class LauncherResource(object):
 
         if not id:
             message = "id not provided in DELETE request"
-            logging.error(message)
+            logger.error(message)
             return aiohttp_web.json_response(
                 {"error": message}, status=STATUS_BAD_REQUEST
             )
 
-        logging.info("DELETE request received for id: %s" % id)
+        logger.info("DELETE request received for id: %s" % id)
 
         session = self.session_manager.getSession(id)
         if not session:
             message = "No session with id: %s" % id
-            logging.error(message)
+            logger.error(message)
             return aiohttp_web.json_response(
                 {"error": message}, status=STATUS_NOT_FOUND
             )
@@ -190,7 +192,7 @@ class LauncherResource(object):
         self.process_manager.stopProcess(id)
 
         message = "Deleted session with id: %s" % id
-        logging.info(message)
+        logger.info(message)
 
         return aiohttp_web.json_response(session, status=STATUS_OK)
 
@@ -213,15 +215,17 @@ def startWebServer(options, config):
     # Setup logging
     logFileName = log_dir + os.sep + "launcherLog.log"
     formatting = "%(asctime)s:%(levelname)s:%(name)s:%(message)s"
-    logging.basicConfig(
-        level=logging.DEBUG, filename=logFileName, filemode="w", format=formatting
-    )
+    # create file handler which logs even debug messages
+    fh = logging.FileHandler(logFileName, mode="w")
+    fh.setLevel(logging.DEBUG)
+    fh.setFormatter(logging.Formatter(formatting))
+    logging.getLogger("wslink").addHandler(fh)
     if options.debug:
         console = logging.StreamHandler(sys.stdout)
         console.setLevel(logging.INFO)
         formatter = logging.Formatter(formatting)
         console.setFormatter(formatter)
-        logging.getLogger("").addHandler(console)
+        logging.getLogger("wslink").addHandler(console)
 
     web_app = aiohttp_web.Application()
 

wslink/backends/aiohttp/relay.py

@@ -377,8 +377,6 @@ def main(host=None, port=None, www_path=None, proxy_route=None, mode=None):
     if proxy_route is None:
         proxy_route = args.proxy_route
 
-    logging.basicConfig()
-
     # Manage routes
     routes = []
 

wslink/backends/generic/core.py

@@ -125,7 +125,7 @@ def startWebServer(*args, **kwargs):
 
 def create_webserver(server_config):
     if "logging_level" in server_config and server_config["logging_level"]:
-        logging.basicConfig(level=server_config["logging_level"])
+        logging.getLogger("wslink").setLevel(server_config["logging_level"])
 
     # Reverse connection
     if "reverse_url" in server_config:

wslink/chunking.py

@@ -0,0 +1,214 @@
+import sys
+import secrets
+import msgpack
+from typing import Dict, Tuple, Union
+if sys.version_info >= (3, 8):
+    from typing import TypedDict  # pylint: disable=no-name-in-module
+else:
+    from typing_extensions import TypedDict
+
+UINT32_LENGTH = 4
+ID_LOCATION = 0
+ID_LENGTH = UINT32_LENGTH
+MESSAGE_OFFSET_LOCATION = ID_LOCATION + ID_LENGTH
+MESSAGE_OFFSET_LENGTH = UINT32_LENGTH
+MESSAGE_SIZE_LOCATION = MESSAGE_OFFSET_LOCATION + MESSAGE_OFFSET_LENGTH
+MESSAGE_SIZE_LENGTH = UINT32_LENGTH
+
+HEADER_LENGTH = ID_LENGTH + MESSAGE_OFFSET_LENGTH + MESSAGE_SIZE_LENGTH
+
+
+def _encode_header(id: bytes, offset: int, size: int) -> bytes:
+    return (
+        id
+        + offset.to_bytes(MESSAGE_OFFSET_LENGTH, "little", signed=False)
+        + size.to_bytes(MESSAGE_SIZE_LENGTH, "little", signed=False)
+    )
+
+
+def _decode_header(header: bytes) -> Tuple[bytes, int, int]:
+    id = header[ID_LOCATION:ID_LENGTH]
+    offset = int.from_bytes(
+        header[
+            MESSAGE_OFFSET_LOCATION : MESSAGE_OFFSET_LOCATION + MESSAGE_OFFSET_LENGTH
+        ],
+        "little",
+        signed=False,
+    )
+    size = int.from_bytes(
+        header[MESSAGE_SIZE_LOCATION : MESSAGE_SIZE_LOCATION + MESSAGE_SIZE_LENGTH],
+        "little",
+        signed=False,
+    )
+    return id, offset, size
+
+
+def generate_chunks(message: bytes, max_size: int):
+    total_size = len(message)
+
+    if max_size == 0:
+        max_content_size = total_size
+    else:
+        max_content_size = max(max_size - HEADER_LENGTH, 1)
+
+    id = secrets.token_bytes(ID_LENGTH)
+
+    offset = 0
+
+    while offset < total_size:
+        header = _encode_header(id, offset, total_size)
+        chunk_content = message[offset : offset + max_content_size]
+
+        yield header + chunk_content
+
+        offset += max_content_size
+
+    return
+
+
+class PendingMessage(TypedDict):
+    received_size: int
+    content: bytearray
+
+
+# This un-chunker is vulnerable to DOS.
+# If it receives a message with a header claiming a large incoming message
+# it will allocate the memory blindly even without actually receiving the content
+# Chunks for a given message can come in any order
+# Chunks across messages can be interleaved.
+class UnChunker:
+    pending_messages: Dict[bytes, PendingMessage]
+    max_message_size: int
+
+    def __init__(self):
+        self.pending_messages = {}
+        self.max_message_size = 512
+
+    def set_max_message_size(self, size):
+        self.max_message_size = size
+
+    def release_pending_messages(self):
+        self.pending_messages = {}
+
+    def process_chunk(self, chunk: bytes) -> Union[bytes, None]:
+        header, chunk_content = chunk[:HEADER_LENGTH], chunk[HEADER_LENGTH:]
+        id, offset, total_size = _decode_header(header)
+
+        pending_message = self.pending_messages.get(id, None)
+
+        if pending_message is None:
+            if total_size > self.max_message_size:
+                raise ValueError(
+                    f"""Total size for message {id} exceeds the allocation limit allowed.
+    Maximum size = {self.max_message_size},
+    Received size = {total_size}."""
+                )
+
+            pending_message = PendingMessage(
+                received_size=0, content=bytearray(total_size)
+            )
+            self.pending_messages[id] = pending_message
+
+        # This should never happen, but still check it
+        if total_size != len(pending_message["content"]):
+            del self.pending_messages[id]
+            raise ValueError(
+                f"Total size in chunk header for message {id} does not match total size declared by previous chunk."
+            )
+
+        content_size = len(chunk_content)
+        content_view = memoryview(pending_message["content"])
+        content_view[offset : offset + content_size] = chunk_content
+        pending_message["received_size"] += content_size
+
+        if pending_message["received_size"] >= total_size:
+            full_message = pending_message["content"]
+            del self.pending_messages[id]
+            return msgpack.unpackb(bytes(full_message))
+
+        return None
+
+
+class StreamPendingMessage(TypedDict):
+    received_size: int
+    total_size: int
+    unpacker: msgpack.Unpacker
+
+
+# This un-chunker is more memory efficient
+# (each chunk is passed immediately to msgpack)
+# and it will only allocate memory when it receives content.
+# Chunks for a given message are expected to come sequentially
+# Chunks across messages can be interleaved.
+class StreamUnChunker:
+    pending_messages: Dict[bytes, StreamPendingMessage]
+
+    def __init__(self):
+        self.pending_messages = {}
+
+    def set_max_message_size(self, _size):
+        pass
+
+    def release_pending_messages(self):
+        self.pending_messages = {}
+
+    def process_chunk(self, chunk: bytes) -> Union[bytes, None]:
+        header, chunk_content = chunk[:HEADER_LENGTH], chunk[HEADER_LENGTH:]
+        id, offset, total_size = _decode_header(header)
+
+        pending_message = self.pending_messages.get(id, None)
+
+        if pending_message is None:
+            pending_message = StreamPendingMessage(
+                received_size=0,
+                total_size=total_size,
+                unpacker=msgpack.Unpacker(max_buffer_size=total_size),
+            )
+            self.pending_messages[id] = pending_message
+
+        # This should never happen, but still check it
+        if offset != pending_message["received_size"]:
+            del self.pending_messages[id]
+            raise ValueError(
+                f"""Received an unexpected chunk for message {id}.
+    Expected offset = {pending_message['received_size']},
+    Received offset = {offset}."""
+            )
+
+        # This should never happen, but still check it
+        if total_size != pending_message["total_size"]:
+            del self.pending_messages[id]
+            raise ValueError(
+                f"""Received an unexpected total size in chunk header for message {id}.
+    Expected size = {pending_message['total_size']},
+    Received size = {total_size}."""
+            )
+
+        content_size = len(chunk_content)
+        pending_message["received_size"] += content_size
+
+        unpacker = pending_message["unpacker"]
+        unpacker.feed(chunk_content)
+
+        full_message = None
+
+        try:
+            full_message = unpacker.unpack()
+        except msgpack.OutOfData:
+            pass  # message is incomplete, keep ingesting chunks
+
+        if full_message is not None:
+            del self.pending_messages[id]
+
+            if pending_message["received_size"] < total_size:
+                # In principle feeding a stream to the unpacker could yield multiple outputs
+                # for example unpacker.feed(b'0123') would yield b'0', b'1', ect
+                # or concatenated packed payloads would yield two or more unpacked objects
+                # but in our use case we expect a full message to be mapped to a single object
+                raise ValueError(
+                    f"""Received a parsable payload shorter than expected for message {id}.
+    Expected size = {total_size},
+    Received size = {pending_message['received_size']}."""
+                )
+
+        return full_message

wslink/launcher.py

@@ -16,6 +16,8 @@ from random import choice
 
 from wslink import backends
 
+logger = logging.getLogger(__name__)
+
 STATUS_OK = 200
 STATUS_BAD_REQUEST = 400
 STATUS_NOT_FOUND = 404
@@ -205,7 +207,7 @@ def checkSanitize(key_pair, sanitize):
             value = key_pair[key]
             if checkItem["type"] == "inList":
                 if not value in checkItem["list"]:
-                    logging.warning(
+                    logger.warning(
                         "key %s: sanitize %s with default" % (key, key_pair[key])
                     )
                     key_pair[key] = checkItem["default"]
@@ -214,7 +216,7 @@ def checkSanitize(key_pair, sanitize):
                     # User is responsible to add begin- and end- string symbols, to make sure entire string is matched.
                     checkItem["compiled"] = re.compile(checkItem["regexp"])
                 if checkItem["compiled"].match(value) == None:
-                    logging.warning(
+                    logger.warning(
                         "key %s: sanitize %s with default" % (key, key_pair[key])
                     )
                     key_pair[key] = checkItem["default"]
@@ -232,7 +234,7 @@ def replaceVariables(template_str, variable_list, sanitize):
         template_str = item_template.safe_substitute(key_pair)
 
     if "$" in template_str:
-        logging.error("Some properties could not be resolved: " + template_str)
+        logger.error("Some properties could not be resolved: " + template_str)
 
     return template_str
 
@@ -441,8 +443,8 @@ class ProcessManager(object):
                 )
                 self.processes[session["id"]] = proc
             except:
-                logging.error("The command line failed")
-                logging.error(" ".join(map(str, session["cmd"])))
+                logger.error("The command line failed")
+                logger.error(" ".join(map(str, session["cmd"])))
                 return None
 
         return proc

wslink/protocol.py

@@ -1,13 +1,14 @@
 import asyncio
 import copy
 import inspect
-import json
 import logging
-import re
+import msgpack
+import os
 import traceback
 
 from wslink import schedule_coroutine
 from wslink.publish import PublishManager
+from wslink.chunking import generate_chunks, UnChunker
 
 # from http://www.jsonrpc.org/specification, section 5.1
 METHOD_NOT_FOUND = -32601
@@ -17,6 +18,11 @@ RESULT_SERIALIZE_ERROR = -32002
 # used in client JS code:
 CLIENT_ERROR = -32099
 
+# 4MB is the default inside aiohttp
+MAX_MSG_SIZE = int(os.environ.get("WSLINK_MAX_MSG_SIZE", 4194304))
+
+logger = logging.getLogger(__name__)
+
 
 class AbstractWebApp:
     def __init__(self, server_config):
@@ -108,7 +114,7 @@ class AbstractWebApp:
 
     def shutdown_cancel(self):
         if self._shutdown_task is not None:
-            logging.info("Canceling shutdown task")
+            logger.info("Canceling shutdown task")
             self._shutdown_task.cancel()
             self._shutdown_task = None
 
@@ -145,6 +151,7 @@ class WslinkHandler(object):
         self.authentified_client_ids = set()
         self.attachment_atomic = asyncio.Lock()
         self.pub_manager = PublishManager()
+        self.unchunkers = {}
 
         # Build the rpc method dictionary, assuming we were given a serverprotocol
         if self.getServerProtocol():
@@ -181,6 +188,8 @@ class WslinkHandler(object):
         return "reverse_connection_client_id"
 
     async def onConnect(self, request, client_id):
+        self.unchunkers[client_id] = UnChunker()
+
         if not self.serverProtocol:
             return
         if hasattr(self.serverProtocol, "onConnect"):
@@ -190,6 +199,8 @@ class WslinkHandler(object):
                 linkProtocol.onConnect(request, client_id)
 
     async def onClose(self, client_id):
+        del self.unchunkers[client_id]
+
         if not self.serverProtocol:
             return
         if hasattr(self.serverProtocol, "onClose"):
@@ -210,9 +221,16 @@ class WslinkHandler(object):
                     and await self.validateToken(args[0]["secret"], client_id)
                 ):
                     self.authentified_client_ids.add(client_id)
+                    # Once a client is authenticated let the unchunker allocate memory unrestricted
+                    self.unchunkers[client_id].set_max_message_size(
+                        4 * 1024 * 1024 * 1024
+                    )  # 4GB
                     await self.sendWrappedMessage(
                         rpcid,
-                        {"clientID": "c{0}".format(client_id)},
+                        {
+                            "clientID": "c{0}".format(client_id),
+                            "maxMsgSize": MAX_MSG_SIZE,
+                        },
                         client_id=client_id,
                     )
                 else:
@@ -233,34 +251,16 @@ class WslinkHandler(object):
         return False
 
     async def onMessage(self, is_binary, msg, client_id):
-        payload = msg.data
-
-        if is_binary:
-            if self.isClientAuthenticated(client_id):
-                # assume all binary messages are attachments
-                try:
-                    key = self.attachmentsRecvQueue.pop(0)
-                    self.attachmentsReceived[key] = payload
-                except:
-                    pass
-                return
+        if not is_binary:
+            return
 
-        # handles issue https://bugs.python.org/issue10976
-        # `payload` is type bytes in Python 3. Unfortunately, json.loads
-        # doesn't support taking bytes until Python 3.6.
-        if type(payload) is bytes:
-            payload = payload.decode("utf-8")
+        full_message = self.unchunkers[client_id].process_chunk(msg.data)
+        if full_message is not None:
+            await self.onCompleteMessage(full_message, client_id)
 
-        rpc = json.loads(payload)
-        logging.debug("wslink incoming msg %s" % self.payloadWithSecretStripped(rpc))
+    async def onCompleteMessage(self, rpc, client_id):
+        logger.debug("wslink incoming msg %s", self.payloadWithSecretStripped(rpc))
         if "id" not in rpc:
-            # should be a binary attachment header
-            if rpc.get("method") == "wslink.binary.attachment":
-                keys = rpc.get("args", [])
-                if isinstance(keys, list):
-                    for k in keys:
-                        # wait for an attachment by it's order
-                        self.attachmentsRecvQueue.append(k)
             return
 
         # TODO validate
@@ -301,73 +301,37 @@ class WslinkHandler(object):
             return
 
         obj, func = self.functionMap[methodName]
+        args.insert(0, obj)
+
         try:
-            # get any attachments
-            def findAttachments(o):
-                if (
-                    isinstance(o, str)
-                    and re.match(r"^wslink_bin\d+$", o)
-                    and o in self.attachmentsReceived
-                ):
-                    attachment = self.attachmentsReceived[o]
-                    del self.attachmentsReceived[o]
-                    return attachment
-                elif isinstance(o, list):
-                    for i, v in enumerate(o):
-                        o[i] = findAttachments(v)
-                elif isinstance(o, dict):
-                    for k in o:
-                        o[k] = findAttachments(o[k])
-                return o
-
-            args = findAttachments(args)
-            kwargs = findAttachments(kwargs)
-
-            args.insert(0, obj)
-
-            try:
-                self.web_app.last_active_client_id = client_id
-                results = func(*args, **kwargs)
-                if inspect.isawaitable(results):
-                    results = await results
-
-                if self.connections[client_id].closed:
-                    # Connection was closed during RPC call.
-                    return
-
-                await self.sendWrappedMessage(
-                    rpcid, results, method=methodName, client_id=client_id
-                )
-            except Exception as e_inst:
-                captured_trace = traceback.format_exc()
-                logging.error("Exception raised")
-                logging.error(repr(e_inst))
-                logging.error(captured_trace)
-                await self.sendWrappedError(
-                    rpcid,
-                    EXCEPTION_ERROR,
-                    "Exception raised",
-                    {
-                        "method": methodName,
-                        "exception": repr(e_inst),
-                        "trace": captured_trace,
-                    },
-                    client_id=client_id,
-                )
+            self.web_app.last_active_client_id = client_id
+            results = func(*args, **kwargs)
+            if inspect.isawaitable(results):
+                results = await results
+
+            if self.connections[client_id].closed:
+                # Connection was closed during RPC call.
+                return
 
-        except Exception as e:
+            await self.sendWrappedMessage(
+                rpcid, results, method=methodName, client_id=client_id
+            )
+        except Exception as e_inst:
+            captured_trace = traceback.format_exc()
+            logger.error("Exception raised")
+            logger.error(repr(e_inst))
+            logger.error(captured_trace)
             await self.sendWrappedError(
                 rpcid,
                 EXCEPTION_ERROR,
                 "Exception raised",
                 {
                     "method": methodName,
-                    "exception": repr(e),
-                    "trace": traceback.format_exc(),
+                    "exception": repr(e_inst),
+                    "trace": captured_trace,
                 },
                 client_id=client_id,
             )
-            return
 
     def payloadWithSecretStripped(self, payload):
         payload = copy.deepcopy(payload)
@@ -426,9 +390,10 @@ class WslinkHandler(object):
             "id": rpcid,
             "result": content,
         }
+
         try:
-            encMsg = json.dumps(wrapper, ensure_ascii=False)
-        except TypeError as e:
+            packed_wrapper = msgpack.packb(wrapper)
+        except Exception:
             # the content which is not serializable might be arbitrarily large, don't include.
             # repr(content) would do that...
             await self.sendWrappedError(
@@ -442,47 +407,14 @@ class WslinkHandler(object):
 
         websockets = self.getAuthenticatedWebsockets(client_id, skip_last_active_client)
 
-        # Check if any attachments in the map go with this message
-        attachments = self.pub_manager.getAttachmentMap()
-        found_keys = []
-        if attachments:
-            for key in attachments:
-                # string match the encoded attachment key
-                if key in encMsg:
-                    if key not in found_keys:
-                        found_keys.append(key)
-                    # increment  for key
-                    self.pub_manager.registerAttachment(key)
-
-            for key in found_keys:
-                # send header
-                header = {
-                    "wslink": "1.0",
-                    "method": "wslink.binary.attachment",
-                    "args": [key],
-                }
-                json_header = json.dumps(header, ensure_ascii=False)
-
-                # aiohttp can not handle pending ws.send_bytes()
-                # tried with semaphore but got exception with >1
-                # https://github.com/aio-libs/aiohttp/issues/2934
-                async with self.attachment_atomic:
-                    for ws in websockets:
-                        if ws is not None:
-                            # Send binary header
-                            await ws.send_str(json_header)
-                            # Send binary message
-                            await ws.send_bytes(attachments[key])
-
-                # decrement for key
-                self.pub_manager.unregisterAttachment(key)
-
-        for ws in websockets:
-            if ws is not None:
-                await ws.send_str(encMsg)
-
-        loop = asyncio.get_event_loop()
-        loop.call_soon(self.pub_manager.freeAttachments, found_keys)
+        # aiohttp can not handle pending ws.send_bytes()
+        # tried with semaphore but got exception with >1
+        # https://github.com/aio-libs/aiohttp/issues/2934
+        async with self.attachment_atomic:
+            for chunk in generate_chunks(packed_wrapper, MAX_MSG_SIZE):
+                for ws in websockets:
+                    if ws is not None:
+                        await ws.send_bytes(chunk)
 
     async def sendWrappedError(self, rpcid, code, message, data=None, client_id=None):
         wrapper = {
@@ -495,15 +427,26 @@ class WslinkHandler(object):
         }
         if data:
             wrapper["error"]["data"] = data
-        encMsg = json.dumps(wrapper, ensure_ascii=False)
+
+        try:
+            packed_wrapper = msgpack.packb(wrapper)
+        except Exception:
+            del wrapper["error"]["data"]
+            packed_wrapper = msgpack.packb(wrapper)
+
         websockets = (
             [self.connections[client_id]]
             if client_id
             else [self.connections[c] for c in self.connections]
         )
-        for ws in websockets:
-            if ws is not None:
-                await ws.send_str(encMsg)
+        # aiohttp can not handle pending ws.send_bytes()
+        # tried with semaphore but got exception with >1
+        # https://github.com/aio-libs/aiohttp/issues/2934
+        async with self.attachment_atomic:
+            for chunk in generate_chunks(packed_wrapper, MAX_MSG_SIZE):
+                for ws in websockets:
+                    if ws is not None:
+                        await ws.send_bytes(chunk)
 
     def publish(self, topic, data, client_id=None, skip_last_active_client=False):
         client_list = [client_id] if client_id else [c_id for c_id in self.connections]

wslink/publish.py

@@ -7,9 +7,6 @@ from . import schedule_coroutine
 class PublishManager(object):
     def __init__(self):
         self.protocols = []
-        self.attachmentMap = {}
-        self.attachmentRefCounts = {}  # keyed same as attachment map
-        self.attachmentId = 0
         self.publishCount = 0
 
     def registerProtocol(self, protocol):
@@ -19,38 +16,12 @@ class PublishManager(object):
         if protocol in self.protocols:
             self.protocols.remove(protocol)
 
-    def getAttachmentMap(self):
-        return self.attachmentMap
-
-    def clearAttachmentMap(self):
-        self.attachmentMap.clear()
-
-    def registerAttachment(self, attachKey):
-        self.attachmentRefCounts[attachKey] += 1
-
-    def unregisterAttachment(self, attachKey):
-        self.attachmentRefCounts[attachKey] -= 1
-
-    def freeAttachments(self, keys=None):
-        keys_to_delete = []
-        keys_to_check = keys if keys is not None else [k for k in self.attachmentMap]
-
-        for key in keys_to_check:
-            if self.attachmentRefCounts.get(key) == 0:
-                keys_to_delete.append(key)
-
-        for key in keys_to_delete:
-            self.attachmentMap.pop(key)
-            self.attachmentRefCounts.pop(key)
-
     def addAttachment(self, payload):
-        # print("attachment", self, self.attachmentId)
-        # use a string flag in place of the binary attachment.
-        binaryId = "wslink_bin{0}".format(self.attachmentId)
-        self.attachmentMap[binaryId] = payload
-        self.attachmentRefCounts[binaryId] = 0
-        self.attachmentId += 1
-        return binaryId
+        """Deprecated method, keeping it to avoid breaking compatibility
+        Now that we use msgpack to pack/unpack messages,
+        We can have binary data directly in the object itself,
+        without needing to transfer it separately from the rest."""
+        return payload
 
     def publish(self, topic, data, client_id=None, skip_last_active_client=False):
         for protocol in self.protocols:

wslink/websocket.py

@@ -9,6 +9,8 @@ import logging
 from . import register as exportRpc
 from . import schedule_callback
 
+logger = logging.getLogger(__name__)
+
 
 # =============================================================================
 #
@@ -114,7 +116,7 @@ class ServerProtocol(object):
         try:
             self.linkProtocols.remove(protocol)
         except ValueError as e:
-            logging.error("Link protocol missing from registered list.")
+            logger.error("Link protocol missing from registered list.")
 
     def getLinkProtocols(self):
         return self.linkProtocols

{wslink-1.12.3.dist-info → wslink-2.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wslink
-Version: 1.12.3
+Version: 2.0.0
 Summary: Python/JavaScript library for communicating over WebSocket
 Home-page: https://github.com/kitware/wslink
 Author: Kitware, Inc.
@@ -22,6 +22,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Dist: aiohttp <4
+Requires-Dist: msgpack <2,>=1
 Provides-Extra: ssl
 Requires-Dist: cryptography ; extra == 'ssl'
 

wslink-2.0.0.dist-info/RECORD

@@ -0,0 +1,25 @@
+wslink/LICENSE,sha256=I44UH7kDVqxDLnnlOWw_hFL2Fz7RjQ_4vPzZv9NYgTU,1483
+wslink/__init__.py,sha256=AbEm-sUSoGL-uLpnbK1rSSjHSvyW-bMsGHWie7FgMHw,2708
+wslink/chunking.py,sha256=BZZ0YAlh6PNI8rQe80NfdxU8pAvn_Klxew47AkvUJow,7130
+wslink/launcher.py,sha256=8VMs3juObLkyGYQFNLjMoo4qFpKIcxWz0kS-af-DKO4,21170
+wslink/protocol.py,sha256=zdf4QthFHpAgEw3hTUyyaOuN76jzHeOJBpvekPbk7aY,15886
+wslink/publish.py,sha256=9G5TXqyGr-LCo_LwHYhzif6lhG2iXDvEBmEgwR8fh1M,1437
+wslink/relay.py,sha256=E8Lzu2Ay7KbOheN1-ArAZawo8lLqdDgJXOZSBuMknYs,86
+wslink/server.py,sha256=FKSJAKHDyfkNVM45-M-y1Zn8hh2TTYto1hTCIJx1pp8,9440
+wslink/ssl_context.py,sha256=hNOJJCdrStws1Qf6vPvY4vTk9Bf8J5d90W3fS0cRv8o,2290
+wslink/uri.py,sha256=woCQ4yChUqTMg9IT6YYDtUYeKmCg7OUCEgeBGA-19DY,384
+wslink/websocket.py,sha256=pBiWqkL8Zn8LuSJ9nv3yA-KjEynbolOQ2gLHtQFJ2Ic,4611
+wslink/backends/__init__.py,sha256=cyJGjm-YyBSyOEX81owyTbJ3YnrA6dB7--B4LnsEtHI,1214
+wslink/backends/aiohttp/__init__.py,sha256=u2UxSnaMJPoiba1CL1fOeM8bTggyOw6znQ2iGd69qJI,8463
+wslink/backends/aiohttp/launcher.py,sha256=gHNMvtgNHEwBN_QBRDSCrTp2B4K1PsfV81rKaHi7Cxo,8897
+wslink/backends/aiohttp/relay.py,sha256=oZAzIQTpsQaObWXaa-_VtoTOUQALC_QLDd9UvWspYaU,13311
+wslink/backends/generic/__init__.py,sha256=Qu65gWsd2xCSsxybnDtEDI5vMjHN-F5jgPZOyNIxnGs,112
+wslink/backends/generic/core.py,sha256=Gwo7nksgph96i7J77nYm4hB5NwysybQ23rRs86f0d6U,4149
+wslink/backends/jupyter/__init__.py,sha256=Qu65gWsd2xCSsxybnDtEDI5vMjHN-F5jgPZOyNIxnGs,112
+wslink/backends/jupyter/core.py,sha256=H73IEEHyom3TsbhkyI5O88bFBbUIDzHVuvqbIF6PAIM,3858
+wslink/backends/tornado/__init__.py,sha256=Qu65gWsd2xCSsxybnDtEDI5vMjHN-F5jgPZOyNIxnGs,112
+wslink/backends/tornado/core.py,sha256=tPMkkhWuO_ovkisVim0zcegwZKEAG4IRUdd_O_0a_R0,2157
+wslink-2.0.0.dist-info/METADATA,sha256=XWnXD2kxZEtIuOy4Y-_8HZ8anq9hYFL7axDQhv96GoE,3045
+wslink-2.0.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+wslink-2.0.0.dist-info/top_level.txt,sha256=N0d8eqvhwhfW1p1yPTmvxlbzhjz7ZyhBfysNvaFqpQY,7
+wslink-2.0.0.dist-info/RECORD,,

{wslink-1.12.3.dist-info → wslink-2.0.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.2)
+Generator: bdist_wheel (0.43.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

wslink-1.12.3.dist-info/RECORD

@@ -1,24 +0,0 @@
-wslink/LICENSE,sha256=I44UH7kDVqxDLnnlOWw_hFL2Fz7RjQ_4vPzZv9NYgTU,1483
-wslink/__init__.py,sha256=AbEm-sUSoGL-uLpnbK1rSSjHSvyW-bMsGHWie7FgMHw,2708
-wslink/launcher.py,sha256=4l_z3l2xkRSiH8hWbZzjtsIbBeujl-P8bTA0TBtJEEg,21137
-wslink/protocol.py,sha256=AXBlQ4BD2iB7e3WHPedIv0fzgd05eBuNLQpXjX74j5U,18132
-wslink/publish.py,sha256=6xLr7tGkvbq5LbXbAvPQ6OplEBbs192WR-uhiTJMPdg,2354
-wslink/relay.py,sha256=E8Lzu2Ay7KbOheN1-ArAZawo8lLqdDgJXOZSBuMknYs,86
-wslink/server.py,sha256=FKSJAKHDyfkNVM45-M-y1Zn8hh2TTYto1hTCIJx1pp8,9440
-wslink/ssl_context.py,sha256=hNOJJCdrStws1Qf6vPvY4vTk9Bf8J5d90W3fS0cRv8o,2290
-wslink/uri.py,sha256=woCQ4yChUqTMg9IT6YYDtUYeKmCg7OUCEgeBGA-19DY,384
-wslink/websocket.py,sha256=REYYyKB4WweMuDruvPBA0oPOfhddPF9G5WIGiK367KA,4574
-wslink/backends/__init__.py,sha256=cyJGjm-YyBSyOEX81owyTbJ3YnrA6dB7--B4LnsEtHI,1214
-wslink/backends/aiohttp/__init__.py,sha256=FUPMwi0Z1nuVVIJhY_sI0ZzRY5ne30tmUILi7JQTgP4,8433
-wslink/backends/aiohttp/launcher.py,sha256=Itblr2-lX_ZpAInvLlGkl1_XrNYjIn3w469BDZ5XyyQ,8737
-wslink/backends/aiohttp/relay.py,sha256=Nn1ELqGU1kwvst4DNwNqYWU4nYd08K3STAMQIhcangY,13338
-wslink/backends/generic/__init__.py,sha256=Qu65gWsd2xCSsxybnDtEDI5vMjHN-F5jgPZOyNIxnGs,112
-wslink/backends/generic/core.py,sha256=nphpXjwZ_U0-OyNDgLGg0IMpi6lfYb8tke_YDti4U4Y,4138
-wslink/backends/jupyter/__init__.py,sha256=Qu65gWsd2xCSsxybnDtEDI5vMjHN-F5jgPZOyNIxnGs,112
-wslink/backends/jupyter/core.py,sha256=H73IEEHyom3TsbhkyI5O88bFBbUIDzHVuvqbIF6PAIM,3858
-wslink/backends/tornado/__init__.py,sha256=Qu65gWsd2xCSsxybnDtEDI5vMjHN-F5jgPZOyNIxnGs,112
-wslink/backends/tornado/core.py,sha256=tPMkkhWuO_ovkisVim0zcegwZKEAG4IRUdd_O_0a_R0,2157
-wslink-1.12.3.dist-info/METADATA,sha256=DeacEUE7uOSFLAJHAf5tHdgPh2CkWDwxZKBl-xSoLpk,3016
-wslink-1.12.3.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-wslink-1.12.3.dist-info/top_level.txt,sha256=N0d8eqvhwhfW1p1yPTmvxlbzhjz7ZyhBfysNvaFqpQY,7
-wslink-1.12.3.dist-info/RECORD,,