def test_malicious_messages():
def echo(text):
return text
servant = Servant([echo])
servant.bind_tcp(HOST, PORT)
servant_thread = Thread(target=servant.run)
servant_thread.start()
ctx = zmq.Context()
socket = ctx.socket(zmq.DEALER)
poll = zmq.Poller()
poll.register(socket, zmq.POLLIN)
socket.connect('tcp://{}:{}'.format(HOST, PORT))
socket.send('foo')
sockets = dict(poll.poll(2000))
assert sockets.get(socket) != zmq.POLLIN
time.sleep(.2)
assert servant.stats['summary']['received_messages'] == 1
assert servant.stats['summary']['processed_messages'] == 0
assert servant.stats['summary']['malicious_messages'] == 1
servant.stop()
servant_thread.join()
def test_default_handler():
ipc_name = 'smite-test-{}'.format(uuid.uuid1().hex)
default_handler = lambda t: t
servant = Servant()
servant.set_default_handler(default_handler)
servant.bind_ipc(ipc_name)
servant.run(True)
msg_num = 10
client = RClient('ipc://{}'.format(ipc_name))
for _ in range(msg_num):
msg_txt = uuid.uuid1().hex
random_msg_name = uuid.uuid1().hex
rep = client.send(random_msg_name, msg_txt)
assert rep == msg_txt
assert servant.stats['summary']['exceptions'] == 0
assert servant.stats['summary']['malicious_messages'] == 0
assert servant.stats['summary']['received_messages'] == msg_num
assert servant.stats['summary']['processed_messages'] == msg_num
client.close()
servant.stop()
def test_proxy():
host = '127.0.0.1'
proxy_port = '9000'
servant_port = '9001'
def echo(text):
return text
servant = Servant({'echo': echo})
servant.bind_tcp(host, servant_port)
servant_thread = Thread(target=servant.run)
servant_thread.start()
proxy = Proxy(host, servant_port)
proxy.bind(host, proxy_port)
proxy_thread = Thread(target=proxy.run)
proxy_thread.start()
class send_msg(object):
def __init__(self, message_name):
self.message_name = message_name
def __call__(self):
time.sleep(.3)
client = Client()
client.connect_tcp(host, proxy_port)
txt = uuid.uuid4().hex
res = client.send(self.message_name, args=(txt,))
assert res == txt
client.close()
messages_num = 10
client_threads = []
for i in xrange(messages_num):
thread = Thread(target=send_msg('echo'))
client_threads.append(thread)
thread.start()
time.sleep(1)
assert servant.stats['summary']['received_messages'] == messages_num
assert servant.stats['summary']['processed_messages'] == messages_num
assert servant.stats['summary']['exceptions'] == 0
servant.stop()
servant_thread.join()
proxy.stop()
proxy_thread.join()
def test_client_timeout():
timeout = 3
client = Client(default_timeout=timeout)
client.connect(CONNECTION_URI)
raised = False
start = time.time()
try:
client.send('dummy_method')
except ClientTimeout:
raised = True
assert raised
# it should take around 3 seconds
assert 2.5 < time.time() - start < 3.5
# servant should not get this message after start, it's just gone
class DummyException(Exception):
pass
def dummy_method():
raise DummyException
servant = Servant({'dummy_method': dummy_method})
servant.bind_tcp(HOST, PORT)
# run servant in separate thread and wait 3 seconds for message
servant_thread = Thread(target=servant.run)
servant_thread.start()
time.sleep(3)
servant.stop()
servant_thread.join()
client.close()
for thread_stats in servant.stats['threads'].values():
assert thread_stats['exceptions'] == 0
assert thread_stats['received_messages'] == 0
assert thread_stats['malicious_messages'] == 0
assert thread_stats['processed_messages'] == 0
assert servant.stats['summary']['exceptions'] == 0
assert servant.stats['summary']['received_messages'] == 0
assert servant.stats['summary']['malicious_messages'] == 0
assert servant.stats['summary']['processed_messages'] == 0
def test_noreply_message():
servant = Servant({'echo': lambda t: t})
servant.bind_tcp(HOST, PORT)
servant_thread = Thread(target=servant.run)
servant_thread.start()
client = Client()
client.connect(CONNECTION_URI)
client.send('echo', args=(uuid.uuid1().hex,), noreply=True)
time.sleep(2)
assert servant.stats['summary']['received_messages'] == 1
assert servant.stats['summary']['processed_messages'] == 1
servant.stop()
servant_thread.join()
client.close()
def test_exception_response():
exc_message = 'This is dummy exception message'
class DummyException(Exception):
pass
def raise_dummy_exc():
raise DummyException(exc_message)
servant = Servant({'raise_dummy_exc': raise_dummy_exc})
servant.bind_tcp(HOST, PORT)
servant.run(True)
client = Client()
client.connect(CONNECTION_URI)
raised = False
try:
client.send('raise_dummy_exc')
except MessageException, e:
assert e.message == exc_message
raised = True
def test_rclient():
ipc_name = 'smite-test-{}'.format(uuid.uuid1().hex)
servant = Servant({'echo': lambda t: t})
servant.bind_ipc(ipc_name)
servant.run(True)
msg_num = 10
client = RClient('ipc://{}'.format(ipc_name))
for _ in range(msg_num):
echo_txt = uuid.uuid1().hex
rep = client.send('echo', echo_txt)
assert rep == echo_txt
assert servant.stats['summary']['exceptions'] == 0
assert servant.stats['summary']['malicious_messages'] == 0
assert servant.stats['summary']['received_messages'] == msg_num
assert servant.stats['summary']['processed_messages'] == msg_num
client.close()
servant.stop()
def test_secure_messaging():
keys_dir = create_keys_dir()
def short_echo(text):
time.sleep(1)
return text
def long_echo(text):
time.sleep(2)
return text
send_msgs = ['short_echo', 'long_echo']
# generate keys for clients
client_secrets = [
utils.create_certificates(keys_dir, 'client-{}'.format(i))[1]
for i in range(2)
]
# generate keys for servant
servant_public, servant_secret = (
utils.create_certificates(keys_dir, 'servant')
)
servant = Servant({'short_echo': short_echo, 'long_echo': long_echo})
servant.enable_security(
os.path.join(keys_dir, 'public_keys'), servant_secret,
)
servant.bind_tcp(HOST, PORT)
servant_thread = Thread(target=servant.run)
servant_thread.start()
class send_msg(object):
def __init__(self, message_name, client_secret):
self.message_name = message_name
self.client_secret = client_secret
def __call__(self):
client = Client()
client.enable_security(self.client_secret, servant_public)
client.connect(CONNECTION_URI)
txt = uuid.uuid4().hex
res = client.send(self.message_name, args=(txt,))
assert res == txt
client.close()
client_threads = []
for client_secret, method_name in zip(client_secrets, send_msgs):
thread = Thread(target=send_msg(method_name, client_secret))
client_threads.append(thread)
thread.start()
# long echo takes 2 seconds
time.sleep(2.5)
assert servant.stats['summary']['received_messages'] == 2
assert servant.stats['summary']['processed_messages'] == 2
assert servant.stats['summary']['exceptions'] == 0
servant.stop()
servant_thread.join()
for client_thread in client_threads:
client_thread.join()
def test_multiple_clients():
def short_echo(text):
time.sleep(1)
return text
def long_echo(text):
time.sleep(2)
return text
called = {
'extract_one': False,
'extract_two': False,
'process_one': False,
'process_two': False,
}
def extract_one(msg):
called['extract_one'] = True
# 'foo:bar:real_message' -> 'bar:real_message'
return msg.split(':', 1)[1]
def extract_two(msg):
called['extract_two'] = True
# 'bar:real_message' -> 'real_message'
return msg.split('-', 1)[1]
servant = Servant(
{'short_echo': short_echo, 'long_echo': long_echo},
message_extractors=(extract_one, extract_two),
)
servant.bind_tcp(HOST, PORT)
servant.run(run_in_background=True)
def process_one(msg):
called['process_one'] = True
return 'bar-{}'.format(msg)
def process_two(msg):
called['process_two'] = True
return 'foo:{}'.format(msg)
def send_msg(msg_name):
client = Client(message_processors=(process_one, process_two))
client.connect(CONNECTION_URI)
msg_txt = uuid.uuid4().hex
res = client.send(msg_name, args=(msg_txt,))
assert res == msg_txt
client.close()
client_threads = []
# send short_echo
thread = Thread(target=send_msg('short_echo'))
client_threads.append(thread)
thread.start()
# send long_echo
thread = Thread(target=send_msg('long_echo'))
client_threads.append(thread)
thread.start()
# long echo takes 2 seconds
time.sleep(2.5)
assert servant.stats['summary']['received_messages'] == 2
assert servant.stats['summary']['processed_messages'] == 2
assert servant.stats['summary']['exceptions'] == 0
for was_called in called.values():
assert was_called
servant.stop()
for client_thread in client_threads:
client_thread.join()