def encodedecode(pkt_size, key):
pkt = Packet()
typ = PacketType.APP if key else PacketType.CLIENT_HELLO
typ = PacketType.APP
pkt.msg = os.urandom(pkt_size)
pkt.hdr = PacketHeader.create(False, 0, typ, SeqNum(1), SeqNum(1), 0)
pkt.hdr.length = len(pkt.msg)
pkt.hdr.count = 1
datagram = pkt.to_bytes(key)
hdr = PacketHeader.from_bytes(True, datagram)
pkt = Packet.from_bytes(hdr, key, datagram)
def setUp(self):
# launch a server thread, connect one client
self.server_sock, self.client_sock = MockUDPSocket.mkpair()
self.client = TestClient(self.client_sock)
self.ctxt = ServerContext(TestHandler())
self.thread = UdpServerThread(self.server_sock, self.ctxt)
self.thread.start()
t0 = time.time()
connected = False
while not connected:
datagram, addr = self.server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(True, datagram)
self.thread.append(addr, hdr, datagram)
self.client.update()
time.sleep(1 / 60)
if time.time() - t0 > .5:
self.fail("failed to connect")
# test that both sides are connected
if self.client.conn.status == ConnectionStatus.CONNECTED:
for addr, other in self.ctxt.connections.items():
if other.status == ConnectionStatus.CONNECTED:
self.server_client = other
connected = True
self.assertEqual(self.client.conn.status, ConnectionStatus.CONNECTED)
def test_server_send_recv_large(self):
# send a large payload to the server
# send the payload back to the client
# test sending fragmented messages in both directions
payload = os.urandom(10 * 1024)
self.client.send(payload)
# TODO: investigate why this takes .5 seconds
# it should be 10KB / MAX_PAYLOAD_SIZE ~= 8 packets
# round tripping 18 packets should take around 16/60 ~= 0.25 seconds
# the test contsistently takes double the expected time
t0 = time.time()
while not self.client.conn.incoming_messages:
datagram, addr = self.server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(True, datagram)
self.thread.append(addr, hdr, datagram)
self.client.update()
time.sleep(1 / 60)
if time.time() - t0 > 1.0:
self.fail("failed to receive message")
self.assertEqual(self.client.conn.incoming_messages[0][1], payload)
def test_conn_receive_datagram_1(self):
key = b"0" * 16
hdr = PacketHeader.create(True, 0, PacketType.APP, SeqNum(1),
SeqNum(1), 0)
msgs = [
PendingMessage(SeqNum(1), PacketType.APP, b"hello world1", None, 0)
]
pkt = Packet.create(hdr, msgs)
datagram = pkt.to_bytes(key)
conn = ConnectionBase(False, None)
conn.session_key_bytes = key
conn._recv_datagram(PacketHeader.from_bytes(False, datagram), datagram)
self.assertEqual(SeqNum(1), conn.bitfield_pkt.current_seqnum)
self.assertEqual(1, len(conn.incoming_messages))
self.assertEqual(msgs[0].payload, conn.incoming_messages[0][1])
def test_server_connection_keepalive(self):
# sleep for 1 seconds to test send/recev keep alive
for i in range(60):
datagram, addr = self.server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(True, datagram)
self.thread.append(addr, hdr, datagram)
self.client.update()
time.sleep(1 / 60)
self.assertTrue(self.client.conn.stats.received >= 10)
def test_server_connect_invalid_key(self):
# test that a client does not connect if the public key does
# not match the server public key
server_sock, client_sock = MockUDPSocket.mkpair()
client = TestClient(client_sock)
key = EllipticCurvePrivateKey.new()
key2 = EllipticCurvePrivateKey.new().getPublicKey()
client.conn.setServerPublicKey(key2)
ctxt = ServerContext(TestHandler(), key)
thread = UdpServerThread(server_sock, ctxt)
thread.start()
timedout=None
def callback(success):
nonlocal timedout
timedout = not success
client.conn.outgoing_timeout = .25
client.conn.connection_callback = callback
t0 = time.time()
connected = False
with self.assertRaises(EllipticCurvePublicKey.InvalidSignature):
while not connected:
datagram, addr = server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(True, datagram)
thread.append(addr, hdr, datagram)
client.update()
time.sleep(1/60)
if time.time() - t0 > .5:
self.fail("failed to connect")
# test that both sides are connected
if client.conn.status == ConnectionStatus.CONNECTED:
for addr, other in ctxt.connections.items():
if other.status == ConnectionStatus.CONNECTED:
connected = True
self.assertEqual(client.conn.status, ConnectionStatus.DISCONNECTED)
self.assertFalse(timedout)
ctxt._active = False
# append an invalid packet to wake up the server
thread.append(("0.0.0.0", 0), None, b"")
thread.join()
def update(self, delta_t=0):
self.conn.update()
datagram, addr = self.sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(False, datagram)
self.conn._recv_datagram(hdr, datagram)
pkt = self.conn._build_packet()
if pkt:
datagram = self.conn._encode_packet(pkt)
self.sock.sendto(datagram, self.addr)
def test_server_send_recv(self):
self.client.send(b"hello")
t0 = time.time()
while not self.client.conn.incoming_messages:
datagram, addr = self.server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(True, datagram)
self.thread.append(addr, hdr, datagram)
self.client.update()
time.sleep(1 / 60)
if time.time() - t0 > .5:
self.fail("failed to receive")
def test_server_connect(self):
# test that a client can connect to the server
server_sock, client_sock = MockUDPSocket.mkpair()
client = TestClient(client_sock)
ctxt = ServerContext(TestHandler())
thread = UdpServerThread(server_sock, ctxt)
thread.start()
timedout = None
def callback(success):
nonlocal timedout
timedout = not success
client.conn.outgoing_timeout = .25
client.conn.connection_callback = callback
t0 = time.time()
connected = False
while not connected:
datagram, addr = server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(True, datagram)
thread.append(addr, hdr, datagram)
client.update()
time.sleep(1 / 60)
if time.time() - t0 > .5:
self.fail("failed to connect")
# test that both sides are connected
if client.conn.status == ConnectionStatus.CONNECTED:
for addr, other in ctxt.connections.items():
if other.status == ConnectionStatus.CONNECTED:
connected = True
self.assertEqual(client.conn.status, ConnectionStatus.CONNECTED)
self.assertFalse(timedout)
ctxt._active = False
# append an invalid packet to wake up the server
thread.append(("0.0.0.0", 0), None, b"")
thread.join()
def test_conn_send_datagram_simple(self):
"""
"""
payload = b"hello world"
conn = ConnectionBase(False, None)
conn._send_type(PacketType.CLIENT_HELLO, payload, RetryMode.NONE, None)
pkt = conn._build_packet()
datagram = conn._encode_packet(pkt)
self.assertEqual(
PacketHeader.SIZE + PacketHeader.CRC_SIZE + Packet.overhead(1) +
len(payload), len(datagram))
hdr = PacketHeader.from_bytes(True, datagram)
pkt = Packet.from_bytes(hdr, None, datagram)
self.assertEqual(pkt.msg[2:], payload)
def test_packet_max_size(self):
# test that forming a packet of max payload size
# correctly builds a packet with the maximum number of
# bytes for the default MTU size (1500)
hdr = PacketHeader()
dat = os.urandom(Packet.MAX_PAYLOAD_SIZE)
msg = PendingMessage(SeqNum(1), PacketType.APP, dat, None, 0)
pkt = Packet.create(hdr, [msg])
datagram = pkt.to_bytes(b"0" * 16)
self.assertEqual(len(datagram), Packet.MAX_SIZE)
self.assertEqual(len(datagram), 1472) # mtu - 28
datagram = pkt.to_bytes(None)
self.assertEqual(len(datagram), Packet.MAX_SIZE_CRC)
def test_server_datagram_corrupt(self):
self.client.send(b"hello1")
t0 = time.time()
while self.server_client.stats.dropped == 0:
datagram, addr = self.server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
# valid header, invalid tag
datagram = datagram[:-16] + b'0' * 16
hdr = PacketHeader.from_bytes(True, datagram)
self.thread.append(addr, hdr, datagram)
self.client.update()
time.sleep(1 / 60)
if time.time() - t0 > .5:
self.fail("failed to receive")
self.assertEqual(self.server_client.stats.dropped, 1)
def test_server_duplicate_datagram(self):
self.client.send(b"hello1")
self.client.send(b"hello2")
self.client.send(b"hello3")
t0 = time.time()
while self.server_client.stats.dropped == 0:
datagram, addr = self.server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(True, datagram)
self.thread.append(addr, hdr, datagram)
self.thread.append(addr, hdr, datagram)
self.client.update()
time.sleep(1 / 60)
if time.time() - t0 > .5:
self.fail("failed to receive")
self.assertEqual(self.server_client.stats.dropped, 1)
def test_server_disconnect(self):
# in the real world the client is not guaranteed to
# receive the reply from the server.
disconnected = False
def onDisconnectCallback(success):
nonlocal disconnected
disconnected = True
self.client.conn.disconnect(onDisconnectCallback)
t0 = time.time()
while not disconnected:
datagram, addr = self.server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
hdr = PacketHeader.from_bytes(True, datagram)
self.thread.append(addr, hdr, datagram)
self.client.update()
time.sleep(1/60)
if time.time() - t0 > .5:
self.fail("failed to receive disconnect")
def test_server_send_recv_multi(self):
self.client.send(b"hello1")
self.client.send(b"hello2")
self.client.send(b"hello3")
received = 0
t0 = time.time()
while len(self.client.conn.incoming_messages) < 3:
datagram, addr = self.server_sock.recvfrom(Packet.RECV_SIZE)
if datagram:
received += 1
hdr = PacketHeader.from_bytes(True, datagram)
self.thread.append(addr, hdr, datagram)
self.client.update()
time.sleep(1 / 60)
if time.time() - t0 > .5:
self.fail("failed to receive")
self.assertEqual(received, 1)
self.assertEqual(len(self.client.conn.incoming_messages), 3)
def test_packet_set_max_size(self):
# test that forming a packet of max payload size
# correctly builds a packet with the maximum number of
# bytes for a custom MTU size.
#
Packet.setMTU(512)
hdr = PacketHeader()
dat = os.urandom(Packet.MAX_PAYLOAD_SIZE)
msg = PendingMessage(SeqNum(1), PacketType.APP, dat, None, 0)
pkt = Packet.create(hdr, [msg])
datagram1 = pkt.to_bytes(b"0" * 16)
datagram2 = pkt.to_bytes(None)
self.assertEqual(len(datagram1), 484) # mtu - 28
self.assertEqual(len(datagram1), Packet.MAX_SIZE)
self.assertEqual(len(datagram2), Packet.MAX_SIZE_CRC)
# reset mtu for subsequent tests
Packet.setMTU(1500)
def test_conn_sendrecv_datagram_encrypted(self):
"""
"""
key = b"0" * 16
payload = b"hello world"
client = ConnectionBase(False, None)
client.session_key_bytes = key
client.status = ConnectionStatus.CONNECTED
server = ConnectionBase(True, None)
server.session_key_bytes = key
client.send(payload)
pkt = client._build_packet()
datagram = client._encode_packet(pkt)
hdr = PacketHeader.from_bytes(True, datagram)
server._recv_datagram(hdr, datagram)
self.assertEqual(1, len(server.incoming_messages))
self.assertEqual(payload, server.incoming_messages[0][1])
def test_conn_handshake(self):
"""
run through the entire handshake
happy path integration test
"""
current_time = 0
def clock():
nonlocal current_time
current_time += 1
return current_time
client = ClientServerConnection(('0.0.0.0', 1234))
client.clock = clock
#client.status = ConnectionStatus.CONNECTED
ctxt = ServerContext(EventHandler(), None)
server = ServerClientConnection(ctxt, ('0.0.0.0', 1235))
server.clock = clock
#server.status = ConnectionStatus.CONNECTED
#with self.subTest(mode='send client hello'):
client._sendClientHello()
datagram = client._encode_packet(client._build_packet())
self.assertEqual(len(datagram),
Packet.MAX_PAYLOAD_SIZE + PacketHeader.CRC_SIZE)
hdr = PacketHeader.from_bytes(True, datagram)
# update the context with a reference to the server side client connection
# user later to finalize the connection
ctxt.temp_connections[server.addr] = server
server._recv_datagram(hdr, datagram)
datagram = server._encode_packet(server._build_packet())
with self.subTest(mode='recv server hello'):
hdr = PacketHeader.from_bytes(False, datagram)
client._recv_datagram(hdr, datagram)
with self.subTest(mode='recv challenge response'):
datagram = client._encode_packet(client._build_packet())
hdr = PacketHeader.from_bytes(True, datagram)
server._recv_datagram(hdr, datagram)
self.assertEqual(ConnectionStatus.CONNECTED, server.status)
self.assertEqual(ConnectionStatus.CONNECTED, client.status)
with self.subTest(mode='server send test'):
greeting = b"hello client"
server.send(greeting)
datagram = server._encode_packet(server._build_packet())
hdr = PacketHeader.from_bytes(False, datagram)
client._recv_datagram(hdr, datagram)
self.assertEqual(1, len(client.incoming_messages))
self.assertEqual(b"hello client", client.incoming_messages[0][1])
with self.subTest(mode='client send test'):
greeting = b"hello server"
client.send(greeting)
datagram = client._encode_packet(client._build_packet())
hdr = PacketHeader.from_bytes(True, datagram)
server._recv_datagram(hdr, datagram)
self.assertEqual(1, len(server.incoming_messages))
self.assertEqual(greeting, server.incoming_messages[0][1])
self.assertTrue(server.latency > 0)
def recvinto(sock, thread):
r, w, x = select.select([sock], [], [], 0)
if r:
datagram, addr = sock.recvfrom(Packet.RECV_SIZE)
hdr = PacketHeader.from_bytes(True, datagram)
thread.append(addr, hdr, datagram)