90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
import json
|
|
import socket
|
|
import ssl
|
|
import unittest
|
|
|
|
from colony.transport.framing import send_frame, recv_frame
|
|
from colony.transport.listener import Listener
|
|
from colony.transport.whisperer import Whisperer
|
|
from tests.helpers import CERT, KEY, start_listener
|
|
|
|
PAYLOAD = {"command": {"role": "X", "name": "y"}}
|
|
|
|
|
|
class StubDispatcher:
|
|
def __init__(self, result=None, error=None):
|
|
self.result, self.error = result, error
|
|
|
|
def dispatch(self, payload):
|
|
if self.error:
|
|
raise self.error
|
|
return self.result
|
|
|
|
|
|
def whisper(listener, **kwargs):
|
|
return Whisperer("127.0.0.1", listener.port, CERT, KEY).send(PAYLOAD, **kwargs)
|
|
|
|
|
|
class TlsTests(unittest.TestCase):
|
|
def test_mutual_tls_round_trip(self):
|
|
listener = start_listener(StubDispatcher(result={"x": 1}))
|
|
self.assertEqual(whisper(listener), {"status": "ok", "result": {"x": 1}})
|
|
|
|
def test_large_payload_round_trip(self):
|
|
big = {"data": "y" * 500_000}
|
|
listener = start_listener(StubDispatcher(result=big))
|
|
self.assertEqual(whisper(listener)["result"], big)
|
|
|
|
def test_client_without_certificate_is_rejected(self):
|
|
listener = start_listener(StubDispatcher(result={"x": 1}))
|
|
|
|
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
|
context.check_hostname = False
|
|
context.load_verify_locations(cafile=CERT)
|
|
with self.assertRaises((ssl.SSLError, OSError)):
|
|
with socket.create_connection(("127.0.0.1", listener.port), timeout=5) as sock:
|
|
with context.wrap_socket(sock) as tls:
|
|
send_frame(tls, json.dumps(PAYLOAD).encode())
|
|
recv_frame(tls)
|
|
|
|
self.assertEqual(whisper(listener)["status"], "ok")
|
|
|
|
def test_plain_tcp_client_does_not_break_the_listener(self):
|
|
listener = start_listener(StubDispatcher(result=1))
|
|
with socket.create_connection(("127.0.0.1", listener.port)) as sock:
|
|
sock.sendall(b"GET / HTTP/1.0\r\n\r\n")
|
|
self.assertEqual(whisper(listener)["result"], 1)
|
|
|
|
def test_errors_do_not_leak_details(self):
|
|
listener = start_listener(StubDispatcher(error=RuntimeError("secret-path /etc/x")))
|
|
with self.assertLogs("Listener", level="ERROR"):
|
|
response = whisper(listener)
|
|
self.assertEqual(response, {"status": "error", "error": "RuntimeError"})
|
|
self.assertNotIn("secret", json.dumps(response))
|
|
listener.dispatcher.error = None
|
|
listener.dispatcher.result = 2
|
|
self.assertEqual(whisper(listener)["result"], 2)
|
|
|
|
def test_unserializable_result_becomes_an_error(self):
|
|
listener = start_listener(StubDispatcher(result={"bad": object()}))
|
|
with self.assertLogs("Listener", level="ERROR"):
|
|
self.assertEqual(whisper(listener)["status"], "error")
|
|
|
|
def test_silent_server_times_out(self):
|
|
server = socket.socket()
|
|
self.addCleanup(server.close)
|
|
server.bind(("127.0.0.1", 0))
|
|
server.listen()
|
|
with self.assertRaises(TimeoutError):
|
|
Whisperer("127.0.0.1", server.getsockname()[1], CERT, KEY).send(PAYLOAD, timeout=0.5)
|
|
|
|
def test_bind_on_used_port_raises(self):
|
|
first = start_listener(StubDispatcher())
|
|
second = Listener(StubDispatcher(), host="127.0.0.1", port=first.port, cert_file=CERT, keyfile=KEY)
|
|
with self.assertRaises(OSError):
|
|
second.bind()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|