Commits

Anonymous committed ca26cf9

fixed werkzeug.test.Client redirect handling

Comments (0)

Files changed (2)

tests/test_test.py

 
 import sys
 from cStringIO import StringIO
+from nose.tools import assert_raises
 from werkzeug.wrappers import Request, Response
 from werkzeug.test import Client, EnvironBuilder, create_environ
+from werkzeug.utils import redirect, get_host
+from werkzeug.datastructures import Headers
+
 
 
 def cookie_app(environ, start_response):
     return response(environ, start_response)
 
 
+def redirect_demo_app(environ, start_response):
+    response = redirect('http://localhost/some/redirect/')
+    return response(environ, start_response)
+
+
+def external_redirect_demo_app(environ, start_response):
+    response = redirect('http://example.org/')
+    return response(environ, start_response)
+
+
 def test_set_cookie_app():
     """Test that a server cookie is set and stored in the client"""
     c = Client(cookie_app)
     builder.files.add_file('blah', SpecialInput())
     builder.close()
     assert len(closed) == 2
+
+
+def test_follow_redirect():
+    env = create_environ('/', base_url='http://localhost')
+    c = Client(redirect_demo_app)
+    headers = Headers(c.open(environ_overrides=env, follow_redirects=True)[2])
+    assert headers['Location'] == 'http://localhost/some/redirect/'
+
+
+def test_follow_external_redirect():
+    env = create_environ('/', base_url='http://localhost')
+    c = Client(external_redirect_demo_app)
+    assert_raises(RuntimeError, lambda: c.open(environ_overrides=env, follow_redirects=True))
 from werkzeug._internal import _empty_stream
 from werkzeug.wrappers import BaseRequest
 from werkzeug.utils import create_environ, run_wsgi_app, get_current_url, \
-     url_encode, url_decode, FileStorage
+     url_encode, url_decode, FileStorage, get_host
 from werkzeug.datastructures import FileMultiDict, MultiDict, \
      CombinedMultiDict, Headers
 
         if self.cookie_jar is not None:
             self.cookie_jar.extract_wsgi(environ, rv[2])
 
+        # handle redirects
+        redirect_chain = []
         status_code = int(rv[1].split(None, 1)[0])
-        if status_code in (301, 302, 303, 305, 307) and follow_redirects:
-            redirect = urlparse.urlunsplit(urlparse.urlsplit(
-                dict(rv[2])['Location'])[:-2] + ('', ''))
+        while status_code in (301, 302, 303, 305, 307) and follow_redirects:
+            redirect = dict(rv[2])['Location']
+            host = get_host(create_environ('/', redirect))
+            if get_host(environ).split(':', 1)[0] != host:
+                raise RuntimeError('%r does not support redirect to '
+                                   'external targets' % self.__class__)
+
+            scheme, netloc, script_root, qs, anchor = urlparse.urlsplit(redirect)
+            redirect_chain.append((redirect, status_code))
+
             kwargs.update({
-                'base_url':         redirect,
+                'base_url':         urlparse.urlunsplit((scheme, host,
+                                    script_root, '', '')).rstrip('/') + '/',
+                'query_string':     qs,
                 'as_tuple':         as_tuple,
                 'buffered':         buffered,
-                'follow_redirects': True
+                'follow_redirects': False
             })
-            return self.open(*args, **kwargs)
+            rv = self.open(*args, **kwargs)
+            status_code = int(rv[1].split(None, 1)[0])
+
+            # Prevent loops
+            if redirect_chain[-1] in redirect_chain[0:-1]:
+                break
 
         response = self.response_wrapper(*rv)
         if as_tuple:
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.