vortexpy 4.0.23__py3-none-any.whl

vortex/VortexServerHttpResource.py

@@ -0,0 +1,193 @@
+"""
+ * Created by Synerty Pty Ltd
+ *
+ * This software is open source, the MIT license applies.
+ *
+ * Website : http://www.synerty.com
+ * Support : support@synerty.com
+"""
+
+import collections
+import logging
+from _pyio import BytesIO
+from typing import Deque
+
+from twisted.internet import task
+from twisted.internet.defer import inlineCallbacks
+from twisted.web.resource import Resource
+from twisted.web.server import NOT_DONE_YET
+
+from vortex.PayloadEnvelope import PayloadEnvelope
+from vortex.VortexServer import HEART_BEAT_PERIOD
+from .PayloadPriority import DEFAULT_PRIORITY
+from .VortexConnectionABC import VortexConnectionABC
+from .VortexServer import VortexServer
+
+logger = logging.getLogger(name=__name__)
+
+
+class VortexServerHttpResource(Resource):
+    """VortexServer Resource
+    This resource is the server endpoint for the vortex.
+    """
+
+    isLeaf = True
+
+    def __init__(self, vortex: VortexServer) -> None:
+        Resource.__init__(self)
+        self.__vortex = vortex
+
+    def render_GET(self, request):
+        request.code = 405
+        return b"You have reached the vortex. It only likes POST methods."
+
+    def render_POST(self, request):
+
+        from vortex.VortexFactory import VortexFactory
+
+        if not VortexFactory.canConnect(str(request.client.host)):
+            return
+
+        remoteVortexUuid = request.args[b"vortexUuid"][0].decode()
+        remoteVortexName = request.args[b"vortexName"][0].decode()
+
+        if self.__vortex.isShutdown():
+            return None
+
+        httpSession = request.getSession()
+        conn = VortexResourceConnection(
+            self.__vortex, remoteVortexUuid, remoteVortexName, request
+        )
+
+        # Send a heart beat down the new connection, tell it who we are.
+        connectPayloadFilt = {}
+        connectPayloadFilt[PayloadEnvelope.vortexUuidKey] = self.__vortex.uuid()
+        connectPayloadFilt[PayloadEnvelope.vortexNameKey] = self.__vortex.name()
+        conn.write(PayloadEnvelope(filt=connectPayloadFilt).toVortexMsg())
+
+        data = request.content.read()
+        if len(data):
+            for vortexStr in data.strip(b".").split(b"."):
+                self._processVortexMsg(
+                    httpSession, conn, vortexStr.decode("UTF-8")
+                )
+
+        # Request will be around for a while, do some cleanups
+        request.content = BytesIO()
+        request.args = {}
+
+        self.__vortex.connectionOpened(httpSession, conn)
+
+        def connClosed(err):
+            logger.debug("VortexServer connection ended by client")
+            self.__vortex.connectionClosed(conn)
+
+        request.notifyFinish().addErrback(connClosed)
+
+        return NOT_DONE_YET
+
+    @inlineCallbacks
+    def _processVortexMsg(self, httpSession, conn, vortexMsg):
+        payload = yield PayloadEnvelope().fromVortexMsgDefer(vortexMsg)
+        self.__vortex.payloadReveived(
+            httpSession=httpSession,
+            vortexUuid=conn.remoteVortexUuid,
+            vortexName=conn.remoteVortexName,
+            payload=payload,
+        )
+
+
+class VortexResourceConnection(VortexConnectionABC):
+    def __init__(
+        self,
+        vortexServer: VortexServer,
+        remoteVortexUuid: str,
+        remoteVortexName: str,
+        request,
+    ) -> None:
+        VortexConnectionABC.__init__(
+            self,
+            logger,
+            vortexServer,
+            remoteVortexUuid=remoteVortexUuid,
+            remoteVortexName=remoteVortexName,
+            httpSessionUuid=request.getSession().uid,
+        )
+
+        self._request = request
+        self._request.responseHeaders.setRawHeaders(
+            b"content-type", [b"text/text"]
+        )
+
+        self._buffer: Deque[bytes] = collections.deque()
+
+        self._requestReadyCheckTimer = task.LoopingCall(
+            self._checkIfRequestIsReady
+        )
+        self._requestReadyCheckTimer.start(0.001)
+
+        # Start our heart beat
+        self._beatLoopingCall = task.LoopingCall(self._beat)
+        d = self._beatLoopingCall.start(HEART_BEAT_PERIOD)
+        d.addErrback(lambda f: logger.exception(f.value))
+
+    def _beat(self):
+        if (
+            self._closed
+            or self._request.finished
+            or self._request.channel is None
+        ):
+            self._beatLoopingCall.stop()
+            return
+
+        # Send the heartbeats
+        self._request.write(b".")
+
+        # Touch the session
+        self._request.getSession().touch()
+
+    @property
+    def ip(self):
+        return self._request.client.host
+
+    @property
+    def port(self):
+        return self._request.client.port
+
+    def _checkIfRequestIsReady(self):
+        # This seems to be the fingerprint of the fact that request will not write
+        # our data yet.
+        if self._request.chunked and not self._request.startedWriting:
+            return
+
+        if self._requestReadyCheckTimer.running:
+            self._requestReadyCheckTimer.stop()
+
+        while self._buffer:
+            self._write(self._buffer.popleft())
+        self._buffer = None
+
+    def write(self, payloadVortexStr: bytes, priority: int = DEFAULT_PRIORITY):
+        assert not self._closed
+        # If we have a buffer, write it.
+        if self._buffer:
+            self._buffer.append(payloadVortexStr)
+        else:
+            self._write(payloadVortexStr)
+
+    def _write(self, payloadVortexStr: bytes):
+        if self._request.finished or self._request.channel is None:
+            logger.error("VORTEX CLOSED, CAN'T WRITE PAYLOAD")
+            self._closed = True
+            self._vortexServer.connectionClosed(self)
+            return
+
+        self._request.write(payloadVortexStr)
+        self._request.write(b".")
+
+    def close(self):
+        if self._beatLoopingCall.running:
+            self._beatLoopingCall.stop()
+
+        self._closed = True
+        self._request.finish()

