Anonymous avatar Anonymous committed c26e33d

Return tuple for all opcodes when not in compatibility mode

- Comment public methods of WebSocketVersion7

Comments (0)

Files changed (2)

geventwebsocket/websocket.py

         self.websocket_closed = False
         self.compatibility_mode = compatibility_mode
         self._fragments = []
+        self._original_opcode = -1
 
     def _read_from_socket(self, count):
         return self.rfile.read(count)
 
-    # TODO: replace all magic numbers with constants
     def wait(self):
+        """Return the next frame from the socket
+
+        If the next frame is invalid, wait closes the socket and returns None.
+        
+        If the next frame is valid and the websocket instance's
+        compatibility_mode attribute is True, then wait ignores PING and PONG
+        frames, returns None when sent a CLOSE frame and returns the payload
+        for data frames.
+        
+        If the next frame is valid and the websocket instance's
+        compatibility_mode attribute is False, it returns a tuple of the form
+        (opcode, payload).
+        """
+        
         while True:
             payload = ''
             if self.websocket_closed:
 
                 self.close(self.REASON_NORMAL, '')
                 if not self.compatibility_mode:
-                    return (reason, message)
+                    return (self.OPCODE_CLOSE, (reason, message))
                 else:
                     return None
-
-            if opcode == self.OPCODE_PING:
+            elif opcode == self.OPCODE_PING:
                 self.send(payload, opcode=self.OPCODE_PONG)
                 if not self.compatibility_mode:
                     return (self.OPCODE_PING, payload)
                     continue
 
             if is_final_frag:
-                payload = ''.join(self._fragments) + payload
-                self._fragments = []
-                return payload
+                if len(self._fragments) > 0:
+                    opcode = self._original_opcode
+                    self._original_opcode = -1
+                    payload = ''.join(self._fragments) + payload
+                    self._fragments = []
+                if not self.compatibility_mode:
+                    return (opcode, payload)
+                else:
+                    return payload
             else:
+                if len(self._fragments) == 0:
+                    self._original_opcode = opcode
                 self._fragments.append(payload)
 
     def _encode_text(self, s):
                 opcode < self.OPCODE_CLOSE) or opcode > self.OPCODE_PONG
 
     def send(self, message, opcode=OPCODE_TEXT):
+        """Send a frame over the websocket with message as its payload
+
+        Keyword args:
+        opcode -- the opcode to use (default OPCODE_TEXT)
+        """
+
         if self.websocket_closed:
             raise Exception('Connection was terminated')
 
         self.socket.sendall(preamble + message)
 
     def close(self, reason, message):
+        """Close the websocket, sending the specified reason and message
+        """
+
         message = self._encode_text(message)
         self.send(struct.pack('!H%ds' % len(message), reason, message), opcode=self.OPCODE_CLOSE)
         self.websocket_closed = True

tests/test__websocket.py

             WebSocketVersion7.FIN | WebSocketVersion7.OPCODE_TEXT, 
             WebSocketVersion7.MASK | length, mask, encoded_msg))
 
-        rxd_msg = self.ws.wait()
+        opcode, rxd_msg = self.ws.wait()
         assert not self.ws.websocket_closed, 'Closed connection when sent a good frame'
+        assert opcode == WebSocketVersion7.OPCODE_TEXT, 'Wrong opcode "%x"' % opcode
         assert rxd_msg == msg, 'Wrong message "%s"' % rxd_msg
 
         fd.close();
             WebSocketVersion7.FIN | WebSocketVersion7.OPCODE_TEXT, 
             WebSocketVersion7.MASK | 126, length, mask, encoded_msg))
 
-        rxd_msg = self.ws.wait()
+        opcode, rxd_msg = self.ws.wait()
         assert not self.ws.websocket_closed, 'Closed connection when sent a good frame'
+        assert opcode == WebSocketVersion7.OPCODE_TEXT, 'Wrong opcode "%x"' % opcode
         assert rxd_msg == msg, 'Wrong message "%s"' % rxd_msg
 
         fd.close();
         fd.write(payload)
 
         self.ws._waiter = 'test'
