From 66267ad2768686f6af0ec20dfa89d1a281fc7f83 Mon Sep 17 00:00:00 2001
From: Steven Van Acker <StevenVanAcker@users.noreply.github.com>
Date: Fri, 22 Apr 2016 20:16:05 +0200
Subject: [PATCH] support for setting/sending multiple cookies (#1091)

* support for setting/sending multiple cookies

* py.test for multiple cookie support
---
 mitmproxy/flow.py           |  9 +++++----
 test/mitmproxy/test_flow.py | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 20cd592ba..bb8bf2eab 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 60f6b1a92..8729cc779 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