Handle redirects in handshake

2 files changed, 17 insertions, 3 deletions

diff --git a/websocket/_core.py b/websocket/_core.py
index aea16e6..04f1fcb 100644
--- a/websocket/_core.py
+++ b/websocket/_core.py
@@ -210,6 +210,7 @@ class WebSocket(object):
                  "http_proxy_auth" - http proxy auth information.
                                      tuple of username and password.
                                      default is None
+                 "redirect_threshold" -> number of redirects to follow.
                  "subprotocols" - array of available sub protocols.
                                   default is None.
                  "socket" - pre-initialized stream socket.
@@ -220,6 +221,13 @@ class WebSocket(object):
 
         try:
             self.handshake_response = handshake(self.sock, *addrs, **options)
+            for attempt in range(options.pop('redirect_threshold', 3)):
+                if self.handshake_response.status in SUPPORTED_REDIRECT_STATUSES:
+                    url = self.handshake_response.headers['location']
+                    self.sock.close()
+                    self.sock, addrs =  connect(url, self.sock_opt, proxy_info(**options),
+                                                options.pop('socket', None))
+                    self.handshake_response = handshake(self.sock, *addrs, **options)
             self.connected = True
         except:
             if self.sock:
@@ -482,6 +490,7 @@ def create_connection(url, timeout=None, class_=WebSocket, **options):
                                     tuple of username and password.
                                     default is None
              "enable_multithread" -> enable lock for multithread.
+             "redirect_threshold" -> number of redirects to follow.
              "sockopt" -> socket options
              "sslopt" -> ssl option
              "subprotocols" - array of available sub protocols.
diff --git a/websocket/_handshake.py b/websocket/_handshake.py
index d9dae47..e28e88d 100644
--- a/websocket/_handshake.py
+++ b/websocket/_handshake.py
@@ -22,6 +22,7 @@ Copyright (C) 2010 Hiroki Ohtani(liris)
 import hashlib
 import hmac
 import os
+from http import HTTPStatus
 
 import six
 
@@ -36,7 +37,7 @@ if six.PY3:
 else:
     from base64 import encodestring as base64encode
 
-__all__ = ["handshake_response", "handshake"]
+__all__ = ["handshake_response", "handshake", "SUPPORTED_REDIRECT_STATUSES"]
 
 if hasattr(hmac, "compare_digest"):
     compare_digest = hmac.compare_digest
@@ -47,6 +48,8 @@ else:
 # websocket supported version.
 VERSION = 13
 
+SUPPORTED_REDIRECT_STATUSES = [HTTPStatus.MOVED_PERMANENTLY, HTTPStatus.FOUND, HTTPStatus.SEE_OTHER]
+
 CookieJar = SimpleCookieJar()
 
 
@@ -67,6 +70,8 @@ def handshake(sock, hostname, port, resource, **options):
     dump("request header", header_str)
 
     status, resp = _get_resp_headers(sock)
+    if status in SUPPORTED_REDIRECT_STATUSES:
+        return handshake_response(status, resp, None)
     success, subproto = _validate(resp, key, options.get("subprotocols"))
     if not success:
         raise WebSocketException("Invalid WebSocket Header")
@@ -134,9 +139,9 @@ def _get_handshake_headers(resource, host, port, options):
     return headers, key
 
 
-def _get_resp_headers(sock, success_status=101):
+def _get_resp_headers(sock, success_statuses=(101, 301, 302, 303)):
     status, resp_headers, status_message = read_headers(sock)
-    if status != success_status:
+    if status not in success_statuses:
         raise WebSocketBadStatusException("Handshake status %d %s", status, status_message)
     return status, resp_headers