mirror of
https://github.com/espressif/esp-protocols.git
synced 2025-07-29 10:17:30 +02:00
websocket_client: added example_test with a local websocket server
- Added a example test that connects to a local python websocket server. - Added readme for websocket_client example. Closes IDF-907 * Original commit: espressif/esp-idf@67c5225c14
This commit is contained in:
committed by
gabsuren
parent
f21a2f32e0
commit
3b0488cfdc
@ -1,15 +1,158 @@
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
import re
|
||||
import os
|
||||
import socket
|
||||
import hashlib
|
||||
import base64
|
||||
from threading import Thread
|
||||
|
||||
import ttfw_idf
|
||||
|
||||
|
||||
@ttfw_idf.idf_example_test(env_tag="Example_WIFI", ignore=True)
|
||||
def get_my_ip():
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
try:
|
||||
# doesn't even have to be reachable
|
||||
s.connect(('10.255.255.255', 1))
|
||||
IP = s.getsockname()[0]
|
||||
except Exception:
|
||||
IP = '127.0.0.1'
|
||||
finally:
|
||||
s.close()
|
||||
return IP
|
||||
|
||||
|
||||
# Simple Websocket server for testing purposes
|
||||
class Websocket:
|
||||
HEADER_LEN = 6
|
||||
|
||||
def __init__(self, port):
|
||||
self.port = port
|
||||
self.socket = socket.socket()
|
||||
self.socket.settimeout(10.0)
|
||||
|
||||
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()
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.server_thread.join()
|
||||
self.socket.close()
|
||||
self.conn.close()
|
||||
|
||||
def run_server(self):
|
||||
self.conn, address = self.socket.accept() # accept new connection
|
||||
self.conn.settimeout(10.0)
|
||||
print("Connection from: {}".format(address))
|
||||
|
||||
self.establish_connection()
|
||||
|
||||
# Echo data until client closes connection
|
||||
self.echo_data()
|
||||
|
||||
def establish_connection(self):
|
||||
while True:
|
||||
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 echo_data(self):
|
||||
while(True):
|
||||
try:
|
||||
header = bytearray(self.conn.recv(self.HEADER_LEN, socket.MSG_WAITALL))
|
||||
if not header:
|
||||
# exit if data is not received
|
||||
return
|
||||
|
||||
# Remove mask bit
|
||||
payload_len = ~(1 << 7) & header[1]
|
||||
|
||||
payload = bytearray(self.conn.recv(payload_len, socket.MSG_WAITALL))
|
||||
frame = header + payload
|
||||
|
||||
decoded_payload = self.decode_frame(frame)
|
||||
|
||||
echo_frame = self.encode_frame(decoded_payload)
|
||||
self.conn.send(echo_frame)
|
||||
except socket.error as err:
|
||||
print("Stopped echoing data: {}".format(err))
|
||||
|
||||
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)
|
||||
frame.append(len(payload))
|
||||
frame += payload
|
||||
|
||||
return frame
|
||||
|
||||
|
||||
def test_echo(dut):
|
||||
dut.expect("WEBSOCKET_EVENT_CONNECTED")
|
||||
for i in range(0, 10):
|
||||
dut.expect(re.compile(r"Received=hello (\d)"))
|
||||
dut.expect("Websocket Stopped")
|
||||
|
||||
|
||||
@ttfw_idf.idf_example_test(env_tag="Example_WIFI")
|
||||
def test_examples_protocol_websocket(env, extra_data):
|
||||
"""
|
||||
steps: |
|
||||
steps:
|
||||
1. join AP
|
||||
2. connect to ws://echo.websocket.org
|
||||
2. connect to uri specified in the config
|
||||
3. send and receive data
|
||||
"""
|
||||
dut1 = env.get_dut("websocket", "examples/protocols/websocket", dut_class=ttfw_idf.ESP32DUT)
|
||||
@ -18,15 +161,33 @@ def test_examples_protocol_websocket(env, extra_data):
|
||||
bin_size = os.path.getsize(binary_file)
|
||||
ttfw_idf.log_performance("websocket_bin_size", "{}KB".format(bin_size // 1024))
|
||||
ttfw_idf.check_performance("websocket_bin_size", bin_size // 1024)
|
||||
|
||||
try:
|
||||
if "CONFIG_WEBSOCKET_URI_FROM_STDIN" in dut1.app.get_sdkconfig():
|
||||
uri_from_stdin = True
|
||||
else:
|
||||
uri = dut1.app.get_sdkconfig()["CONFIG_WEBSOCKET_URI"].strip('"')
|
||||
uri_from_stdin = False
|
||||
|
||||
except Exception:
|
||||
print('ENV_TEST_FAILURE: Cannot find uri settings in sdkconfig')
|
||||
raise
|
||||
|
||||
# start test
|
||||
dut1.start_app()
|
||||
dut1.expect("Waiting for wifi ...")
|
||||
dut1.expect("Connection established...", timeout=30)
|
||||
dut1.expect("WEBSOCKET_EVENT_CONNECTED")
|
||||
for i in range(0, 10):
|
||||
dut1.expect(re.compile(r"Sending hello (\d)"))
|
||||
dut1.expect(re.compile(r"Received=hello (\d)"))
|
||||
dut1.expect("Websocket Stopped")
|
||||
|
||||
if uri_from_stdin:
|
||||
server_port = 4455
|
||||
with Websocket(server_port):
|
||||
uri = "ws://{}:{}".format(get_my_ip(), server_port)
|
||||
print("DUT connecting to {}".format(uri))
|
||||
dut1.expect("Please enter uri of websocket endpoint", timeout=30)
|
||||
dut1.write(uri)
|
||||
test_echo(dut1)
|
||||
|
||||
else:
|
||||
print("DUT connecting to {}".format(uri))
|
||||
test_echo(dut1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
Reference in New Issue
Block a user