diff --git a/examples/protocols/websocket/example_test.py b/examples/protocols/websocket/example_test.py index 6dfdccf5d8..0bb085d265 100644 --- a/examples/protocols/websocket/example_test.py +++ b/examples/protocols/websocket/example_test.py @@ -3,12 +3,10 @@ from __future__ import unicode_literals import re import os import socket -import select -import hashlib -import base64 -import queue import random import string +from SimpleWebSocketServer import SimpleWebSocketServer, WebSocket +from tiny_test_fw import Utility from threading import Thread, Event import ttfw_idf @@ -26,159 +24,45 @@ def get_my_ip(): return IP +class TestEcho(WebSocket): + + def handleMessage(self): + self.sendMessage(self.data) + print('Server sent: {}'.format(self.data)) + + def handleConnected(self): + print('Connection from: {}'.format(self.address)) + + def handleClose(self): + print('{} closed the connection'.format(self.address)) + + # Simple Websocket server for testing purposes -class Websocket: - HEADER_LEN = 6 +class Websocket(object): + + def send_data(self, data): + for nr, conn in self.server.connections.items(): + conn.sendMessage(data) + + def run(self): + self.server = SimpleWebSocketServer('', self.port, TestEcho) + while not self.exit_event.is_set(): + self.server.serveonce() def __init__(self, port): self.port = port - self.socket = socket.socket() - self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.socket.settimeout(10.0) - self.send_q = queue.Queue() - self.shutdown = Event() + self.exit_event = Event() + self.thread = Thread(target=self.run) + self.thread.start() def __enter__(self): - try: - self.socket.bind(('', self.port)) - except socket.error as e: - print("Bind failed:{}".format(e)) - raise - - self.socket.listen(1) - self.server_thread = Thread(target=self.run_server) - self.server_thread.start() - return self def __exit__(self, exc_type, exc_value, traceback): - self.shutdown.set() - self.server_thread.join() - self.socket.close() - self.conn.close() - - def run_server(self): - self.conn, address = self.socket.accept() # accept new connection - self.socket.settimeout(10.0) - - print("Connection from: {}".format(address)) - - self.establish_connection() - print("WS established") - # Handle connection until client closes it, will echo any data received and send data from send_q queue - self.handle_conn() - - def establish_connection(self): - while not self.shutdown.is_set(): - try: - # receive data stream. it won't accept data packet greater than 1024 bytes - data = self.conn.recv(1024).decode() - if not data: - # exit if data is not received - raise - - if "Upgrade: websocket" in data and "Connection: Upgrade" in data: - self.handshake(data) - return - - except socket.error as err: - print("Unable to establish a websocket connection: {}".format(err)) - raise - - def handshake(self, data): - # Magic string from RFC - MAGIC_STRING = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" - headers = data.split("\r\n") - - for header in headers: - if "Sec-WebSocket-Key" in header: - client_key = header.split()[1] - - if client_key: - resp_key = client_key + MAGIC_STRING - resp_key = base64.standard_b64encode(hashlib.sha1(resp_key.encode()).digest()) - - resp = "HTTP/1.1 101 Switching Protocols\r\n" + \ - "Upgrade: websocket\r\n" + \ - "Connection: Upgrade\r\n" + \ - "Sec-WebSocket-Accept: {}\r\n\r\n".format(resp_key.decode()) - - self.conn.send(resp.encode()) - - def handle_conn(self): - while not self.shutdown.is_set(): - r,w,e = select.select([self.conn], [], [], 1) - try: - if self.conn in r: - self.echo_data() - - if not self.send_q.empty(): - self._send_data_(self.send_q.get()) - - except socket.error as err: - print("Stopped echoing data: {}".format(err)) - raise - - def echo_data(self): - header = bytearray(self.conn.recv(self.HEADER_LEN, socket.MSG_WAITALL)) - if not header: - # exit if socket closed by peer - return - - # Remove mask bit - payload_len = ~(1 << 7) & header[1] - - payload = bytearray(self.conn.recv(payload_len, socket.MSG_WAITALL)) - - if not payload: - # exit if socket closed by peer - return - frame = header + payload - - decoded_payload = self.decode_frame(frame) - print("Sending echo...") - self._send_data_(decoded_payload) - - def _send_data_(self, data): - frame = self.encode_frame(data) - self.conn.send(frame) - - def send_data(self, data): - self.send_q.put(data.encode()) - - def decode_frame(self, frame): - # Mask out MASK bit from payload length, this len is only valid for short messages (<126) - payload_len = ~(1 << 7) & frame[1] - - mask = frame[2:self.HEADER_LEN] - - encrypted_payload = frame[self.HEADER_LEN:self.HEADER_LEN + payload_len] - payload = bytearray() - - for i in range(payload_len): - payload.append(encrypted_payload[i] ^ mask[i % 4]) - - return payload - - def encode_frame(self, payload): - # Set FIN = 1 and OP_CODE = 1 (text) - header = (1 << 7) | (1 << 0) - - frame = bytearray([header]) - payload_len = len(payload) - - # If payload len is longer than 125 then the next 16 bits are used to encode length - if payload_len > 125: - frame.append(126) - frame.append(payload_len >> 8) - frame.append(0xFF & payload_len) - - else: - frame.append(payload_len) - - frame += payload - - return frame + self.exit_event.set() + self.thread.join(10) + if self.thread.is_alive(): + Utility.console_log('Thread cannot be joined', 'orange') def test_echo(dut): @@ -188,6 +72,11 @@ def test_echo(dut): print("All echos received") +def test_close(dut): + code = dut.expect(re.compile(r"WEBSOCKET: Received closed message with code=(\d*)"), timeout=60)[0] + print("Received close frame with code {}".format(code)) + + def test_recv_long_msg(dut, websocket, msg_len, repeats): send_msg = ''.join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for _ in range(msg_len)) @@ -246,6 +135,7 @@ def test_examples_protocol_websocket(env, extra_data): test_echo(dut1) # Message length should exceed DUT's buffer size to test fragmentation, default is 1024 byte test_recv_long_msg(dut1, ws, 2000, 3) + test_close(dut1) else: print("DUT connecting to {}".format(uri))