Search code examples
pythonwebsockettornadohandshake

How to refuse a websocket handshake in a tornado websocket handler


So I'm doing a unittest of a tornado server I programmed which has a websocketHandler. This is my testing code:

def test_unauthorized_websocket(self):
    message = "Hey bro"     
    ws = websocket.create_connection('ws://localhost:8888/ws', header=['Authorization: false'])
    send_dict = {'command': test_command, 'message': message}
    serialized_dict = json.dumps(send_dict)
    ws.send(serialized_dict)
    response = ws.recv()
    #test response with assert

My goal with this test was to prove that my tornado server correctly refuses and closes this websocket connection because of wrong authentication header. This is my tornado websocketHandler code:

class WebSocketHandler(tornado.websocket.WebSocketHandler):

   def open(self):
      #some code
      headers = self.request.headers
      try:
          auth = headers['Authorization']
      except KeyError:
          self.close(code = 1002, reason = "Unauthorized websocket")
          print("IT'S CLOSED")
          return

       if auth == "true":
          print("Authorized Websocket!")
          #some code
       else:
          print("Unauthorized Websocket... :-(")
          self.close(code = 1002, reason = "Unauthorized websocket")

So when the authentication is wrong, self.close() is called (not sure I need code and reason). This should close the websocket. But this doesn't actually happens in the "client" side. After I call create_connection() in the "client" the ws.connected variable is still True, and when I do ws.send() the on_message method of the websocketHandler is still called and when it tries to make a response with self.write_message() it raises a WebSocketClosedError. And only then the "client" actually closes its side of the websocket, ws.recv() returns nothing and ws.connected turns to False after that. Is there a way I can communicate to the client side (through handshake headers or something) that the websocket is meant to be closed earlier on its side?


Solution

  • You can override prepare() and raise a tornado.web.HTTPError, instead of overriding open() and calling self.close(). This will reject the connection at the first opportunity and this will be reported on the client side as a part of create_connection() instead of on a later read.