Prefix dot only for suffix check and add test

tirkarthi · tirkarthi · commit dfdc776358e8 · 2018-12-24T13:41:07.000+05:30
diff --git a/Lib/http/cookiejar.py b/Lib/http/cookiejar.py
@@ -1169,13 +1169,14 @@ def domain_return_ok(self, domain, request):
         # Liberal check of.  This is here as an optimization to avoid
         # having to load lots of MSIE cookie files unless necessary.
         req_host, erhn = eff_request_host(request)
+        suffix_check_domain = domain
         if not req_host.startswith("."):
             req_host = "."+req_host
         if not erhn.startswith("."):
             erhn = "."+erhn
-        if not domain.startswith("."):
-            domain = "."+domain
-        if not (req_host.endswith(domain) or erhn.endswith(domain)):
+        if suffix_check_domain and not suffix_check_domain.startswith("."):
+            suffix_check_domain = "." + suffix_check_domain
+        if not (req_host.endswith(suffix_check_domain) or erhn.endswith(suffix_check_domain)):
             #_debug("   request domain %s does not match cookie domain %s",
             #       req_host, domain)
             return False
diff --git a/Lib/test/test_http_cookiejar.py b/Lib/test/test_http_cookiejar.py
@@ -961,6 +961,31 @@ def test_domain_block(self):
         c.add_cookie_header(req)
         self.assertFalse(req.has_header("Cookie"))
 
+        c.clear()
+
+        pol.set_blocked_domains([])
+        req = urllib.request.Request("http://acme.com/")
+        res = FakeResponse(headers, "http://acme.com/")
+        c.extract_cookies(res, req)
+        self.assertEqual(len(c), 1)
+
+        req = urllib.request.Request("http://acme.com/")
+        c.add_cookie_header(req)
+        self.assertTrue(req.has_header("Cookie"))
+
+        req = urllib.request.Request("http://badacme.com/")
+        c.add_cookie_header(req)
+        self.assertFalse(req.has_header("Cookie"))
+
+        p = pol.set_blocked_domains(["acme.com"])
+        req = urllib.request.Request("http://acme.com/")
+        c.add_cookie_header(req)
+        self.assertFalse(req.has_header("Cookie"))
+
+        req = urllib.request.Request("http://badacme.com/")
+        c.add_cookie_header(req)
+        self.assertFalse(req.has_header("Cookie"))
+
     def test_secure(self):
         for ns in True, False:
             for whitespace in " ", "":
