fixed closing frame support

Lawouach · Lawouach · commit 3fcb166daa57 · 2011-12-15T19:05:53.000+01:00
diff --git a/ws4py/messaging.py b/ws4py/messaging.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import os
+import struct
 
 from ws4py.framing import Frame, OPCODE_CONTINUATION, OPCODE_TEXT, \
      OPCODE_BINARY, OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG
@@ -115,7 +116,13 @@ def is_text(self):
 
 class CloseControlMessage(Message):
     def __init__(self, code=1000, reason=''):
-        Message.__init__(self, OPCODE_CLOSE, reason)
+        data = ""
+        if code:
+            data += struct.pack("!H", code)
+        if reason:
+            data += reason.encode('utf-8')
+            
+        Message.__init__(self, OPCODE_CLOSE, data)
         self.code = code
         self.reason = reason
 
diff --git a/ws4py/server/handler/threadedhandler.py b/ws4py/server/handler/threadedhandler.py
@@ -215,17 +215,19 @@ def _receive(self):
                 for pong in s.pongs:
                     self.ponged(pong)
                 s.pongs = []
-                    
         except:
             print "".join(traceback.format_exception(*exc_info()))
         finally:
             self.client_terminated = self.server_terminated = True
-            self.close_connection()
-        if self.stream.closing:
-            self.closed(self.stream.closing.code, self.stream.closing.reason)
-        else:
-            self.closed(1006)
 
+        try:
+            if not self.server_terminated:
+                if self.stream.closing:
+                    self.closed(self.stream.closing.code, self.stream.closing.reason)
+                else:
+                    self.closed(1006)
+        finally:
+            self.close_connection()
 
 class EchoWebSocketHandler(WebSocketHandler):
     """
diff --git a/ws4py/streaming.py b/ws4py/streaming.py
@@ -1,4 +1,6 @@
 # -*- coding: utf-8 -*-
+import struct
+
 from ws4py.messaging import TextMessage, BinaryMessage, CloseControlMessage,\
      PingControlMessage, PongControlMessage
 from ws4py.framing import Frame, OPCODE_CONTINUATION, OPCODE_TEXT, \
@@ -203,7 +205,14 @@ def receiver(self):
                             m.extend(bytes)
 
                     elif frame.opcode == OPCODE_CLOSE:
-                        self.closing = CloseControlMessage(reason=bytes.decode("utf-8", "replace"))
+                        code = 1000
+                        reason = ""
+                        bytes = str(bytes)
+                        if len(bytes) > 1:
+                            code = struct.unpack("!H", bytes[0:2])[0]
+                            if len(bytes) > 2:
+                                reason = bytes[2:].decode("utf-8", "replace")
+                        self.closing = CloseControlMessage(code=code, reason=reason)
                         
                     elif frame.opcode == OPCODE_PING:
                         self.pings.append(PingControlMessage(bytes))
