support for setting/sending multiple cookies (#1091)

* support for setting/sending multiple cookies

* py.test for multiple cookie support

2 files changed, 24 insertions, 4 deletions

diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 20cd592b..bb8bf2ea 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -22,6 +22,7 @@ from .proxy.config import HostMatcher
 from .protocol.http_replay import RequestReplayThread
 from .exceptions import Kill
 from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest
+from collections import defaultdict
 
 
 class AppRegistry:
@@ -309,7 +310,7 @@ class StickyCookieState:
         """
             flt: Compiled filter.
         """
-        self.jar = {}
+        self.jar = defaultdict(dict)
         self.flt = flt
 
     def ckey(self, m, f):
@@ -337,7 +338,7 @@ class StickyCookieState:
             for m in c.values():
                 k = self.ckey(m, f)
                 if self.domain_match(f.request.host, k[0]):
-                    self.jar[k] = m
+                    self.jar[k][m.key] = m
 
     def handle_request(self, f):
         l = []
@@ -349,10 +350,10 @@ class StickyCookieState:
                     f.request.path.startswith(i[2])
                 ]
                 if all(match):
-                    l.append(self.jar[i].output(header="").strip())
+                    l.extend([m.output(header="").strip() for m in self.jar[i].values()])
         if l:
             f.request.stickycookie = True
-            f.request.headers.set_all("cookie", l)
+            f.request.headers["cookie"] = "; ".join(l)
 
 
 class StickyAuthState:
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 60f6b1a9..8729cc77 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -69,6 +69,25 @@ class TestStickyCookieState:
         s, f = self._response("SSID=mooo", "www.google.com")
         assert s.jar.keys()[0] == ('www.google.com', 80, '/')
 
+        # Test setting of multiple cookies
+        c1 = "somecookie=test; Path=/"
+        c2 = "othercookie=helloworld; Path=/"
+        s, f = self._response(c1, "www.google.com")
+        f.response.headers["Set-Cookie"] = c2
+        s.handle_response(f)
+        googlekey = s.jar.keys()[0]
+        assert len(s.jar[googlekey].keys()) == 2
+
+        # Test overwriting of a cookie value
+        c1 = "somecookie=helloworld; Path=/"
+        c2 = "somecookie=newvalue; Path=/"
+        s, f = self._response(c1, "www.google.com")
+        f.response.headers["Set-Cookie"] = c2
+        s.handle_response(f)
+        googlekey = s.jar.keys()[0]
+        assert len(s.jar[googlekey].keys()) == 1
+        assert s.jar[googlekey]["somecookie"].value == "newvalue"
+
     def test_handle_request(self):
         s, f = self._response("SSID=mooo", "www.google.com")
         assert "cookie" not in f.request.headers