Python Server - Processing WebSocket.close()

2.3k Views Asked by At

I'm working on a small server from the code provided below. The server will interact with a websocket to take a message from a websocket, send a message back to the websocket, and once the websocket calls .close() the server will "shutdown" or close the connection.

However at the moment, when the websocket calls .close() the time it takes to disconnect is long. ( I'm assuming it's the 300 second time out for websockets that occurs when ping/pong messages aren't sent to keep the connection alive). In turn, I think this means that the connection is not closing completely and is just getting lost.

Is there a way to determine and fix this?

I've been reading up on the python socket server documentation and other various sources to figure it out and have come up short. My only solution is the hack comment in the code below.

  1. https://docs.python.org/2/library/socketserver.html
  2. http://pymotw.com/2/SocketServer/

    import struct
    import SocketServer
    from base64 import b64encode
    from hashlib import sha1
    from mimetools import Message
    from StringIO import StringIO
    
    class WebSocketsHandler(SocketServer.StreamRequestHandler):
    magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
    
    def setup(self):
        SocketServer.StreamRequestHandler.setup(self)
        print "connection established", self.client_address
        self.handshake_done = False
    
    def handle(self):
        while True:
            if not self.handshake_done:
                self.handshake()
            else:
                self.read_next_message()
    
    def read_next_message(self):
        length = ord(self.rfile.read(2)[1]) & 127
        if length == 126:
            length = struct.unpack(">H", self.rfile.read(2))[0]
        elif length == 127:
            length = struct.unpack(">Q", self.rfile.read(8))[0]
        masks = [ord(byte) for byte in self.rfile.read(4)]
        decoded = ""
        for char in self.rfile.read(length):
            decoded += chr(ord(char) ^ masks[len(decoded) % 4])
        self.on_message(decoded)
    
    def send_message(self, message):
        self.request.send(chr(129))
        length = len(message)
        if length <= 125:
            self.request.send(chr(length))
        elif length >= 126 and length <= 65535:
            self.request.send(chr(126))
            self.request.send(struct.pack(">H", length))
        else:
            self.request.send(chr(127))
            self.request.send(struct.pack(">Q", length))
        self.request.send(message)
    
    def handshake(self):
        data = self.request.recv(1024).strip()
        headers = Message(StringIO(data.split('\r\n', 1)[1]))
        if headers.get("Upgrade", None) != "websocket":
            return
        print 'Handshaking...'
        key = headers['Sec-WebSocket-Key']
        digest = b64encode(sha1(key + self.magic).hexdigest().decode('hex'))
        response = 'HTTP/1.1 101 Switching Protocols\r\n'
        response += 'Upgrade: websocket\r\n'
        response += 'Connection: Upgrade\r\n'
        response += 'Sec-WebSocket-Accept: %s\r\n\r\n' % digest
        self.handshake_done = self.request.send(response)
    
    def on_message(self, message):
        print message
        self.send_message("How do you do?")
    
    
    if __name__ == "__main__":
        server = SocketServer.TCPServer(
        ("", 9999), WebSocketsHandler)
    server.serve_forever()
    

[Slightly] modified from: https://gist.github.com/jkp/3136208

NOTE: I am aware of the server.serve_forever() and how that may be an issue. I'm looking for a suggestion or direction along the lines of :

def on_close(): ...

1

There are 1 best solutions below

0
On

I think you need to implement an onClose() method on your server side. Why you do not use some python framework that has websocket support such as Tornado or Authobahn.

Here is an example for websocket server in python using Tornado https://developer.mbed.org/cookbook/Websockets-Server

I hope this will help you.