Commit b3045051 authored by Taddeüs Kroes's avatar Taddeüs Kroes

Started adding asynchronous methods

parent 364f446c
from websocket import websocket
from websocket import websocket, STATE_INIT, STATE_READ, STATE_WRITE, \
STATE_CLOSE
from server import Server
from frame import Frame, ControlFrame, OPCODE_CONTINUATION, OPCODE_TEXT, \
OPCODE_BINARY, OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG, CLOSE_NORMAL, \
CLOSE_GOING_AWAY, CLOSE_PROTOCOL_ERROR, CLOSE_NOACCEPT_DTYPE, \
CLOSE_INVALID_DATA, CLOSE_POLICY, CLOSE_MESSAGE_TOOBIG, \
CLOSE_MISSING_EXTENSIONS, CLOSE_UNABLE
CLOSE_MISSING_EXTENSIONS, CLOSE_UNABLE, read_frame, pop_frame, \
contains_frame
from connection import Connection
from message import Message, TextMessage, BinaryMessage
from errors import SocketClosed, HandshakeError, PingError, SSLError
from extension import Extension
from deflate_frame import DeflateFrame, WebkitDeflateFrame
#from multiplex import Multiplex
......@@ -194,12 +194,8 @@ class ControlFrame(Frame):
return code, reason
def receive_frame(sock):
"""
Receive a single frame on socket `sock`. The frame scheme is explained in
the docs of Frame.pack().
"""
b1, b2 = struct.unpack('!BB', recvn(sock, 2))
def decode_frame(reader):
b1, b2 = struct.unpack('!BB', reader.recvn(2))
final = bool(b1 & 0x80)
rsv1 = bool(b1 & 0x40)
......@@ -211,16 +207,16 @@ def receive_frame(sock):
payload_len = b2 & 0x7F
if payload_len == 126:
payload_len = struct.unpack('!H', recvn(sock, 2))
payload_len = struct.unpack('!H', reader.recvn(2))
elif payload_len == 127:
payload_len = struct.unpack('!Q', recvn(sock, 8))
payload_len = struct.unpack('!Q', reader.recvn(8))
if masked:
masking_key = recvn(sock, 4)
payload = mask(masking_key, recvn(sock, payload_len))
masking_key = reader.recvn(4)
payload = mask(masking_key, reader.recvn(payload_len))
else:
masking_key = ''
payload = recvn(sock, payload_len)
payload = reader.recvn(payload_len)
# Control frames have most significant bit 1
cls = ControlFrame if opcode & 0x8 else Frame
......@@ -229,14 +225,43 @@ def receive_frame(sock):
rsv1=rsv1, rsv2=rsv2, rsv3=rsv3)
def recvn(sock, n):
def receive_frame(sock):
return decode_frame(BlockingSocket(sock))
def read_frame(data):
reader = BufferedReader(data)
frame = decode_frame(reader)
return frame, len(data) - reader.offset
def pop_frame(data):
frame, l = read_frame(data)
class BufferedReader(object):
def __init__(self, data):
self.data = data
self.offset = 0
def recvn(self, n):
assert len(self.data) - self.offset >= n
self.offset += n
return self.data[self.offset - n:self.offset]
class BlockingSocket(object):
def __init__(self, sock):
self.sock = sock
def recvn(self, n):
"""
Keep receiving data from `sock` until exactly `n` bytes have been read.
Keep receiving data until exactly `n` bytes have been read.
"""
data = ''
while len(data) < n:
received = sock.recv(n - len(data))
received = self.sock.recv(n - len(data))
if not len(received):
raise socket.error('no data read from socket')
......@@ -246,6 +271,32 @@ def recvn(sock, n):
return data
def contains_frame(data):
"""
Read the frame length from the start of `data` and check if the data is
long enough to contain the entire frame.
"""
if len(data) < 2:
return False
b2 = struct.unpack('!B', data[1])
payload_len = b2 & 0x7F
payload_start = 2
if payload_len == 126:
if len(data) > 4:
payload_len = struct.unpack('!H', data[2:4])
payload_start = 4
elif payload_len == 127:
if len(data) > 12:
payload_len = struct.unpack('!Q', data[4:12])
payload_start = 12
return len(data) >= payload_len + payload_start
def mask(key, original):
"""
Mask an octet string using the given masking key.
......
import socket
import ssl
from frame import receive_frame
from frame import receive_frame, pop_frame, contains_frame
from handshake import ServerHandshake, ClientHandshake
from errors import SSLError
......@@ -11,6 +11,11 @@ INHERITED_ATTRS = ['bind', 'close', 'listen', 'fileno', 'getpeername',
'settimeout', 'gettimeout', 'shutdown', 'family', 'type',
'proto']
STATE_INIT = 0
STATE_READ = 1
STATE_WRITE = 2
STATE_CLOSE = 4
class websocket(object):
"""
......@@ -88,6 +93,11 @@ class websocket(object):
self.hooks_send = []
self.hooks_recv = []
self.state = STATE_INIT
self.sendbuf = ''
self.recvbuf = ''
self.recv_callbacks = []
self.sock = sock or socket.socket(sfamily, socket.SOCK_STREAM, sproto)
def __getattr__(self, name):
......@@ -153,6 +163,62 @@ class websocket(object):
"""
return [self.recv() for i in xrange(n)]
def queue_send(self, frame):
"""
Enqueue `frame` to the send buffer so that it is send on the next
`do_async_send`.
"""
for hook in self.hooks_send:
frame = hook(frame)
self.sendbuf += frame.pack()
self.state |= STATE_WRITE
def queue_recv(self, callback):
"""
Enqueue `callback` to be called when the next frame is recieved by
`do_async_recv`.
"""
self.recv_callbacks.push(callback)
self.state |= STATE_READ
def queue_close(self):
self.state |= STATE_CLOSE
def do_async_send(self):
"""
Send any queued data. If all data is sent, STATE_WRITE is removed from
the state mask.
"""
assert self.state | STATE_WRITE
assert len(self.sendbuf)
nwritten = self.sock.send(self.sendbuf)
self.sendbuf = self.sendbuf[nwritten:]
if len(self.sendbuf) == 0:
self.state ^= STATE_WRITE
def do_async_recv(self, bufsize):
"""
"""
assert self.state | STATE_READ
self.recvbuf += self.sock.recv(bufsize)
while contains_frame(self.recvbuf):
frame, self.recvbuf = pop_frame(self.recvbuf)
if len(self.recv_callbacks) == 0:
raise IndexError('no callback installed for received frame %s'
% frame)
cb = self.recv_callbacks.pop(0)
cb(frame)
if len(self.recvbuf) == 0:
self.state ^= STATE_READ
def enable_ssl(self, *args, **kwargs):
"""
Transforms the regular socket.socket to an ssl.SSLSocket for secure
......@@ -180,6 +246,7 @@ class websocket(object):
being sent and removes the instance for received data. This way, data
can be sent and received as if on a regular socket.
>>> import wspy
>>> sock = wspy.websocket()
>>> sock.add_hook(lambda data: tswpy.Frame(tswpy.OPCODE_TEXT, data),
>>> lambda frame: frame.payload)
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment