Anonymous avatar Anonymous committed c073c4d

Implement send and close and beginnings of wait

Send doesn't support fragmenting, but does support all opcode types and checks
for invalid ones and for invalid encoding. Close sends a close frame.

Comments (0)

Files changed (2)

geventwebsocket/websocket.py

 # The version as of August 6, 2010 will be implementend once Firefox or
 # Webkit-trunk support this version.
 
+import binascii
+import struct
+
 class WebSocket(object):
     def __init__(self, sock, rfile, environ):
         self.rfile = rfile
                 raise IOError("Reveiced an invalid message")
 
 class WebSocketVersion7(WebSocket):
+    FIN = int("10000000", 2)
+    RSV = int("01110000", 2)
+    OPCODE = int("00001111", 2)
+    MASK = int("10000000", 2)
+    PAYLOAD = int("01111111", 2)
+
+    OPCODE_TEXT = 0x1
+    OPCODE_BINARY = 0x2
+    OPCODE_CLOSE = 0x8
+    OPCODE_PING = 0x9
+    OPCODE_PONG = 0xA
+
+    REASON_NORMAL = 1000
+    REASON_GOING_AWAY = 1001
+
+    LEN_16 = 126
+    LEN_64 = 127
+
     def __init__(self, sock, rfile, environ):
         self.rfile = rfile
         self.socket = sock
         self.protocol = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL', 'unknown')
         self.path = environ.get('PATH_INFO')
         self.websocket_closed = False
+
+    def wait(self):
+        msg = ""
+        while True:
+            if self.websocket_closed:
+                return None
+
+            opcode, length = struct.unpack('!BB', self.rfile.read(2))
+
+            if self.RSV & opcode:
+                self.close(1002, 'Reserved bits cannot be set')
+                return None
+
+            is_final_frag = (self.FIN & opcode) != 0
+
+    def _encodeText(self, s):
+        if isinstance(s, unicode):
+            return s.encode('utf-8')
+        elif isinstance(s, str):
+            return unicode(s).encode('utf-8')
+        else:
+            raise Exception('Invalid encoding')
+
+    def send(self, opcode, message):
+        if self.websocket_closed:
+            raise Exception('Connection was terminated')
+
+        if opcode < self.OPCODE_TEXT or (opcode > self.OPCODE_BINARY and 
+                opcode < self.OPCODE_CLOSE) or opcode > self.OPCODE_PONG:
+            raise Exception('Invalid opcode %d' % opcode)
+
+        if opcode == self.OPCODE_TEXT:
+            message = self._encodeText(message)
+
+        length = len(message)
+
+        if opcode == self.OPCODE_TEXT:
+            message = struct.pack('!%ds' % length, message)
+
+        if length < 126:
+            preamble = struct.pack('!BB', self.FIN | opcode, length)
+        elif length < 2 ** 16:
+            preamble = struct.pack('!BBH', self.FIN | opcode, self.LEN_16, length)
+        else:
+            preamble = struct.pack('!BBQ', self.FIN | opcode, self.LEN_64, length)
+
+        self.socket.sendall(preamble + message)
+
+    def close(self, reason, message):
+        message = self._encodeText(message)
+        self.send(self.OPCODE_CLOSE, struct.pack('!H%ds' % len(message), reason, message))
+        self.websocket_closed = True
+
+        # based on gevent/pywsgi.py
+        # see http://pypi.python.org/pypi/gevent#downloads
+        if self.socket is not None:
+            try:
+                self.socket._sock.close()
+                self.socket.close()
+            except socket.error:
+                pass

tests/test__websocket.py

 from gevent import monkey
 monkey.patch_all(thread=False)
 
+import binascii
 import base64
+import struct
 import sys
 import greentest
 import gevent
 from gevent import socket
 from geventwebsocket.handler import WebSocketHandler
+from geventwebsocket.websocket import WebSocketVersion7
 
 
 CONTENT_LENGTH = 'Content-Length'
         fd.close()
 
 class TestWebSocketVersion7(TestCase):
