kopia lustrzana https://github.com/espressif/esp-idf
ws_client tests: Updated example test to use WebsSocket package
Added a new test for closing connection with close framespull/5919/head
David Cermak
bot
rodzic
5e9f8b52e7
commit
44c553fd14
|
|
@ -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))
|
||||
|
|
|
|||
Ładowanie…
Reference in New Issue