示例#1
0
async def handleHandshake(state: ConnectionState, ws, app: Dict, pdu: JsonDict,
                          serializedPdu: str):
    authMethod = pdu.get('body', {}).get('method')
    if authMethod != 'role_secret':
        errMsg = f'invalid auth method: {authMethod}'
        logging.warning(errMsg)
        response = {
            "action": f"auth/handshake/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        await state.respond(ws, response)
        return

    role = pdu.get('body', {}).get('data', {}).get('role')
    state.role = role
    state.nonce = generateNonce()

    response = {
        "action": "auth/handshake/ok",
        "id": pdu.get('id', 1),
        "body": {
            "data": {
                "nonce": state.nonce,
                "version": getVersion(),
                "connection_id": state.connection_id,
                "node": platform.uname().node,
            }
        },
    }
    await state.respond(ws, response)
示例#2
0
async def handleRead(
    state: ConnectionState, ws, app: Dict, pdu: JsonDict, serializedPdu: str
):

    body = pdu.get('body', {})
    position = body.get('position')
    channel = body.get('channel')

    appChannel = '{}::{}'.format(state.appkey, channel)

    redis = app['redis_clients'].makeRedisClient()

    try:
        # Handle read
        message = await kvStoreRead(redis, appChannel, position, state.log)
    except Exception as e:
        errMsg = f'read: cannot connect to redis {e}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/read/error",
            "id": pdu.get('id', 1),
            "body": {"error": errMsg},
        }
        await state.respond(ws, response)
        return

    app['stats'].updateReads(state.role, len(serializedPdu))

    # Correct path
    response = {
        "action": "rtm/read/ok",
        "id": pdu.get('id', 1),
        "body": {"message": message},
    }
    await state.respond(ws, response)
示例#3
0
async def toggleFileLogging(state: ConnectionState, app: Dict,
                            params: JsonDict):
    found = False
    for connectionId, (st, websocket) in app['connections'].items():
        if connectionId == params.get('connection_id', ''):
            st.fileLogging = not st.fileLogging
            found = True

    return {'found': found, 'params': params}
示例#4
0
async def handleProducerMessage(
    state: ConnectionState, ws, app: Dict, pdu: JsonDict, serializedPdu: str, path: str
):
    tokens = path.split('/')
    if len(tokens) != 8:
        await badFormat(state, ws, app, f'Invalid uri -> {path}')
        return

    tenant = tokens[5]
    namespace = tokens[6]
    # FIXME: topic can contain multiple /
    topic = tokens[7]

    chan = f'{tenant}::{namespace}::{topic}'
    payload = pdu.get('payload')  # FIXME error if missing
    context = pdu.get('context')  # FIXME error if missing ??

    args = ['payload', payload, 'context', context]

    appkey = state.appkey
    redis = app['redis_clients'].getRedisClient(appkey)

    try:
        maxLen = app['channel_max_length']
        stream = '{}::{}'.format(appkey, chan)
        streamId = await redis.xaddRaw(stream, maxLen, *args)
    except Exception as e:
        # await publishers.erasePublisher(appkey, chan)  # FIXME

        errMsg = f'publish: cannot connect to redis {e}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/publish/error",
            "id": pdu.get('id', 1),
            "body": {"error": errMsg},
        }
        await state.respond(ws, response)
        return

    app['stats'].updateChannelPublished(chan, len(serializedPdu))

    response = {"result": "ok", "messageId": streamId.decode(), "context": context}
    await state.respond(ws, response)