+
+    GOOD_HEADERS = "" \
+        "GET /echo HTTP/1.1\r\n" \
+        "Host: localhost\r\n" \
+        "Upgrade: WebSocket\r\n" \
+        "Connection: Upgrade\r\n" \
+        "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" \
+        "Sec-WebSocket-Origin: http://localhost\r\n" \
+        "Sec-WebSocket-Protocol: chat, superchat\r\n" \
+        "Sec-WebSocket-Version: 7\r\n" \
+        "\r\n"
+
     def application(self, environ, start_response):
         if environ['PATH_INFO'] == "/echo":
             try:
                 ws = environ['wsgi.websocket']
+                self.ws = ws
             except KeyError:
                 start_response("400 Bad Request", [])
                 return []
 
+            """
             while True:
                 message = ws.wait()
                 if message is None:
                     break
                 ws.send(message)
+            """
 
             return []
 
 
     def test_good_handshake(self):
         fd = self.connect().makefile(bufsize=1)
-        headers = "" \
-            "GET /echo HTTP/1.1\r\n" \
-            "Host: localhost\r\n" \
-            "Upgrade: WebSocket\r\n" \
-            "Connection: Upgrade\r\n" \
-            "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" \
-            "Sec-WebSocket-Origin: http://localhost\r\n" \
-            "Sec-WebSocket-Protocol: chat, superchat\r\n" \
-            "Sec-WebSocket-Version: 7\r\n" \
-            "\r\n"
 
-        fd.write(headers)
+        fd.write(self.GOOD_HEADERS)
         response = read_http(fd, code=101, reason="Switching Protocols")
         response.assertHeader("Upgrade", "websocket")
         response.assertHeader("Connection", "Upgrade")
 
         fd.close();
 
+    def test_send_short_frame(self):
+        fd = self.connect().makefile(bufsize=1)
 
