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'