vortex/VortexServerTcp.py

@@ -0,0 +1,104 @@
+"""
+ * Created by Synerty Pty Ltd
+ *
+ * This software is open source, the MIT license applies.
+ *
+ * Website : http://www.synerty.com
+ * Support : support@synerty.com
+"""
+import logging
+from typing import Union
+
+from twisted.internet.defer import succeed
+from twisted.internet.protocol import connectionDone, Factory
+
+from vortex.VortexPayloadProtocol import VortexPayloadProtocol
+from vortex.VortexServerConnection import VortexServerConnection
+from .PayloadEnvelope import VortexMsgList, PayloadEnvelope
+from .PayloadPriority import DEFAULT_PRIORITY
+from .VortexServer import VortexServer
+
+logger = logging.getLogger(name=__name__)
+
+
+class VortexTcpServerProtocol(VortexPayloadProtocol):
+    def __init__(self, vortex: VortexServer, addr) -> None:
+        VortexPayloadProtocol.__init__(self, logger)
+        self._vortex = vortex
+        self._addr = addr
+
+        self._conn = None
+
+        self._dataBuffer = b""
+        self._remoteVortexUuid = None
+        self._remoteVortexName = None
+        self._httpSession = None
+
+    def _beat(self):
+        if self._conn:
+            self._conn.beatReceived()
+
+    def _nameAndUuidReceived(self, name, uuid):
+        self._remoteVortexUuid = uuid
+        self._remoteVortexName = name
+        self._conn = VortexServerConnection(
+            self._vortex,
+            self._remoteVortexUuid,
+            self._remoteVortexName,
+            self._httpSession,
+            self.transport,
+            self._addr,
+        )
+
+        # Send a heart beat down the new connection, tell it who we are.
+        connectPayloadFilt = {}
+        connectPayloadFilt[PayloadEnvelope.vortexUuidKey] = self._vortex.uuid()
+        connectPayloadFilt[PayloadEnvelope.vortexNameKey] = self._vortex.name()
+        self._conn.write(PayloadEnvelope(filt=connectPayloadFilt).toVortexMsg())
+
+        self._vortex.connectionOpened(self._httpSession, self._conn)
+
+    def _createResponseSenderCallable(self):
+        def sendResponse(
+            vortexMsgs: Union[VortexMsgList, bytes],
+            priority: int = DEFAULT_PRIORITY,
+        ):
+            if isinstance(vortexMsgs, bytes):
+                vortexMsgs = [vortexMsgs]
+
+            for vortexMsg in vortexMsgs:
+                self._conn.write(vortexMsg, priority)
+
+            return succeed(True)
+
+        return sendResponse
+
+    def dataReceived(self, data):
+
+        if self._vortex.isShutdown():
+            return None
+
+        VortexPayloadProtocol.dataReceived(self, data)
+
+    def connectionLost(self, reason=connectionDone):
+        if self._conn:
+            self._conn.transportClosed()
+
+        VortexPayloadProtocol.connectionLost(self, reason)
+
+
+class VortexTcpServerFactory(Factory):
+    protocol = None
+
+    def __init__(self, vortexServer: VortexServer) -> None:
+        self._vortexServer = vortexServer
+
+    def buildProtocol(self, addr):
+        from vortex.VortexFactory import VortexFactory
+
+        if not VortexFactory.canConnect(str(addr)):
+            return
+
+        p = VortexTcpServerProtocol(self._vortexServer, addr)
+        p.factory = self
+        return p

