Merge pull request #185 from fschulze/fix-strict-cookie-policy

Fix strict cookie policy.

3 files changed, 39 insertions, 5 deletions

diff --git a/tests/test_app.py b/tests/test_app.py
index ed442b7..0c611fb 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -357,6 +357,39 @@ class TestCookies(unittest.TestCase):
         self.assertEqual(res.request.environ['HTTP_COOKIE'], 'spam=eggs')
         self.assertEqual(dict(res.request.cookies), {'spam': 'eggs'})
 
+    def test_cookie_policy(self):
+        from six.moves import http_cookiejar
+
+        def cookie_app(environ, start_response):
+            status = to_bytes("200 OK")
+            body = 'Cookie.'
+            headers = [
+                ('Content-Type', 'text/plain'),
+                ('Content-Length', str(len(body))),
+                ('Set-Cookie',
+                 'spam=eggs; secure; Domain=.example.org;'),
+            ]
+            start_response(status, headers)
+            return [to_bytes(body)]
+
+        policy = webtest.app.CookiePolicy()
+        flags = (
+            policy.DomainStrictNoDots |
+            policy.DomainRFC2965Match |
+            policy.DomainStrictNonDomain)
+        policy.strict_ns_domain |= flags
+        cookiejar = http_cookiejar.CookieJar(policy=policy)
+        app = webtest.TestApp(
+            cookie_app,
+            cookiejar=cookiejar,
+            extra_environ={'HTTP_HOST': 'example.org'})
+        res = app.get('/')
+        res = app.get('/')
+        self.assertFalse(app.cookies,
+                        'Response should not have set cookies')
+        self.assertNotIn('HTTP_COOKIE', res.request.environ)
+        self.assertEqual(dict(res.request.cookies), {})
+
 
 class TestEnviron(unittest.TestCase):
 
diff --git a/webtest/app.py b/webtest/app.py
index 16aeee1..1c19e94 100644
--- a/webtest/app.py
+++ b/webtest/app.py
@@ -69,13 +69,13 @@ class CookiePolicy(http_cookiejar.DefaultCookiePolicy):
     Domain=localhost."""
 
     def return_ok_domain(self, cookie, request):
-        if cookie.domain.endswith(request.origin_req_host):
+        if cookie.domain == '.localhost':
             return True
         return http_cookiejar.DefaultCookiePolicy.return_ok_domain(
             self, cookie, request)
 
     def set_ok_domain(self, cookie, request):
-        if cookie.domain.endswith(request.origin_req_host):
+        if cookie.domain == '.localhost':
             return True
         return http_cookiejar.DefaultCookiePolicy.set_ok_domain(
             self, cookie, request)
@@ -231,9 +231,10 @@ class TestApp(object):
         Sets a cookie to be passed through with requests.
 
         """
-        cookie_domain = self.extra_environ.get('HTTP_HOST', 'localhost')
+        cookie_domain = self.extra_environ.get('HTTP_HOST', '.localhost')
         cookie_domain = cookie_domain.split(':', 1)[0]
-        cookie_domain = '.' + cookie_domain
+        if '.' not in cookie_domain:
+            cookie_domain = "%s.local" % cookie_domain
         value = escape_cookie_value(value)
         cookie = http_cookiejar.Cookie(
             version=0,
diff --git a/webtest/utils.py b/webtest/utils.py
index adb01f3..3eb7f90 100644
--- a/webtest/utils.py
+++ b/webtest/utils.py
@@ -97,7 +97,7 @@ class _RequestCookieAdapter(object):
     """
     def __init__(self, request):
         self._request = request
-        self.origin_req_host = request.host.split(':')[0]
+        self.origin_req_host = request.host
 
     def is_unverifiable(self):
         return True  # sure? Why not?