示例#5
0
async def handleAdminCloseConnection(state: ConnectionState, ws, app: Dict,
                                     pdu: JsonDict, serializedPdu: str):
    action = pdu['action']
    body = pdu.get('body', {})
    targetConnectionId = body.get('connection_id')

    if targetConnectionId is None:
        errMsg = f'Missing connection id'
        logging.warning(errMsg)
        response = {
            "action": f"{action}/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        await state.respond(ws, response)
        return

    found = False
    for connectionId, (_, websocket) in app['connections'].items():
        if connectionId == targetConnectionId:
            targetWebSocket = websocket
            found = True

    if not found:
        errMsg = f'Cannot find connection id'
        logging.warning(errMsg)
        response = {
            "action": f"{action}/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        await state.respond(ws, response)
        return

    await targetWebSocket.close()

    response = {"action": f"{action}/ok", "id": pdu.get('id', 1), "body": {}}
    await state.respond(ws, response)
示例#6
0
async def handleUnSubscribe(state: ConnectionState, ws, app: Dict,
                            pdu: JsonDict, serializedPdu: str):
    '''
    Cancel a subscription
    '''
    body = pdu.get('body', {})

    subscriptionId = body.get('subscription_id')
    if subscriptionId is None:
        errMsg = 'Body Missing subscriptionId'
        logging.warning(errMsg)
        response = {
            "action": "rtm/unsubscribe/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        await state.respond(ws, response)
        return

    key = subscriptionId + state.connection_id
    item = state.subscriptions.get(key, (None, None))
    task, _ = item
    if task is None:
        errMsg = f'Invalid subscriptionId: {subscriptionId}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/unsubscribe/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        await state.respond(ws, response)
        return

    # Correct path
    response = {"action": "rtm/unsubscribe/ok", "id": pdu.get('id', 1)}
    await state.respond(ws, response)

    task.cancel()
示例#7
0
async def handleAdminGetConnections(state: ConnectionState, ws, app: Dict,
                                    pdu: JsonDict, serializedPdu: str):
    action = pdu['action']
    connections = list(app['connections'].keys())

    response = {
        "action": f"{action}/ok",
        "id": pdu.get('id', 1),
        "body": {
            'connections': connections
        },
    }
    await state.respond(ws, response)
示例#8
0
async def handleAuth(state: ConnectionState, ws, app: Dict, pdu: JsonDict,
                     serializedPdu: str):
    try:
        secret = app['apps_config'].getRoleSecret(state.appkey, state.role)
    except KeyError:
        reason = 'invalid_role'
        success = False
    else:
        serverHash = computeHash(secret, state.nonce.encode('ascii'))
        clientHash = pdu.get('body', {}).get('credentials', {}).get('hash')

        state.log(f'server hash {serverHash}')
        state.log(f'client hash {clientHash}')

        success = clientHash == serverHash
        if not success:
            reason = 'challenge_failed'

    if success:
        response = {
            "action": "auth/authenticate/ok",
            "id": pdu.get('id', 1),
            "body": {},
        }

        state.authenticated = True
        state.permissions = app['apps_config'].getPermissions(
            state.appkey, state.role)
    else:
        logging.warning(f'auth error: {reason}')
        response = {
            "action": "auth/authenticate/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": "authentication_failed",
                "reason": reason
            },
        }
    await state.respond(ws, response)
示例#9
0
async def handleAdminCloseAllConnection(state: ConnectionState, ws, app: Dict,
                                        pdu: JsonDict, serializedPdu: str):
    action = pdu['action']

    websocketLists = []
    for connectionId, (st, websocket) in app['connections'].items():
        if connectionId != st.connection_id:
            websocketLists.append(websocket)

    for websocket in websocketLists:
        await websocket.close()  # should this be shielded ?

    response = {"action": f"{action}/ok", "id": pdu.get('id', 1), "body": {}}
    await state.respond(ws, response)
示例#10
0
async def handleDelete(
    state: ConnectionState, ws, app: Dict, pdu: JsonDict, serializedPdu: str
):
    # Missing channel
    channel = pdu.get('body', {}).get('channel')
    if channel is None:
        errMsg = 'delete: missing channel field'
        logging.warning(errMsg)
        response = {
            "action": "rtm/delete/error",
            "id": pdu.get('id', 1),
            "body": {"error": errMsg},
        }
        await state.respond(ws, response)
        return

    appChannel = '{}::{}'.format(state.appkey, channel)

    appkey = state.appkey
    redis = app['redis_clients'].getRedisClient(appkey)

    try:
        await redis.delete(appChannel)
    except Exception as e:
        errMsg = f'delete: cannot connect to redis {e}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/delete/error",
            "id": pdu.get('id', 1),
            "body": {"error": errMsg},
        }
        await state.respond(ws, response)
        return

    response = {"action": f"rtm/delete/ok", "id": pdu.get('id', 1), "body": {}}
    await state.respond(ws, response)
示例#11
0
async def handlePublish(state: ConnectionState, ws, app: Dict, pdu: JsonDict,
                        serializedPdu: str):
    '''Here we don't write back a result to the client for efficiency.
    Client doesn't really needs it.
    '''
    # Potentially add extra channels with channel builder rules
    rules = app['apps_config'].getChannelBuilderRules(state.appkey)
    pdu = updateMsg(rules, pdu)

    # Missing message
    message = pdu.get('body', {}).get('message')
    if message is None:
        errMsg = 'publish: empty message'
        logging.warning(errMsg)
        response = {
            "action": "rtm/publish/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        await state.respond(ws, response)
        return

    # Missing channels
    channel = pdu.get('body', {}).get('channel')
    channels = pdu.get('body', {}).get('channels')
    if channel is None and channels is None:
        errMsg = 'publish: no channel or channels field'
        logging.warning(errMsg)
        response = {
            "action": "rtm/publish/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        await state.respond(ws, response)
        return

    if channels is None:
        channels = [channel]

    streams = {}

    appkey = state.appkey
    redis = app['redis_clients'].getRedisClient(appkey)

    for chan in channels:

        # sanity check to skip empty channels
        if chan is None:
            continue

        try:
            maxLen = app['channel_max_length']
            stream = '{}::{}'.format(appkey, chan)
            streamId = await redis.xadd(stream, 'json', serializedPdu, maxLen)

            streams[chan] = streamId
        except Exception as e:
            # await publishers.erasePublisher(appkey, chan)  # FIXME

            errMsg = f'publish: cannot connect to redis {e}'
            logging.warning(errMsg)
            response = {
                "action": "rtm/publish/error",
                "id": pdu.get('id', 1),
                "body": {
                    "error": errMsg
                },
            }
            await state.respond(ws, response)
            return

        app['stats'].updateChannelPublished(chan, len(serializedPdu))

    response = {
        "action": "rtm/publish/ok",
        "id": pdu.get('id', 1),
        "body": {
            'channels': channels
        },
    }
    await state.respond(ws, response)

    # Stats
    app['stats'].updatePublished(state.role, len(serializedPdu))
示例#12
0
async def handleSubscribe(state: ConnectionState, ws, app: Dict, pdu: JsonDict,
                          serializedPdu: str):
    '''
    Client doesn't really needs it.
    '''
    body = pdu.get('body', {})
    channel = body.get('channel')

    subscriptionId = body.get('subscription_id')

    if channel is None and subscriptionId is None:
        errMsg = 'missing channel and subscription_id'
        logging.warning(errMsg)
        response = {
            "action": "rtm/subscribe/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        await state.respond(ws, response)
        return

    maxSubs = app['max_subscriptions']
    if maxSubs >= 0 and len(state.subscriptions) + 1 > maxSubs:
        errMsg = f'subscriptions count over max limit: {maxSubs}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/subscribe/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        state.ok = False
        state.error = response
        await state.respond(ws, response)
        return

    if channel is None:
        channel = subscriptionId

    if subscriptionId is None:
        subscriptionId = channel

    filterStr = body.get('filter')
    hasFilter = filterStr not in ('', None)

    try:
        streamSQLFilter = StreamSqlFilter(filterStr) if hasFilter else None
    except InvalidStreamSQLError:
        errMsg = f'Invalid SQL expression {filterStr}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/subscribe/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        state.error = response
        await state.respond(ws, response)
        return

    if hasFilter and streamSQLFilter is not None:
        channel = streamSQLFilter.channel

    position = body.get('position')
    if not validatePosition(position):
        errMsg = f'Invalid position: {position}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/subscribe/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        state.ok = False
        state.error = response
        await state.respond(ws, response)
        return

    batchSize = body.get('batch_size', 1)
    try:
        batchSize = int(batchSize)
    except ValueError:
        errMsg = f'Invalid batch size: {batchSize}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/subscribe/error",
            "id": pdu.get('id', 1),
            "body": {
                "error": errMsg
            },
        }
        state.ok = False
        state.error = response
        await state.respond(ws, response)
        return

    response = {
        "action": "rtm/subscribe/ok",
        "id": pdu.get('id', 1),
        "body": {
            # FIXME: we should set the position by querying
            # the redis stream, inside the MessageHandler
            "position": "1519190184-559034812775",
            "subscription_id": subscriptionId,
        },
    }

    class MessageHandlerClass(RedisSubscriberMessageHandlerClass):
        def __init__(self, args):
            self.cnt = 0
            self.cntPerSec = 0
            self.throttle = Throttle(seconds=1)
            self.ws = args['ws']
            self.subscriptionId = args['subscription_id']
            self.hasFilter = args['has_filter']
            self.streamSQLFilter = args['stream_sql_filter']
            self.appkey = args['appkey']
            self.serverStats = args['stats']
            self.state = args['state']
            self.subscribeResponse = args['subscribe_response']
            self.app = args['app']
            self.channel = args['channel']
            self.batchSize = args['batch_size']
            self.idIterator = itertools.count()

            self.messages = []

        def log(self, msg):
            self.state.log(msg)

        async def on_init(self, initInfo):
            response = self.subscribeResponse
            response['body'].update(initInfo)

            if not initInfo.get('success', False):
                msgId = response['id']
                response = {
                    'action': 'rtm/subscribe/error',
                    'id': msgId,
                    'body': {
                        'error':
                        'subscribe error: server cannot connect to redis'
                    },
                }

            # Send response.
            await self.state.respond(self.ws, response)

        async def handleMsg(self, msg: dict, position: str,
                            payloadSize: int) -> bool:

            # Input msg is the full serialized publish pdu.
            # Extract the real message out of it.
            msg = msg.get('body', {}).get('message')

            self.serverStats.updateSubscribed(self.state.role, payloadSize)
            self.serverStats.updateChannelSubscribed(self.channel, payloadSize)

            if self.hasFilter:
                filterOutput = self.streamSQLFilter.match(
                    msg.get('messages') or msg)  # noqa
                if not filterOutput:
                    return True
                else:
                    msg = filterOutput

            self.messages.append(msg)
            if len(self.messages) < self.batchSize:
                return True

            assert position is not None

            pdu = {
                "action": "rtm/subscription/data",
                "id": next(self.idIterator),
                "body": {
                    "subscription_id": self.subscriptionId,
                    "messages": self.messages,
                    "position": position,
                },
            }
            serializedPdu = json.dumps(pdu)
            self.state.log(f"> {serializedPdu} at position {position}")

            await self.ws.send(serializedPdu)

            self.cnt += len(self.messages)
            self.cntPerSec += len(self.messages)

            self.messages = []

            if self.throttle.exceedRate():
                return True

            self.state.log(f"#messages {self.cnt} msg/s {self.cntPerSec}")
            self.cntPerSec = 0

            return True

    appChannel = '{}::{}'.format(state.appkey, channel)

    # We need to create a new connection as reading from it will be blocking
    redisClient = app['redis_clients'].makeRedisClient()

    task = asyncio.ensure_future(
        redisSubscriber(
            redisClient.redis,
            appChannel,
            position,
            MessageHandlerClass,
            {
                'ws': ws,
                'subscription_id': subscriptionId,
                'has_filter': hasFilter,
                'stream_sql_filter': streamSQLFilter,
                'appkey': state.appkey,
                'stats': app['stats'],
                'state': state,
                'subscribe_response': response,
                'app': app,
                'channel': channel,
                'batch_size': batchSize,
            },
        ))
    addTaskCleanup(task)

    key = subscriptionId + state.connection_id
    state.subscriptions[key] = (task, state.role)

    app['stats'].incrSubscriptions(state.role)
示例#13
0
async def handleWrite(
    state: ConnectionState, ws, app: Dict, pdu: JsonDict, serializedPdu: str
):
    # Missing message
    message = pdu.get('body', {}).get('message')
    if message is None:
        errMsg = 'write: empty message'
        logging.warning(errMsg)
        response = {
            "action": "rtm/write/error",
            "id": pdu.get('id', 1),
            "body": {"error": errMsg},
        }
        await state.respond(ws, response)
        return

    # Missing channel
    channel = pdu.get('body', {}).get('channel')
    if channel is None:
        errMsg = 'write: missing channel field'
        logging.warning(errMsg)
        response = {
            "action": "rtm/write/error",
            "id": pdu.get('id', 1),
            "body": {"error": errMsg},
        }
        await state.respond(ws, response)
        return

    # Extract the message. This is what will be published
    message = pdu['body']['message']

    appkey = state.appkey
    redis = app['redis_clients'].getRedisClient(appkey)

    try:
        appChannel = '{}::{}'.format(state.appkey, channel)

        serializedPdu = json.dumps(message)
        streamId = await redis.xadd(appChannel, 'json', serializedPdu, maxLen=1)

    except Exception as e:
        errMsg = f'write: cannot connect to redis {e}'
        logging.warning(errMsg)
        response = {
            "action": "rtm/write/error",
            "id": pdu.get('id', 1),
            "body": {"error": errMsg},
        }
        await state.respond(ws, response)
        return

    # Stats
    app['stats'].updateWrites(state.role, len(serializedPdu))

    response = {
        "action": f"rtm/write/ok",
        "id": pdu.get('id', 1),
        "body": {"stream": streamId.decode()},
    }
    await state.respond(ws, response)