fix #2228

3 files changed, 20 insertions, 24 deletions

diff --git a/mitmproxy/addons/stickycookie.py b/mitmproxy/addons/stickycookie.py
index 04d99975..e58e0a58 100644
--- a/mitmproxy/addons/stickycookie.py
+++ b/mitmproxy/addons/stickycookie.py
@@ -1,14 +1,14 @@
 import collections
 from http import cookiejar
+from typing import List, Tuple, Dict, Optional  # noqa
 
+from mitmproxy import http, flowfilter, ctx, exceptions
 from mitmproxy.net.http import cookies
 
-from mitmproxy import exceptions
-from mitmproxy import flowfilter
-from mitmproxy import ctx
+TOrigin = Tuple[str, int, str]
 
 
-def ckey(attrs, f):
+def ckey(attrs: Dict[str, str], f: http.HTTPFlow) -> TOrigin:
     """
         Returns a (domain, port, path) tuple.
     """
@@ -21,18 +21,18 @@ def ckey(attrs, f):
     return (domain, f.request.port, path)
 
 
-def domain_match(a, b):
-    if cookiejar.domain_match(a, b):
+def domain_match(a: str, b: str) -> bool:
+    if cookiejar.domain_match(a, b):  # type: ignore
         return True
-    elif cookiejar.domain_match(a, b.strip(".")):
+    elif cookiejar.domain_match(a, b.strip(".")):  # type: ignore
         return True
     return False
 
 
 class StickyCookie:
     def __init__(self):
-        self.jar = collections.defaultdict(dict)
-        self.flt = None
+        self.jar = collections.defaultdict(dict)  # type: Dict[TOrigin, Dict[str, str]]
+        self.flt = None  # type: Optional[flowfilter.TFilter]
 
     def configure(self, updated):
         if "stickycookie" in updated:
@@ -46,7 +46,7 @@ class StickyCookie:
             else:
                 self.flt = None
 
-    def response(self, flow):
+    def response(self, flow: http.HTTPFlow):
         if self.flt:
             for name, (value, attrs) in flow.response.cookies.items(multi=True):
                 # FIXME: We now know that Cookie.py screws up some cookies with
@@ -63,24 +63,21 @@ class StickyCookie:
                         if not self.jar[dom_port_path]:
                             self.jar.pop(dom_port_path, None)
                     else:
-                        b = attrs.copy()
-                        b.insert(0, name, value)
-                        self.jar[dom_port_path][name] = b
+                        self.jar[dom_port_path][name] = value
 
-    def request(self, flow):
+    def request(self, flow: http.HTTPFlow):
         if self.flt:
-            l = []
+            cookie_list = []  # type: List[Tuple[str,str]]
             if flowfilter.match(self.flt, flow):
-                for domain, port, path in self.jar.keys():
+                for (domain, port, path), c in self.jar.items():
                     match = [
                         domain_match(flow.request.host, domain),
                         flow.request.port == port,
                         flow.request.path.startswith(path)
                     ]
                     if all(match):
-                        c = self.jar[(domain, port, path)]
-                        l.extend([cookies.format_cookie_header(c[name].items(multi=True)) for name in c.keys()])
-            if l:
+                        cookie_list.extend(c.items())
+            if cookie_list:
                 # FIXME: we need to formalise this...
-                flow.request.stickycookie = True
-                flow.request.headers["cookie"] = "; ".join(l)
+                flow.metadata["stickycookie"] = True
+                flow.request.headers["cookie"] = cookies.format_cookie_header(cookie_list)
diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py
index d2cefe79..d8aa593b 100644
--- a/test/mitmproxy/addons/test_dumper.py
+++ b/test/mitmproxy/addons/test_dumper.py
@@ -68,7 +68,6 @@ def test_simple():
         ctx.configure(d, flow_detail=4)
         flow = tflow.tflow()
         flow.request = tutils.treq()
-        flow.request.stickycookie = True
         flow.client_conn = mock.MagicMock()
         flow.client_conn.address[0] = "foo"
         flow.response = tutils.tresp(content=None)
diff --git a/test/mitmproxy/addons/test_stickycookie.py b/test/mitmproxy/addons/test_stickycookie.py
index 9092e09b..f77d019d 100644
--- a/test/mitmproxy/addons/test_stickycookie.py
+++ b/test/mitmproxy/addons/test_stickycookie.py
@@ -110,8 +110,8 @@ class TestStickyCookie:
             f.response.headers["Set-Cookie"] = c2
             sc.response(f)
             googlekey = list(sc.jar.keys())[0]
-            assert len(sc.jar[googlekey].keys()) == 1
-            assert list(sc.jar[googlekey]["somecookie"].items())[0][1] == "newvalue"
+            assert len(sc.jar[googlekey]) == 1
+            assert sc.jar[googlekey]["somecookie"] == "newvalue"
 
     def test_response_delete(self):
         sc = stickycookie.StickyCookie()