From 188f834e0ae305df36438c9a9ad1e2a4b6e98098 Mon Sep 17 00:00:00 2001
From: Michael Barton <mike@weirdlooking.com>
Date: Thu, 8 Nov 2012 14:17:40 -0800
Subject: [PATCH] use Host: for location rewrites

If the Host header is available, use it for making relative URLs absolute.
Otherwise, continue using SERVER_NAME.

Change-Id: Ifc028264ad1b122a2d5dff9d5528cb369090429f
---
 swift/common/swob.py          | 23 ++++++++++++--
 test/unit/common/test_swob.py | 57 +++++++++++++++++++++++++++++++++--
 2 files changed, 75 insertions(+), 5 deletions(-)

diff --git a/swift/common/swob.py b/swift/common/swob.py
index d30cbb8754..1e02f676af 100755
--- a/swift/common/swob.py
+++ b/swift/common/swob.py
@@ -955,12 +955,29 @@ class Response(object):
                 return [body]
         return ['']
 
+    def absolute_location(self):
+        """
+        Attempt to construct an absolute location.
+        """
+        if not self.location.startswith('/'):
+            return self.location
+        if 'HTTP_HOST' in self.environ:
+            host = self.environ['HTTP_HOST']
+        else:
+            host = '%s:%s' % (self.environ['SERVER_NAME'],
+                              self.environ['SERVER_PORT'])
+        scheme = self.environ.get('wsgi.url_scheme', 'http')
+        if scheme == 'http' and host.endswith(':80'):
+            host, port = host.rsplit(':', 1)
+        elif scheme == 'https' and host.endswith(':443'):
+            host, port = host.rsplit(':', 1)
+        return '%s://%s%s' % (scheme, host, self.location)
+
     def __call__(self, env, start_response):
         self.environ = env
         app_iter = self._response_iter(self.app_iter, self._body)
-        if 'location' in self.headers and self.location.startswith('/'):
-            self.location = self.environ['wsgi.url_scheme'] + '://' \
-                + self.environ['SERVER_NAME'] + self.location
+        if 'location' in self.headers:
+            self.location = self.absolute_location()
         start_response(self.status, self.headers.items())
         return app_iter
 
diff --git a/test/unit/common/test_swob.py b/test/unit/common/test_swob.py
index 587d7bec78..6c3f5250c3 100755
--- a/test/unit/common/test_swob.py
+++ b/test/unit/common/test_swob.py
@@ -437,11 +437,64 @@ class TestResponse(unittest.TestCase):
     def test_location_rewrite(self):
         def start_response(env, headers):
             pass
-        req = swift.common.swob.Request.blank('/')
+        req = swift.common.swob.Request.blank(
+            '/', environ={'HTTP_HOST': 'somehost'})
         resp = self._get_response()
         resp.location = '/something'
         body = ''.join(resp(req.environ, start_response))
-        self.assertEquals(resp.location, 'http://localhost/something')
+        self.assertEquals(resp.location, 'http://somehost/something')
+
+        req = swift.common.swob.Request.blank(
+            '/', environ={'HTTP_HOST': 'somehost:80'})
+        resp = self._get_response()
+        resp.location = '/something'
+        body = ''.join(resp(req.environ, start_response))
+        self.assertEquals(resp.location, 'http://somehost/something')
+
+        req = swift.common.swob.Request.blank(
+            '/', environ={'HTTP_HOST': 'somehost:443',
+                          'wsgi.url_scheme': 'http'})
+        resp = self._get_response()
+        resp.location = '/something'
+        body = ''.join(resp(req.environ, start_response))
+        self.assertEquals(resp.location, 'http://somehost:443/something')
+
+        req = swift.common.swob.Request.blank(
+            '/', environ={'HTTP_HOST': 'somehost:443',
+                          'wsgi.url_scheme': 'https'})
+        resp = self._get_response()
+        resp.location = '/something'
+        body = ''.join(resp(req.environ, start_response))
+        self.assertEquals(resp.location, 'https://somehost/something')
+
+    def test_location_rewrite_no_host(self):
+        def start_response(env, headers):
+            pass
+        req = swift.common.swob.Request.blank(
+            '/', environ={'SERVER_NAME': 'local', 'SERVER_PORT': 80})
+        del req.environ['HTTP_HOST']
+        resp = self._get_response()
+        resp.location = '/something'
+        body = ''.join(resp(req.environ, start_response))
+        self.assertEquals(resp.location, 'http://local/something')
+
+        req = swift.common.swob.Request.blank(
+            '/', environ={'SERVER_NAME': 'local', 'SERVER_PORT': 81})
+        del req.environ['HTTP_HOST']
+        resp = self._get_response()
+        resp.location = '/something'
+        body = ''.join(resp(req.environ, start_response))
+        self.assertEquals(resp.location, 'http://local:81/something')
+
+    def test_location_no_rewrite(self):
+        def start_response(env, headers):
+            pass
+        req = swift.common.swob.Request.blank(
+            '/', environ={'HTTP_HOST': 'somehost'})
+        resp = self._get_response()
+        resp.location = 'http://www.google.com/'
+        body = ''.join(resp(req.environ, start_response))
+        self.assertEquals(resp.location, 'http://www.google.com/')
 
     def test_app_iter(self):
         def start_response(env, headers):