class TestBrokerPrometheusEndpoint(unittest.TestCase):
def setUp(self):
prometheus.reset()
authenticator = Authenticator({
'test': {
'secret': 'secret',
'subchans': ['test-chan'],
'pubchans': ['test-chan'],
'owner': 'some-owner',
}
})
self.server = Server(authenticator, exporter='127.0.0.1:20001')
self.server.add_endpoint_legacy('127.0.0.1:20000')
def test_metrics_server(self):
async def inner():
server_future = asyncio.ensure_future(self.server.serve_forever())
await self.server.when_started
async with aiohttp.ClientSession() as session:
async with session.get('http://127.0.0.1:20001/metrics') as resp:
metrics = await resp.text()
print(metrics)
assert 'hpfeeds_broker_client_connections 0.0' in metrics
assert 'hpfeeds_broker_connection_send_buffer_size{' not in metrics
sock = socket.socket()
sock.connect(('127.0.0.1', 20000))
async with aiohttp.ClientSession() as session:
async with session.get('http://127.0.0.1:20001/metrics') as resp:
metrics = await resp.text()
print(metrics)
assert 'hpfeeds_broker_client_connections 1.0' in metrics
assert 'hpfeeds_broker_connection_send_buffer_size{' not in metrics
sock.close()
async with ClientSession('127.0.0.1', 20000, 'test', 'secret'):
async with aiohttp.ClientSession() as session:
async with session.get('http://127.0.0.1:20001/metrics') as resp:
metrics = await resp.text()
print(metrics)
assert 'hpfeeds_broker_client_connections 1.0' in metrics
assert 'hpfeeds_broker_connection_send_buffer_size{ident="test"} 0.0' in metrics
server_future.cancel()
await server_future
asyncio.get_event_loop().run_until_complete(inner())
def main():
parser = argparse.ArgumentParser(description='Run a hpfeeds broker')
parser.add_argument('--bind', default=None, action='store')
parser.add_argument('--exporter', default='', action='store')
parser.add_argument('--name', default='hpfeeds', action='store')
parser.add_argument('--debug', default=False, action='store_true')
parser.add_argument('--auth', default=None, action='append')
parser.add_argument('--tlscert', default=None, action='store')
parser.add_argument('--tlskey', default=None, action='store')
parser.add_argument('-e', '--endpoint', default=None, action='append')
args = parser.parse_args()
if (args.tlscert and not args.tlskey) or (args.tlskey
and not args.tlscert):
parser.error('Must specify --tlskey AND --tlscert')
return
logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO, )
auth = multi.Authenticator()
auths = args.auth if args.auth else ['sqlite']
for a in auths:
try:
auth.add(get_authenticator(a))
except ServerException as e:
print(str(e))
sys.exit(1)
broker = Server(
auth=auth,
exporter=args.exporter,
name=args.name,
)
if args.bind or not args.endpoint:
bind = args.bind or '0.0.0.0:20000'
broker.add_endpoint_legacy(bind,
tlscert=args.tlscert,
tlskey=args.tlskey)
if args.endpoint:
for endpoint in args.endpoint:
broker.add_endpoint_str(endpoint)
return aiorun.run(broker.serve_forever())
class TestBrokerConnection(unittest.TestCase):
def setUp(self):
prometheus.reset()
authenticator = Authenticator({
'test': {
'secret': 'secret',
'subchans': ['test-chan'],
'pubchans': ['test-chan'],
'owner': 'some-owner',
}
})
self.server = Server(authenticator)
self.server.add_endpoint_legacy('127.0.0.1:20000')
def make_connection(self):
transport = mock.Mock()
transport.get_extra_info.side_effect = lambda name: (
'127.0.0.1', 80) if name == 'peername' else None
connection = Connection(self.server)
connection.connection_made(transport)
return connection
def test_sends_challenge(self):
c = self.make_connection()
assert parse(c.transport.write)[0][1][:-4] == b'\x07hpfeeds'
def test_must_auth(self):
c = self.make_connection()
c.data_received(msgpublish('a', 'b', b'c'))
assert parse(c.transport.write)[0][1][:-4] == b'\x07hpfeeds'
assert parse(c.transport.write)[1][1] == b'First message was not AUTH'
def test_auth_failure_wrong_secret(self):
c = self.make_connection()
name, rand = readinfo(parse(c.transport.write)[0][1])
c.data_received(msgauth(rand, 'test', 'secret2'))
assert parse(
c.transport.write)[1][1] == b'Authentication failed for test'
def test_auth_failure_no_such_ident(self):
c = self.make_connection()
name, rand = readinfo(parse(c.transport.write)[0][1])
c.data_received(msgauth(rand, 'test2', 'secret'))
assert parse(
c.transport.write)[1][1] == b'Authentication failed for test2'
def test_permission_to_sub(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-chan2'))
assert parse(c.transport.write)[1][
1] == b'Authkey not allowed to sub here. ident=test, chan=test-chan2'
def test_permission_to_pub(self):
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-chan2', b'c'))
assert parse(c.transport.write)[1][
1] == b'Authkey not allowed to pub here. ident=test, chan=test-chan2'
def test_pub_ident_checked(self):
c = self.make_connection()
name, rand = readinfo(parse(c.transport.write)[0][1])
c.data_received(msgauth(rand, 'test', 'secret'))
c.data_received(msgpublish('wrong-ident', 'test-chan2', b'c'))
assert parse(
c.transport.write
)[1][1] == b'Invalid authkey in message, ident=wrong-ident'
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 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_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
