def test_correct_request(self):
test_host = 'frob.nitz'
test_path = '/fnord'
ws = WSConnection(SERVER)
nonce = bytes(random.getrandbits(8) for x in range(0, 16))
nonce = base64.b64encode(nonce)
request = b"GET " + test_path.encode('ascii') + b" HTTP/1.1\r\n"
request += b'Host: ' + test_host.encode('ascii') + b'\r\n'
request += b'Connection: Upgrade\r\n'
request += b'Upgrade: WebSocket\r\n'
request += b'Sec-WebSocket-Version: 13\r\n'
request += b'Sec-WebSocket-Key: ' + nonce + b'\r\n'
request += b'\r\n'
ws.receive_bytes(request)
event = next(ws.events())
assert isinstance(event, ConnectionRequested)
ws.accept(event)
data = ws.bytes_to_send()
response, headers = data.split(b'\r\n', 1)
version, code, reason = response.split(b' ')
headers = parse_headers(headers)
accept_token = ws._generate_accept_token(nonce)
assert int(code) == 101
assert headers['connection'].lower() == 'upgrade'
assert headers['upgrade'].lower() == 'websocket'
assert headers['sec-websocket-accept'] == accept_token.decode('ascii')
def test_accept_subprotocol(self):
test_host = 'frob.nitz'
test_path = '/fnord'
ws = WSConnection(SERVER)
nonce = bytes(random.getrandbits(8) for x in range(0, 16))
nonce = base64.b64encode(nonce)
request = b'GET ' + test_path.encode('ascii') + b' HTTP/1.1\r\n'
request += b'Host: ' + test_host.encode('ascii') + b'\r\n'
request += b'Connection: Upgrade\r\n'
request += b'Upgrade: WebSocket\r\n'
request += b'Sec-WebSocket-Version: 13\r\n'
request += b'Sec-WebSocket-Key: ' + nonce + b'\r\n'
request += b'Sec-WebSocket-Protocol: one, two\r\n'
request += b'\r\n'
ws.receive_bytes(request)
event = next(ws.events())
assert isinstance(event, ConnectionRequested)
assert event.proposed_subprotocols == ['one', 'two']
ws.accept(event, 'two')
data = ws.bytes_to_send()
response, headers = data.split(b'\r\n', 1)
version, code, reason = response.split(b' ')
headers = parse_headers(headers)
assert int(code) == 101
assert headers['sec-websocket-protocol'] == 'two'
def new_conn(sock):
global count
print("test_server.py received connection {}".format(count))
count += 1
ws = WSConnection(SERVER, extensions=[PerMessageDeflate()])
closed = False
while not closed:
try:
data = sock.recv(65535)
except socket.error:
data = None
ws.receive_bytes(data or None)
for event in ws.events():
if isinstance(event, ConnectionRequested):
ws.accept(event)
elif isinstance(event, DataReceived):
ws.send_data(event.data, event.message_finished)
elif isinstance(event, ConnectionClosed):
closed = True
if not data:
closed = True
try:
data = ws.bytes_to_send()
sock.sendall(data)
except socket.error:
closed = True
sock.close()
def test_unwanted_extension_negotiation(self):
test_host = 'frob.nitz'
test_path = '/fnord'
ext = FakeExtension(accept_response=False)
ws = WSConnection(SERVER, extensions=[ext])
nonce = bytes(random.getrandbits(8) for x in range(0, 16))
nonce = base64.b64encode(nonce)
request = b"GET " + test_path.encode('ascii') + b" HTTP/1.1\r\n"
request += b'Host: ' + test_host.encode('ascii') + b'\r\n'
request += b'Connection: Upgrade\r\n'
request += b'Upgrade: WebSocket\r\n'
request += b'Sec-WebSocket-Version: 13\r\n'
request += b'Sec-WebSocket-Key: ' + nonce + b'\r\n'
request += b'Sec-WebSocket-Extensions: pretend\r\n'
request += b'\r\n'
ws.receive_bytes(request)
event = next(ws.events())
assert isinstance(event, ConnectionRequested)
ws.accept(event)
data = ws.bytes_to_send()
response, headers = data.split(b'\r\n', 1)
version, code, reason = response.split(b' ')
headers = parse_headers(headers)
assert 'sec-websocket-extensions' not in headers
def test_accept_wrong_subprotocol(self):
test_host = 'frob.nitz'
test_path = '/fnord'
ws = WSConnection(SERVER)
nonce = bytes(random.getrandbits(8) for x in range(0, 16))
nonce = base64.b64encode(nonce)
request = b'GET ' + test_path.encode('ascii') + b' HTTP/1.1\r\n'
request += b'Host: ' + test_host.encode('ascii') + b'\r\n'
request += b'Connection: Upgrade\r\n'
request += b'Upgrade: WebSocket\r\n'
request += b'Sec-WebSocket-Version: 13\r\n'
request += b'Sec-WebSocket-Key: ' + nonce + b'\r\n'
request += b'Sec-WebSocket-Protocol: one, two\r\n'
request += b'\r\n'
ws.receive_bytes(request)
event = next(ws.events())
assert isinstance(event, ConnectionRequested)
assert event.proposed_subprotocols == ['one', 'two']
with pytest.raises(ValueError):
ws.accept(event, 'three')
def new_conn(reader, writer):
ws = WSConnection(SERVER, extensions=[PerMessageDeflate()])
closed = False
while not closed:
try:
data = yield from reader.read(65535)
except ConnectionError:
data = None
ws.receive_bytes(data or None)
for event in ws.events():
if isinstance(event, ConnectionRequested):
ws.accept(event)
elif isinstance(event, DataReceived):
ws.send_data(event.data, event.final)
elif isinstance(event, ConnectionClosed):
closed = True
if data is None:
break
try:
data = ws.bytes_to_send()
writer.write(data)
yield from writer.drain()
except (ConnectionError, OSError):
closed = True
if closed:
break
writer.close()
def new_conn(reader, writer):
global count
print("test_server.py received connection {}".format(count))
count += 1
ws = WSConnection(SERVER, extensions=[PerMessageDeflate()])
closed = False
while not closed:
try:
data = yield from reader.read(65535)
except ConnectionError:
data = None
ws.receive_bytes(data or None)
for event in ws.events():
if isinstance(event, ConnectionRequested):
ws.accept(event)
elif isinstance(event, DataReceived):
ws.send_data(event.data, event.message_finished)
elif isinstance(event, ConnectionClosed):
closed = True
if not data:
closed = True
try:
data = ws.bytes_to_send()
writer.write(data)
yield from writer.drain()
except (ConnectionError, OSError):
closed = True
writer.close()
def create_connection(self):
server = WSConnection(SERVER)
client = WSConnection(CLIENT, host='localhost', resource='foo')
server.receive_bytes(client.bytes_to_send())
event = next(server.events())
assert isinstance(event, ConnectionRequested)
server.accept(event)
client.receive_bytes(server.bytes_to_send())
assert isinstance(next(client.events()), ConnectionEstablished)
return client, server
async def ws_adapter(in_q, out_q, client, _):
"""A simple, queue-based Curio-Sans-IO websocket bridge."""
client.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)
wsconn = WSConnection(SERVER)
closed = False
while not closed:
wstask = await spawn(client.recv, 65535)
outqtask = await spawn(out_q.get)
async with TaskGroup([wstask, outqtask]) as g:
task = await g.next_done()
result = await task.join()
await g.cancel_remaining()
if task is wstask:
wsconn.receive_bytes(result)
for event in wsconn.events():
cl = event.__class__
if cl in DATA_TYPES:
await in_q.put(event.data)
elif cl is ConnectionRequested:
# Auto accept. Maybe consult the handler?
wsconn.accept(event)
elif cl is ConnectionClosed:
# The client has closed the connection.
await in_q.put(None)
closed = True
else:
print(event)
await client.sendall(wsconn.bytes_to_send())
else:
# We got something from the out queue.
if result is None:
# Terminate the connection.
print("Closing the connection.")
wsconn.close()
closed = True
else:
wsconn.send_data(result)
payload = wsconn.bytes_to_send()
await client.sendall(payload)
print("Bridge done.")
def test_extension_negotiation_with_our_parameters(self):
test_host = 'frob.nitz'
test_path = '/fnord'
offered_params = 'parameter1=value3; parameter2=value4'
ext_params = 'parameter1=value1; parameter2=value2'
ext = FakeExtension(accept_response=ext_params)
ws = WSConnection(SERVER, extensions=[ext])
nonce = bytes(random.getrandbits(8) for x in range(0, 16))
nonce = base64.b64encode(nonce)
request = b"GET " + test_path.encode('ascii') + b" HTTP/1.1\r\n"
request += b'Host: ' + test_host.encode('ascii') + b'\r\n'
request += b'Connection: Upgrade\r\n'
request += b'Upgrade: WebSocket\r\n'
request += b'Sec-WebSocket-Version: 13\r\n'
request += b'Sec-WebSocket-Key: ' + nonce + b'\r\n'
request += b'Sec-WebSocket-Extensions: ' + \
ext.name.encode('ascii') + b'; ' + \
offered_params.encode('ascii') + b'\r\n'
request += b'\r\n'
ws.receive_bytes(request)
event = next(ws.events())
assert isinstance(event, ConnectionRequested)
ws.accept(event)
data = ws.bytes_to_send()
response, headers = data.split(b'\r\n', 1)
version, code, reason = response.split(b' ')
headers = parse_headers(headers)
assert ext.offered == '%s; %s' % (ext.name, offered_params)
assert headers['sec-websocket-extensions'] == \
'%s; %s' % (ext.name, ext_params)
def handle_connection(stream):
'''
Handle a connection.
The server operates a request/response cycle, so it performs a synchronous
loop:
1) Read data from network into wsproto
2) Get next wsproto event
3) Handle event
4) Send data from wsproto to network
:param stream: a socket stream
'''
ws = WSConnection(ConnectionType.SERVER)
# events is a generator that yields websocket event objects. Usually you
# would say `for event in ws.events()`, but the synchronous nature of this
# server requires us to use next(event) instead so that we can interleave
# the network I/O.
events = ws.events()
running = True
while running:
# 1) Read data from network
in_data = stream.recv(RECEIVE_BYTES)
print('Received {} bytes'.format(len(in_data)))
ws.receive_bytes(in_data)
# 2) Get next wsproto event
try:
event = next(events)
except StopIteration:
print('Client connection dropped unexpectedly')
return
# 3) Handle event
if isinstance(event, ConnectionRequested):
# Negotiate new WebSocket connection
print('Accepting WebSocket upgrade')
ws.accept(event)
elif isinstance(event, ConnectionClosed):
# Print log message and break out
print('Connection closed: code={}/{} reason={}'.format(
event.code.value, event.code.name, event.reason))
running = False
elif isinstance(event, TextReceived):
# Reverse text and send it back to wsproto
print('Received request and sending response')
ws.send_data(event.data[::-1])
elif isinstance(event, PingReceived):
# wsproto handles ping events for you by placing a pong frame in
# the outgoing buffer. You should not call pong() unless you want to
# send an unsolicited pong frame.
print('Received ping and sending pong')
else:
print('Unknown event: {!r}'.format(event))
# 4) Send data from wsproto to network
out_data = ws.bytes_to_send()
print('Sending {} bytes'.format(len(out_data)))
stream.send(out_data)
class WebSocketEndpoint(AbstractEndpoint):
"""
Implements websocket endpoints.
Subprotocol negotiation is currently not supported.
"""
__slots__ = ('ctx', '_client', '_ws')
def __init__(self, ctx: Context, client: BaseHTTPClientConnection):
self.ctx = ctx
self._client = client
self._ws = WSConnection(ConnectionType.SERVER)
def _process_ws_events(self):
for event in self._ws.events():
if isinstance(event, ConnectionRequested):
self._ws.accept(event)
self.on_connect()
elif isinstance(event, DataReceived):
self.on_data(event.data)
elif isinstance(event, ConnectionClosed):
self.on_close()
bytes_to_send = self._ws.bytes_to_send()
if bytes_to_send:
self._client.write(bytes_to_send)
def begin_request(self, request: HTTPRequest):
trailing_data = self._client.upgrade()
self._ws.receive_bytes(trailing_data)
self._process_ws_events()
def receive_body_data(self, data: bytes) -> None:
self._ws.receive_bytes(data)
self._process_ws_events()
def send_message(self, payload: Union[str, bytes]) -> None:
"""
Send a message to the client.
:param payload: either a unicode string or a bytestring
"""
self._ws.send_data(payload)
bytes_to_send = self._ws.bytes_to_send()
self._client.write(bytes_to_send)
def close(self) -> None:
"""Close the websocket."""
self._ws.close()
self._process_ws_events()
def on_connect(self) -> None:
"""Called when the websocket handshake has been done."""
def on_close(self) -> None:
"""Called when the connection has been closed."""
def on_data(self, data: bytes) -> None:
"""Called when there is new data from the peer."""