-        rxd_msg = self.ws.wait()
+        opcode, rxd_msg = self.ws.wait()
         assert not self.ws.websocket_closed, 'Closed connection when sent a good frame'
+        assert opcode == WebSocketVersion7.OPCODE_TEXT, 'Wrong opcode "%x"' % opcode
         assert rxd_msg == msg, 'Wrong message "%s"' % rxd_msg
 
         fd.close();
 
         rxd_msg = self.ws.wait()
         assert self.ws.websocket_closed, 'Did not close connection when sent a close frame'
-        assert rxd_msg == (reason, msg), 'Wrong result "%s"' % (rxd_msg,)
+        assert rxd_msg == (WebSocketVersion7.OPCODE_CLOSE, (reason, msg)), \
+                'Wrong result "%s"' % (rxd_msg,)
 
         preamble = fd.read(2)
         opcode, length = struct.unpack('!BB', preamble)
         rxd_msg = self.ws.wait()
         assert self.ws.websocket_closed, \
                 'Did not close connection when sent a close frame with no payload'
-        assert rxd_msg == (None, None), 'Wrong result "%s"' % (rxd_msg,)
+        assert rxd_msg == (WebSocketVersion7.OPCODE_CLOSE, (None, None)), \
+                'Wrong result "%s"' % (rxd_msg,)
 
         preamble = fd.read(2)
         opcode, length = struct.unpack('!BB', preamble)
             WebSocketVersion7.FIN | WebSocketVersion7.OPCODE_FRAG, 
             WebSocketVersion7.MASK | length, mask, encoded_msg))
 
-        rxd_msg = self.ws.wait()
+        opcode, rxd_msg = self.ws.wait()
         assert not self.ws.websocket_closed, 'Closed connection when sent a good frame'
+        assert opcode == WebSocketVersion7.OPCODE_TEXT, 'Wrong opcode "%x"' % opcode
         assert rxd_msg == expected_msg, 'Wrong message "%s"' % rxd_msg
 
         fd.close();
             WebSocketVersion7.FIN | WebSocketVersion7.OPCODE_FRAG, 
             WebSocketVersion7.MASK | length, mask, encoded_msg))
 
-        rxd_msg = self.ws.wait()
+        opcode, rxd_msg = self.ws.wait()
         assert not self.ws.websocket_closed, 'Closed connection when sent a good frame'
+        assert opcode == WebSocketVersion7.OPCODE_TEXT, 'Wrong opcode "%x"' % opcode
         assert rxd_msg == expected_msg, 'Wrong message "%s"' % rxd_msg
 
         fd.close();
         assert txd_msg == ping_msg.encode('utf-8'), \
                 'Wrong message "%s", expected "%s"' % (txd_msg, ping_msg)
 
-        rxd_msg = self.ws.wait()
+        opcode, rxd_msg = self.ws.wait()
         assert not self.ws.websocket_closed, 'Closed connection when sent a good frame'
+        assert opcode == WebSocketVersion7.OPCODE_TEXT, 'Wrong opcode "%x"' % opcode
         assert rxd_msg == expected_msg, \
                 'Wrong message "%s", expected "%s"' % (rxd_msg, expected_msg)
 
                 WebSocketVersion7.FIN | WebSocketVersion7.OPCODE_FRAG, 
                 WebSocketVersion7.MASK | length, mask, encoded_msg))
 
-            rxd_msg = self.ws.wait()
+            opcode, rxd_msg = self.ws.wait()
             assert not self.ws.websocket_closed, 'Closed connection when sent a good frame'
+            assert opcode == WebSocketVersion7.OPCODE_TEXT, 'Wrong opcode "%x"' % opcode
             assert rxd_msg == expected_msg, 'Wrong message "%s"' % rxd_msg
 
         fd.close();
+
     """
     TODO: write fragmentation test cases:
           - too large of a message
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.