def handle(self):
# first send the info message
self.authrand = authrand = os.urandom(4)
self.write(msginfo(config.FBNAME, authrand))
while True:
self.unpacker.feed(self.s.recv(BUFSIZ))
for opcode, data in self.unpacker:
if opcode != OP_AUTH:
self.error('First message was not AUTH.')
raise BadClient()
ident, rhash = readauth(data)
self.authkey_check(ident, rhash)
break
while True:
self.unpacker.feed(self.recv(BUFSIZ))
for opcode, data in self.unpacker:
if opcode == OP_PUBLISH:
self.do_publish(*readpublish(data))
elif opcode == OP_SUBSCRIBE:
self.do_subscribe(*readsubscribe(data))
elif opcode == OP_UNSUBSCRIBE:
self.do_unsubscribe(*readsubscribe(data))
else:
self.error(
"Unknown message type.",
opcode=opcode,
length=len(data),
)
raise BadClient()
def test_auth_success(self):
c = self.make_connection()
name, rand = readinfo(parse(c.transport.write)[0][1])
c.data_received(msgauth(rand, 'test', 'secret'))
c.data_received(msgsubscribe('test', 'test-chan'))
c.data_received(msgpublish('test', 'test-chan', b'c'))
assert readpublish(parse(c.transport.write)[1][1]) == ('test',
'test-chan',
b'c')
def messageReceived(self, opcode, data):
if opcode == OP_ERROR:
return self.onError(readerror(data))
elif opcode == OP_INFO:
return self.onInfo(*readinfo(data))
elif opcode == OP_AUTH:
return self.onAuth(*readauth(data))
elif opcode == OP_PUBLISH:
return self.onPublish(*readpublish(data))
elif opcode == OP_SUBSCRIBE:
return self.onSubscribe(*readsubscribe(data))
elif opcode == OP_UNSUBSCRIBE:
return self.onUnsubscribe(*readunsubscribe(data))
# Can't recover from an unknown opcode, so drop connection
self.protocolError('Unknown message opcode: {!r}'.format(opcode))
self.transport.loseConnection()
def test_multiple_subscribers(self):
subscribers = []
for i in range(5):
c = self.make_connection()
name, rand = readinfo(parse(c.transport.write)[0][1])
c.data_received(msgauth(rand, 'test', 'secret'))
c.data_received(msgsubscribe('test', 'test-chan'))
subscribers.append(c)
c = self.make_connection()
name, rand = readinfo(parse(c.transport.write)[0][1])
c.data_received(msgauth(rand, 'test', 'secret'))
c.data_received(msgpublish('test', 'test-chan', b'c'))
for c in subscribers:
msgs = parse(c.transport.write)
assert readpublish(msgs[1][1]) == ('test', 'test-chan', b'c')
def test_subscribe_and_publish(self):
c = client.new('127.0.0.1', self.server.port, 'test', 'secret')
c.subscribe('test-chan')
c._subscribe()
c.publish('test-chan', b'data')
opcode, data = c._read_message()
assert opcode == 3
assert readpublish(data) == ('test', 'test-chan', b'data')
self.log.debug('Stopping client')
c.stop()
self.log.debug('Closing client')
c.close()
def test_auth_unsubscribe(self):
c = self.make_connection()
name, rand = readinfo(parse(c.transport.write)[0][1])
c.data_received(msgauth(rand, 'test', 'secret'))
c.data_received(msgsubscribe('test', 'test-chan'))
c.data_received(msgpublish('test', 'test-chan', b'c'))
c.data_received(msgunsubscribe('test', 'test-chan'))
c.data_received(msgpublish('test', 'test-chan', b'c'))
c.data_received(msgsubscribe('test', 'test-chan'))
c.data_received(msgpublish('test', 'test-chan', b'c'))
messages = parse(c.transport.write)
for msg in messages[1:]:
assert readpublish(msg[1]) == ('test', 'test-chan', b'c')
# 1 auth and 2 publish
assert len(messages) == 3
def test_subscribe_and_publish(self):
c = client.new('127.0.0.1', self.port, 'test', 'secret')
c.subscribe('test-chan')
c._subscribe()
# If we have subscribed to a channel we should be able to see a
# connection in monitoring
assert prometheus.REGISTRY.get_sample_value(
'hpfeeds_broker_client_connections') == 1
c.publish('test-chan', b'data')
opcode, data = c._read_message()
assert opcode == 3
assert readpublish(data) == ('test', 'test-chan', b'data')
# We managed to publish a message - check this is reflected in stats
assert 1 == prometheus.REGISTRY.get_sample_value(
'hpfeeds_broker_receive_publish_count', {
'ident': 'test',
'chan': 'test-chan'
})
# If we managed to read a message from the broker then we must be subscribed
# Check this is reflected in stats
assert 1 == prometheus.REGISTRY.get_sample_value(
'hpfeeds_broker_subscriptions', {
'ident': 'test',
'chan': 'test-chan'
})
self.log.debug('Stopping client')
c.stop()
self.log.debug('Closing client')
c.close()
def test_readpublish(self):
ident, chan, data = readpublish(b'\x05ident\x04chansomedata')
assert ident == 'ident'
assert chan == 'chan'
assert data == b'somedata'