-    """
-    def test_handshake(self):
-        fd = self.connect().makefile(bufsize=1)
-        headers = "" \
-            "GET /echo HTTP/1.1\r\n" \
-            "Host: localhost\r\n" \
-            "Upgrade: WebSocket\r\n" \
-            "Connection: Upgrade\r\n" \
-            "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" \
-            "Sec-WebSocket-Origin: http://localhost\r\n" \
-            "Sec-WebSocket-Protocol: chat, superchat\r\n" \
-            "Sec-WebSocket-Version: 7\r\n" \
-            "\r\n"
+        fd.write(self.GOOD_HEADERS)
+        read_http(fd, code=101, reason="Switching Protocols")
 
-        fd.write(headers)
-        response = read_http(fd, code=101, reason="Web Socket Protocol Handshake")
+        msg = 'Hello, websocket'
+        self.ws.send(1, msg)
 
-        fd.write(self.message)
-        message = fd.read(len(self.message))
-        assert message == self.message, \
-               'Unexpected message: %r (expected %r)\n%s' % (message, self.message, self)
+        preamble = fd.read(2)
+        opcode, length = struct.unpack('!BB', preamble)
+
+        assert opcode & WebSocketVersion7.FIN, 'FIN must be set'
+        assert (opcode & WebSocketVersion7.OPCODE) == 1, 'Opcode must be 0x1'
+        assert (length & WebSocketVersion7.MASK) == 0, 'MASK must not be set'
+        assert length == len(msg), 'Wrong length %d, expected %d' % (length, len(msg))
+
+        rxd_msg = fd.read(length).decode('utf-8', 'replace')
+        assert rxd_msg == msg, 'Wrong message "%s"' % rxd_msg
 
         fd.close()
-    """
+
+    def test_send_medium_frame(self):
+        fd = self.connect().makefile(bufsize=1)
+
+        fd.write(self.GOOD_HEADERS)
+        read_http(fd, code=101, reason="Switching Protocols")
+
+        msg = 'Hello, websocket' * 8
+        self.ws.send(1, msg)
+
+        preamble = fd.read(4)
+        opcode, length_code, length = struct.unpack('!BBH', preamble)
+
+        assert opcode & WebSocketVersion7.FIN, 'FIN must be set'
+        assert (opcode & WebSocketVersion7.OPCODE) == 1, 'Opcode must be 0x1'
+        assert (length_code & WebSocketVersion7.MASK) == 0, 'MASK must not be set'
+        assert length_code == 126, 'The length code must be 126'
+        assert length == len(msg), 'Wrong length %d, expected %d' % (length, len(msg))
+
+        rxd_msg = fd.read(length).decode('utf-8', 'replace')
+        assert rxd_msg == msg, 'Wrong message "%s"' % rxd_msg
+
+        fd.close()
+
+    def test_send_long_frame(self):
+        fd = self.connect().makefile(bufsize=1)
+
+        fd.write(self.GOOD_HEADERS)
+        read_http(fd, code=101, reason="Switching Protocols")
+
+        msg = 'Hello, websocket' * 4097
+        self.ws.send(1, msg)
+
+        preamble = fd.read(10)
+        opcode, length_code, length = struct.unpack('!BBQ', preamble)
+
+        assert opcode & WebSocketVersion7.FIN, 'FIN must be set'
+        assert (opcode & WebSocketVersion7.OPCODE) == 1, 'Opcode must be 0x1'
+        assert (length_code & WebSocketVersion7.MASK) == 0, 'MASK must not be set'
+        assert length_code == 127, 'The length code must be 127'
+        assert length == len(msg), 'Wrong length %d, expected %d' % (length, len(msg))
+
+        rxd_msg = fd.read(length).decode('utf-8', 'replace')
+        assert rxd_msg == msg, 'Wrong message "%s"' % rxd_msg
+
+        fd.close()
+
+    def test_binary_frame(self):
+        fd = self.connect().makefile(bufsize=1)
+
+        fd.write(self.GOOD_HEADERS)
+        read_http(fd, code=101, reason="Switching Protocols")
+
+        msg = struct.pack('!BHB', 129, 23, 42)
+        self.ws.send(2, msg)
+
+        frame = fd.read(6)
+        opcode, length, first, second, third = struct.unpack('!BBBHB', frame)
+
+        assert opcode & WebSocketVersion7.FIN, 'FIN must be set'
+        assert (opcode & WebSocketVersion7.OPCODE) == 2, 'Opcode must be 0x2'
+        assert (length & WebSocketVersion7.MASK) == 0, 'MASK must not be set'
+        assert length == 4, 'Wrong length %d, expected 4' % length
+        assert first == 129, 'Expected first value to be 129, but got %d' % first
+        assert second == 23, 'Expected second value to be 23, but got %d' % second
+        assert third == 42, 'Expected third value to be 42, but got %d' % third
+
+        fd.close()
+
+    def test_wait_bad_framing_reserved_bits(self):
+        fd = self.connect().makefile(bufsize=1)
+
+        fd.write(self.GOOD_HEADERS)
+        read_http(fd, code=101, reason="Switching Protocols")
+
+        expected_msg = 'Reserved bits cannot be set'
+
+        fd.write(struct.pack("!BB", int("11000001", 2), int("10000000", 2)))
+
+        frame = self.ws.wait()
+        assert self.ws.websocket_closed, "Failed to close connection when sent a frame with RSV1 set"
+
+        preamble = fd.read(2)
+
+        opcode, length = struct.unpack('!BB', preamble)
+        assert opcode & WebSocketVersion7.FIN, 'FIN must be set'
+        assert (opcode & WebSocketVersion7.OPCODE) == 8, 'Opcode must be 0x8'
+        assert (length & WebSocketVersion7.MASK) == 0, 'MASK must not be set'
+
+        reason = fd.read(2)
+        reason = struct.unpack('!H', reason)[0]
+        assert reason == 1002, 'Expected reason to be 1002, but got %d' % reason
+
+        rxd_msg = fd.read(length - 2).decode('utf-8', 'replace')
+        assert rxd_msg == expected_msg, 'Wrong message "%s"' % rxd_msg
+
+        fd.close();
 
 if __name__ == '__main__':
     greentest.main()
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.