Jeffrey Gelens avatar Jeffrey Gelens committed 586413b

Added some consistency checks to receiving

Comments (0)

Files changed (2)

 This Websocket library for Gevent is written and maintained by
 
   Jeffrey Gelens <jeffrey at noppo.pro>
+
+
+Contributors:
+
+  Denis Bilenko
+  Lon Ingram

geventwebsocket/websocket.py

     pass
 
 
-
 class ProtocolException(Exception):
     pass
 
         self.protocol = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL', 'unknown')
         self.path = environ.get('PATH_INFO')
         self.websocket_closed = False
-        self.compatibility_mode = compatibility_mode
-        self._fragments = []
-        self._original_opcode = -1
+        self._chunks = bytearray()
+        self._first_opcode = None
 
     def _read_from_socket(self, count):
         return self.rfile.read(count)
 
-    def wait(self):
+    def receive(self):
         """Return the next frame from the socket
 
         If the next frame is invalid, wait closes the socket and returns None.
             if opcode > 0x7 and fin == 0:
                 raise ProtocolException('Control frames cannot be fragmented')
 
-            #if len(self._fragments) > 0 and not is_fin and opcode != self.OPCODE_FRAG:
-            #    self.close(self.REASON_PROTOCOL_ERROR,
-            #            'Received new fragment frame with non-zero opcode')
+            if len(self._chunks) > 0 and \
+                    fin == 0 and opcode != self.OPCODE_CONTINUATION:
+                self.close(self.REASON_PROTOCOL_ERROR,
+                        'Received new fragment frame with non-zero opcode')
+                return
 
-            #if len(self._fragments) > 0 and is_fin and (
-            #        self.OPCODE_TEXT <= opcode <= self.OPCODE_BINARY):
-            #    self.close(self.REASON_PROTOCOL_ERROR,
-            #            'Received new unfragmented data frame during fragmented message')
+            if len(self._chunks) > 0 and \
+                    fin == 1 and (self.OPCODE_TEXT <= opcode <= self.OPCODE_BINARY):
+                self.close(self.REASON_PROTOCOL_ERROR,
+                        'Received new unfragmented data frame during fragmented message')
 
             mask = (second_byte >> 7) & 1
             payload_length = (second_byte) & 0x7f
 
-            #if not self.MASK & length_octet:
+            #if not self.MASK & length_octet: # TODO: where is this in the docs?
             #    self.close(self.REASON_PROTOCOL_ERROR, 'MASK must be set')
 
-            #length_code = length_octet & self.PAYLOAD
-
             # Control frames MUST have a payload length of 125 bytes or less
             if opcode > 0x7 and payload_length > 125:
                 raise FrameTooLargeException("Control frame payload cannot be larger than 125 bytes")
 
             payload = ""
 
+            # Unmask the payload if necessary
             if mask:
                 masking_key = struct.unpack('!BBBB', self._read_from_socket(4))
                 masked_payload = self._read_from_socket(length)
 
                 payload = masked_payload
 
-
+            # Read application data
             if opcode == self.OPCODE_TEXT:
-                self._fragments.append(payload.decode("utf-8", "replace"))
+                self._first_opcode = opcode
+                self._chunks.extend(payload.decode("utf-8", "replace"))
 
             elif opcode == self.OPCODE_BINARY:
-                self._fragments.append(payload)
+                self._first_opcode = opcode
+                self._chunks.extend(payload)
 
             elif opcode == self.OPCODE_CONTINUATION:
-                if len(self._fragments) != 0:
+                if len(self._chunks) != 0:
                     raise ProtocolException("Cannot continue a non started message")
 
-                if opcode == self.OPCODE_TEXT:
-                    self._fragments.append(payload.decode("utf-8", "replace"))
+                if self._first_opcode == self.OPCODE_TEXT:
+                    self._chunks.extend(payload.decode("utf-8", "replace"))
                 else:
-                    self._fragments.append(payload)
+                    self._chunks.extend(payload)
 
             elif opcode == self.OPCODE_CLOSE:
                 if length >= 2:
                     return (self.OPCODE_PONG, payload)
                 else:
                     continue
+            else:
+                raise Exception("Shouldn't happen")
 
             if fin == 1:
-                if len(self._fragments) > 0:
-                    msg = ''.join(self._fragments)
-                    self._fragments = []
-                if not self.compatibility_mode:
-                    return (opcode, msg)
-                else:
-                    return msg
+                msg = ''.join(self._chunks)
+
+                self._first_opcode = False
+                self._chunks = bytearray()
+
+                return msg
+
 
     def _encode_text(self, s):
         if isinstance(s, unicode):
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.