vortex/VortexServerWebsocket.py

@@ -0,0 +1,287 @@
+"""
+ * Created by Synerty Pty Ltd
+ *
+ * This software is open source, the MIT license applies.
+ *
+ * Website : http://www.synerty.com
+ * Support : support@synerty.com
+"""
+
+import logging
+from collections import deque
+from random import random
+from urllib.parse import parse_qs
+from urllib.parse import urlparse
+
+from twisted.internet import reactor
+from twisted.internet.defer import inlineCallbacks
+from twisted.internet.protocol import Factory
+from twisted.internet.protocol import Protocol
+from twisted.internet.protocol import connectionDone
+from twisted.protocols.policies import WrappingFactory
+from twisted.web import resource
+from twisted.web.server import NOT_DONE_YET
+
+from txwebsocket.txws import WebSocketProtocol
+from vortex.DeferUtil import vortexLogFailure, nonConcurrentMethod
+from vortex.PayloadEnvelope import PayloadEnvelope
+from vortex.VortexServerConnection import VortexServerConnection
+from .PayloadPriority import DEFAULT_PRIORITY
+from .VortexServer import VortexServer
+from .VortexWritePushProducer import VortexWritePushProducer
+
+logger = logging.getLogger(name=__name__)
+
+
+class VortexWebsocketServerProtocol(Protocol):
+    _DEBUG_LOGGING = False
+
+    def __init__(self, vortex: VortexServer, addr, httpSession=None) -> None:
+        self._vortex = vortex
+        self._addr = addr
+        self._logger = logger
+
+        self._conn = None
+
+        self._remoteVortexUuid = None
+        self._remoteVortexName = None
+        self._httpSession = httpSession
+
+        def httpSessionExpiredCallback():
+            self._httpSession = None
+            self.transport.loseConnection()
+
+        if self._httpSession:
+            self._httpSession.notifyOnExpire(httpSessionExpiredCallback)
+
+        # Most messages don't need a buffer, but websockets can split messages
+        # around the 128kb mark
+        self._data = b""
+
+        self._vortexMsgsQueue = deque()
+
+    def __initConnection(self):
+        # self.transport.setBinaryMode(True)
+
+        params = parse_qs(urlparse(self.transport.location).query)
+
+        if "vortexUuid" not in params or "vortexName" not in params:
+            raise Exception(
+                "This isn't a vortex capable websocket. Check the URL"
+            )
+
+        self._remoteVortexUuid = params["vortexUuid"][0]
+        self._remoteVortexName = params["vortexName"][0]
+        self._conn = VortexServerConnection(
+            self._vortex,
+            self._remoteVortexUuid,
+            self._remoteVortexName,
+            self._httpSession,
+            self.transport,
+            self._addr,
+        )
+
+        # noinspection PyProtectedMember
+        assert isinstance(self.transport, WebSocketProtocol) and isinstance(
+            self.transport._producerForWrite, VortexWritePushProducer
+        ), "The Websocket transport is not how we expected it."
+
+        self.transport._producerForWrite.setRemoteVortexName(
+            f"{self._remoteVortexName}({self._remoteVortexUuid})"
+        )
+
+        # Send a heart beat down the new connection, tell it who we are.
+        connectPayloadFilt = {}
+        connectPayloadFilt[PayloadEnvelope.vortexUuidKey] = self._vortex.uuid()
+        connectPayloadFilt[PayloadEnvelope.vortexNameKey] = self._vortex.name()
+        self._conn.write(
+            PayloadEnvelope(filt=connectPayloadFilt).toVortexMsg(),
+            DEFAULT_PRIORITY,
+        )
+
+        self._vortex.connectionOpened(self._httpSession, self._conn)
+
+    def dataReceived(self, data):
+        if not self._conn:
+            self.__initConnection()
+
+        if self._vortex.isShutdown():
+            return None
+
+        if self._httpSession:
+            self._httpSession.touch()
+
+        self._conn.beatReceived()
+
+        if data.startswith(b"{"):
+            # If the data starts with a { then the payloadEnvelope is al in
+            # the websocket message.
+            assert data.endswith(b"}"), (
+                "%s payloadEnvelope doesn't fit in websocket message" % __name__
+            )
+            self._vortexMsgsQueue.append(data)
+        else:
+            # Else it's a base64 encoded payloadEnvelope, and it will be separated
+            # by b'.'
+            self._data += data
+
+        self._processData()
+
+    def connectionLost(self, reason=connectionDone):
+        if self._DEBUG_LOGGING:
+            logger.debug(
+                "VortexWebsocketServerProtocol.connectionLost reason=%s", reason
+            )
+        self._processData()
+        if self._conn:
+            self._conn.transportClosed()
+            self._vortex.connectionClosed(self._conn)
+
+    def _processData(self):
+        def getNextChunkIter():
+            try:
+                while True:
+                    yield self._data.index(b".")
+
+            except ValueError:
+                # There is no '.' in it, wait for more data.
+                return
+
+        if self._data:
+            for nextChunk in getNextChunkIter():
+                vortexMsg = self._data[:nextChunk]
+                self._data = self._data[nextChunk + 1 :]
+
+                # If we get two heartbeats in a row, this will be false
+                if len(vortexMsg):
+                    self._vortexMsgsQueue.append(vortexMsg)
+
+        if self._vortexMsgsQueue and not self._processVortexMsgs.running:
+            reactor.callLater(0, self._processVortexMsgs)
+
+    @inlineCallbacks
+    @nonConcurrentMethod
+    def _processVortexMsgs(self):
+        while self._vortexMsgsQueue:
+            vortexMsg = self._vortexMsgsQueue.popleft()
+
+            try:
+                payloadEnvelope = yield PayloadEnvelope().fromVortexMsgDefer(
+                    vortexMsg
+                )
+
+                self._deliverPayload(payloadEnvelope)
+
+            except Exception as e:
+                print(vortexMsg)
+                print(e)
+                self._logger.exception(e)
+                raise
+
+    def _deliverPayload(self, payloadEnvelope):
+        self._vortex.payloadReveived(
+            httpSession=self._httpSession,
+            vortexUuid=self._remoteVortexUuid,
+            vortexName=self._remoteVortexName,
+            payload=payloadEnvelope,
+        )
+
+
+class VortexWebsocketServerFactory(Factory):
+    protocol = None
+
+    def __init__(self, vortexServer: VortexServer) -> None:
+        self._vortexServer = vortexServer
+
+    def buildProtocol(self, addr, httpSession=None):
+        # VortexWebSocketUpgradeResource provides this check
+        # if we're connected via a HTTP upgrade
+        if not httpSession:
+            from vortex.VortexFactory import VortexFactory
+
+            if not VortexFactory.canConnect(str(addr)):
+                return
+
+        p = VortexWebsocketServerProtocol(self._vortexServer, addr, httpSession)
+        p.factory = self
+        return p
+
+
+class VortexWrappedWebSocketFactory(WrappingFactory):
+    """
+    Factory which wraps another factory to provide WebSockets transports for
+    all of its protocols.
+    """
+
+    protocol = WebSocketProtocol
+
+    def buildProtocol(self, addr, httpSession):
+        return self.protocol(
+            self, self.wrappedFactory.buildProtocol(addr, httpSession)
+        )
+
+
+class VortexWebSocketUpgradeResource(resource.Resource):
+    """Vortex Websocket Upgrade Resource
+
+    If this resource is hit, it will attempt to upgrade the connection to a websocket.
+
+    """
+
+    isLeaf = 1
+
+    def __init__(self, websocketFactory):
+        """Constructor
+
+        @:param websocketFactory: A factory that will build a WebsocketProtocol (above)
+        """
+        resource.Resource.__init__(self)
+        self._websocketFactory = websocketFactory
+
+    def render(self, request):
+        from vortex.VortexFactory import VortexFactory
+
+        if not VortexFactory.canConnect(str(request.client.host)):
+            request.setResponseCode(503)
+            request.setHeader("Retry-After", str(max(10, int(random() * 100))))
+            return b"Rejected due to new connections throttling."
+
+        httpSession = request.getSession()
+
+        websocketProtocol = self._websocketFactory.buildProtocol(
+            request.client, httpSession
+        )
+        websocketProtocol.makeConnection(request.channel.transport)
+        websocketProtocol.initFromRequest(request)
+        request.channel.upgradeToWebsocket(websocketProtocol)
+
+        def unregister():
+            request.channel.unregisterProducer()
+            websocketProtocol.unregisterProducer()
+
+        producer = VortexWritePushProducer(
+            request.channel.transport,
+            stopProducingCallback=lambda: unregister(),
+            # The websocket frames this before this producer gets it
+            writeWholeFrames=False,
+            terminateFrameWithDot=False,
+            splitFrames=True,
+        )
+        producer.WARNING_DATA_LENGTH = 10 * 1024 * 1024
+        producer.ERROR_DATA_LENGTH = 50 * 1024 * 1024
+
+        unregister()
+        request.channel.registerProducer(producer, streaming=True)
+        websocketProtocol.registerProducer(producer, streaming=True)
+
+        def closedError(failure):
+            # logger.error("Got closedError %s" % failure)
+            websocketProtocol.connectionLost(failure.value)
+
+        def closedOk(data):
+            pass
+            # logger.debug("Got closedOk %s" % data)
+
+        request.notifyFinish().addCallbacks(closedOk, closedError)
+
+        return NOT_DONE_YET