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)
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)
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}
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)
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)
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()
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)
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)
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)
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)
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))
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)
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)