Significantly updated web socket API, see README for details on the new API

README.md

@@ -1,10 +1,14 @@
 **twspy** is a standalone implementation of web sockets for Python, defined by
 [RFC 6455](http://tools.ietf.org/html/rfc6455).
 
-- The websocket.WebSocket class upgrades a regular socket to a web socket.
-- message.py contains classes that abstract messages sent over the socket.
-  Sent messages are automatically converted to frames, and received frames are
-  converted to messages. Fragmented messages are also supported.
-- The server.Server class can be used to support multiple clients to open a
-  web socket simultaneously in different threads, which is often desirable in
-  web-based applications.
+- The websocket class upgrades a regular socket to a web socket. A websocket
+  instance is a single end point of a connection. A websocket instance sends
+  and receives frames (Frame instances) as opposed to bytes (which is received
+  in a regular socket).
+
+- A Connection instance represents a connection between two end points, based
+  on a websocket instance. A connection handles control frames properly, and
+  sends/receives *messages* (which are higher-level than frames). Messages are
+  automatically converted to frames, and received frames are converted to
+  messages. Fragmented messages (messages consisting of multiple frames) are
+  also supported.

connection.py

@@ -0,0 +1,194 @@
+import struct
+
+from frame import ControlFrame, OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG
+from message import create_message
+from exceptions import SocketClosed, PingError
+
+
+class Connection(object):
+    """
+    A Connection uses a websocket instance to send and receive (optionally
+    fragmented) messages, which are Message instances. Control frames are
+    handled automatically in the way specified by RFC 6455.
+
+    To use the Connection class, it should be extended and the exxtending class
+    should implement the on*() handlers.
+    """
+    def __init__(self, sock):
+        """
+        `sock` is a websocket instance which has completed its handshake.
+        """
+        self.sock = sock
+
+        self.received_close_params = None
+        self.close_frame_sent = False
+
+        self.ping_sent = False
+        self.ping_payload = None
+
+        self.onopen()
+
+    def send(self, message, fragment_size=None):
+        """
+        Send a message. If `fragment_size` is specified, the message is
+        fragmented into multiple frames whose payload size does not extend
+        `fragment_size`.
+        """
+        if fragment_size is None:
+            self.sock.send(message.frame())
+        else:
+            self.sock.send(*message.fragment(fragment_size))
+
+    def receive(self):
+        """
+        Receive a message. A message may consist of multiple (ordered) data
+        frames. A control frame may be delivered at any time, also when
+        expecting the next data frame of a fragmented message. These control
+        frames are handled immediately bu handle_control_frame().
+        """
+        fragments = []
+
+        while not len(fragments) or not fragments[-1].final:
+            frame = self.sock.recv()
+
+            if isinstance(frame, ControlFrame):
+                self.handle_control_frame(frame)
+
+                # No more receiving data after a close message
+                if frame.opcode == OPCODE_CLOSE:
+                    break
+            else:
+                fragments.append(frame)
+
+        payload = ''.join([f.payload for f in fragments])
+        return create_message(fragments[0].opcode, payload)
+
+    def handle_control_frame(self, frame):
+        """
+        Handle a control frame as defined by RFC 6455.
+        """
+        if frame.opcode == OPCODE_CLOSE:
+            # Set parameters and keep receiving the current fragmented frame
+            # chain, assuming that the CLOSE frame will be handled by
+            # handle_close() as soon as possible
+            self.received_close_params = frame.unpack_close()
+
+        elif frame.opcode == OPCODE_PING:
+            # Respond with a pong message with identical payload
+            self.sock.send(ControlFrame(OPCODE_PONG, frame.payload))
+
+        elif frame.opcode == OPCODE_PONG:
+            # Assert that the PONG payload is identical to that of the PING
+            if not self.ping_sent:
+                raise PingError('received PONG while no PING was sent')
+
+            self.ping_sent = False
+
+            if frame.payload != self.ping_payload:
+                raise PingError('received PONG with invalid payload')
+
+            self.ping_payload = None
+            self.onpong(frame.payload)
+
+    def receive_forever(self):
+        """
+        Receive and handle messages in an endless loop. A message may consist
+        of multiple data frames, but this is not visible for onmessage().
+        Control messages (or control frames) are handled automatically.
+        """
+        while True:
+            try:
+                self.onmessage(self, self.receive())
+
+                if self.received_close_params is not None:
+                    self.handle_close(*self.received_close_params)
+                    break
+            except SocketClosed:
+                self.onclose(None, '')
+                break
+            except Exception as e:
+                self.onexception(e)
+
+    def send_close(self, code, reason):
+        """
+        Send a CLOSE control frame.
+        """
+        payload = '' if code is None else struct.pack('!H', code) + reason
+        self.sock.send(ControlFrame(OPCODE_CLOSE, payload))
+        self.close_frame_sent = True
+
+    def send_ping(self, payload=''):
+        """
+        Send a PING control frame with an optional payload.
+        """
+        self.sock.send(ControlFrame(OPCODE_PING, payload))
+        self.ping_payload = payload
+        self.ping_sent = True
+        self.onping(payload)
+
+    def handle_close(self, code=None, reason=''):
+        """
+        Handle a close message by sending a response close message if no CLOSE
+        frame was sent before, and closing the connection. The onclose()
+        handler is called afterwards.
+        """
+        if not self.close_frame_sent:
+            payload = '' if code is None else struct.pack('!H', code)
+            self.sock.send(ControlFrame(OPCODE_CLOSE, payload))
+
+        self.sock.close()
+        self.onclose(code, reason)
+
+    def close(self, code=None, reason=''):
+        """
+        Close the socket by sending a CLOSE frame and waiting for a response
+        close message. The onclose() handler is called after the CLOSE frame
+        has been sent, but before the response has been received.
+        """
+        self.send_close(code, reason)
+        # FIXME: swap the two lines below?
+        self.onclose(code, reason)
+        frame = self.sock.recv()
+        self.sock.close()
+
+        if frame.opcode != OPCODE_CLOSE:
+            raise ValueError('expected CLOSE frame, got %s instead' % frame)
+
+    def onopen(self):
+        """
+        Called after the connection is initialized.
+        """
+        pass
+
+    def onmessage(self, message):
+        """
+        Called when a message is received. `message` is a Message object, which
+        can be constructed from a single frame or multiple fragmented frames.
+        """
+        return NotImplemented
+
+    def onping(self, payload):
+        """
+        Called after a PING control frame has been sent. This handler could be
+        used to start a timeout handler for a PONG frame that is not received
+        in time.
+        """
+        pass
+
+    def onpong(self, payload):
+        """
+        Called when a PONG control frame is received.
+        """
+        pass
+
+    def onclose(self, code, reason):
+        """
+        Called when the socket is closed by either end point.
+        """
+        pass
+
+    def onexception(self, e):
+        """
+        Handle a raised exception.
+        """
+        pass

frame.py

@@ -26,8 +26,7 @@ class Frame(object):
     """
     A Frame instance represents a web socket data frame as defined in RFC 6455.
     To encoding a frame for sending it over a socket, use Frame.pack(). To
-    receive and decode a frame from a socket, use receive_frame() (or,
-    preferably, receive_fragments()).
+    receive and decode a frame from a socket, use receive_frame().
     """
     def __init__(self, opcode, payload, masking_key='', final=True, rsv1=False,
             rsv2=False, rsv3=False):
@@ -101,8 +100,14 @@ class Frame(object):
 
     def fragment(self, fragment_size, mask=False):
         """
-        Fragment the frame into a chain of fragment frames, as explained in the
-        docs of the function receive_fragments().
+        Fragment the frame into a chain of fragment frames:
+        - An initial frame with non-zero opcode
+        - Zero or more frames with opcode = 0 and final = False
+        - A final frame with opcode = 0 and final = True
+
+        The first and last frame may be the same frame, having a non-zero
+        opcode and final = True. Thus, this function returns a list containing
+        at least a single frame.
 
         `fragment_size` indicates the maximum payload size of each fragment.
         The payload of the original frame is split into one or more parts, and
@@ -171,38 +176,6 @@ class ControlFrame(Frame):
         return code, reason
 
 
-def receive_fragments(sock, control_frame_handler):
-    """
-    Receive a sequence of frames that belong together on socket `sock`:
-    - An initial frame with non-zero opcode
-    - Zero or more frames with opcode = 0 and final = False
-    - A final frame with opcode = 0 and final = True
-
-    The first and last frame may be the same frame, having a non-zero opcode
-    and final = True. Thus, this function returns a list of at least a single
-    frame.
-
-    `control_frame_handler` is a callback function taking a single argument,
-    which is a ControlFrame instance in case a control frame is received (this
-    may occur in the middle of a fragment chain).
-    """
-    fragments = []
-
-    while not len(fragments) or not fragments[-1].final:
-        frame = receive_frame(sock)
-
-        if isinstance(frame, ControlFrame):
-            control_frame_handler(frame)
-
-            # No more receiving data after a close message
-            if frame.opcode == OPCODE_CLOSE:
-                break
-        else:
-            fragments.append(frame)
-
-    return fragments
-
-
 def receive_frame(sock):
     """
     Receive a single frame on socket `sock`. The frame schme is explained in

websocket.py

@@ -1,69 +1,86 @@
 import re
-import struct
 import socket
 from hashlib import sha1
 
-from frame import ControlFrame, receive_fragments, receive_frame, \
-        OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG
-from message import create_message
-from exceptions import InvalidRequest, SocketClosed, PingError
+from frame import receive_frame
+from exceptions import InvalidRequest
 
 
 WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
 WS_VERSION = '13'
 
 
-class WebSocket(object):
+class websocket(object):
     """
-    A WebSocket upgrades a regular TCP socket to a web socket. The class
-    implements the handshake protocol as defined by RFC 6455, provides
-    abstracted methods for sending (optionally fragmented) messages, and
-    automatically handles control messages.
+    Implementation of web socket, upgrades a regular TCP socket to a websocket
+    using the HTTP handshakes and frame (un)packing, as specified by RFC 6455.
+
+    Server example:
+    >>> sock = websocket()
+    >>> sock.bind(('', 80))
+    >>> sock.listen()
+
+    >>> client = sock.accept()
+    >>> client.send(Frame(...))
+    >>> frame = client.recv()
+
+    Client example:
+    >>> sock = websocket()
+    >>> sock.connect(('kompiler.org', 80))
     """
-    def __init__(self, sock):
+    def __init__(self, wsprotocols=[], family=socket.AF_INET, proto=0):
         """
-        `sock` is a regular TCP socket instance.
+        Create aregular TCP socket of family `family` and protocol
+        `wsprotocols` is a list of supported protocol names.
         """
-        self.sock = sock
+        self.sock = socket.socket(family, socket.SOCK_STREAM, proto)
+        self.protocols = wsprotocols
 
-        self.received_close_params = None
-        self.close_frame_sent = False
+    def bind(self, address):
+        self.sock.bind(address)
 
-        self.ping_sent = False
-        self.ping_payload = None
+    def listen(self, backlog):
+        self.sock.listen(backlog)
 
-    def send_message(self, message, fragment_size=None):
-        if fragment_size is None:
-            self.send_frame(message.frame())
-        else:
-            map(self.send_frame, message.fragment(fragment_size))
+    def accept(self):
+        client, address = socket.socket.accept(self)
+        client = websocket(client)
+        client.server_handshake()
+        return client, address
 
-    def send_frame(self, frame):
-        self.sock.sendall(frame.pack())
+    def connect(self, address):
+        """
+        Equivalent to socket.connect(), but sends an HTTP handshake request
+        after connecting.
+        """
+        self.sock.sonnect(address)
+        self.client_handshake()
 
-    def handle_control_frame(self, frame):
-        if frame.opcode == OPCODE_CLOSE:
-            self.received_close_params = frame.unpack_close()
-        elif frame.opcode == OPCODE_PING:
-            # Respond with a pong message with identical payload
-            self.send_frame(ControlFrame(OPCODE_PONG, frame.payload))
-        elif frame.opcode == OPCODE_PONG:
-            # Assert that the PONG payload is identical to that of the PING
-            if not self.ping_sent:
-                raise PingError('received PONG while no PING was sent')
+    def send(self, *args):
+        """
+        Send a number of frames.
+        """
+        for frame in args:
+            self.sock.sendall(frame.pack())
 
-            self.ping_sent = False
+    def recv(self, n=1):
+        """
+        Receive exactly `n` frames. These can be either data frames or control
+        frames, or a combination of both.
+        """
+        return [receive_frame(self.sock) for i in xrange(n)]
 
-            if frame.payload != self.ping_payload:
-                raise PingError('received PONG with invalid payload')
+    def getpeername(self):
+        return self.sock.getpeername()
 
-            self.ping_payload = None
-            self.onpong(frame.payload)
+    def getsockname(self):
+        return self.sock.getpeername()
 
-    def receive_message(self):
-        frames = receive_fragments(self.sock, self.handle_control_frame)
-        payload = ''.join([f.payload for f in frames])
-        return create_message(frames[0].opcode, payload)
+    def setsockopt(self, level, optname, value):
+        self.sock.setsockopt(level, optname, value)
+
+    def getsockopt(self, level, optname):
+        return self.sock.getsockopt(level, optname)
 
     def server_handshake(self):
         """
@@ -120,168 +137,6 @@ class WebSocket(object):
 
         self.sock.send(shake + '\r\n')
 
-        self.onopen()
-
-    def receive_forever(self):
-        """
-        Receive and handle messages in an endless loop. A message may consist
-        of multiple data frames, but this is not visible for onmessage().
-        Control messages (or control frames) are handled automatically.
-        """
-        while True:
-            try:
-                self.onmessage(self, self.receive_message())
-
-                if self.received_close_params is not None:
-                    self.handle_close(*self.received_close_params)
-                    break
-            except SocketClosed:
-                self.onclose(None, '')
-                break
-            except Exception as e:
-                self.onexception(e)
-
-    def send_close(self, code, reason):
-        """
-        Send a close control frame.
-        """
-        payload = '' if code is None else struct.pack('!H', code) + reason
-        self.send_frame(ControlFrame(OPCODE_CLOSE, payload))
-        self.close_frame_sent = True
-
-    def send_ping(self, payload=''):
-        """
-        Send a ping control frame with an optional payload.
-        """
-        self.send_frame(ControlFrame(OPCODE_PING, payload))
-        self.ping_payload = payload
-        self.ping_sent = True
-        self.onping(payload)
-
-    def handle_close(self, code=None, reason=''):
-        """
-        Handle a close message by sending a response close message if no close
-        message was sent before, and closing the connection. The onclose()
-        handler is called afterwards.
-        """
-        if not self.close_frame_sent:
-            payload = '' if code is None else struct.pack('!H', code)
-            self.send_frame(ControlFrame(OPCODE_CLOSE, payload))
-
-        self.sock.close()
-        self.onclose(code, reason)
-
-    def close(self, code=None, reason=''):
-        """
-        Close the socket by sending a close message and waiting for a response
-        close message. The onclose() handler is called after the close message
-        has been sent, but before the response has been received.
-        """
-        self.send_close(code, reason)
-        # FIXME: swap the two lines below?
-        self.onclose(code, reason)
-        frame = receive_frame(self.sock)
-        self.sock.close()
-
-        if frame.opcode != OPCODE_CLOSE:
-            raise ValueError('expected close frame, got %s instead' % frame)
-
-    def onopen(self):
-        """
-        Called after the handshake has completed.
-        """
-        pass
-
-    def onmessage(self, message):
-        """
-        Called when a message is received. `message` is a Message object, which
-        can be constructed from a single frame or multiple fragmented frames.
-        """
-        return NotImplemented
-
-    def onping(self, payload):
-        """
-        Called after a ping control frame has been sent. This handler could be
-        used to start a timeout handler for a pong message that is not received
-        in time.
-        """
-        pass
-
-    def onpong(self, payload):
-        """
-        Called when a pong control frame is received.
-        """
-        pass
-
-    def onclose(self, code, reason):
-        """
-        Called when the socket is closed by either end point.
-        """
-        pass
-
-    def onexception(self, e):
-        """
-        Handle a raised exception.
-        """
-        pass
-
-
-class websocket(WebSocket):
-    """
-    Alternative implementation of web socket, extending the regular socket
-    object.
-    """
-    def __init__(self, family=socket.AF_INET, proto=0):
-        sock = socket.socket(family, socket.SOCK_STREAM, proto)
-        WebSocket.__init__(self, sock)
-
-    def bind(self, address):
-        self.sock.bind(address)
-
-    def listen(self, backlog):
-        self.sock.listen(backlog)
-
-    def accept(self):
-        client, address = socket.socket.accept(self)
-        client = websocket(client)
-        client.handshake()
-        return client, address
-
-    def recv(self):
-        """
-        Receive a sinfle frame.
-        """
-        return receive_frame(self.sock)
-
-    def send(self, frame):
-        """
-        Send a single frame.
-        """
-        self.send_frame(frame)
-
-    def sendall(self, frames):
-        """
-        Send a list of frames.
-        """
-        for frame in frames:
-            self.send(frame)
-
-    def getpeername(self):
-        return self.sock.getpeername()
-
-    def getsockname(self):
-        return self.sock.getpeername()
-
-    def setsockopt(self, level, optname, value):
-        self.sock.setsockopt(level, optname, value)
-
-    def getsockopt(self, level, optname):
-        return self.sock.getsockopt(level, optname)
-
-
-if __name__ == '__main__':
-    sock = websocket()
-    sock.bind(('', 80))
-    sock.listen()
-
-    client = sock.accept()
+    def client_handshake(self):
+        # TODO: implement HTTP request headers for client handshake
+        raise NotImplementedError()