From db658b12edf9a44e40ca79209a652e839ffa78dd Mon Sep 17 00:00:00 2001
From: Rajat Gupta <35985127+rjt-gupta@users.noreply.github.com>
Date: Tue, 4 Dec 2018 02:35:59 +0530
Subject: [PATCH] fix query array

fixes #3072
closes #3254
---
 mitmproxy/contentviews/base.py            | 44 ++++++++++++++++-------
 mitmproxy/contentviews/query.py           |  2 +-
 mitmproxy/contentviews/urlencoded.py      |  3 +-
 test/mitmproxy/contentviews/test_base.py  | 15 ++++++++
 test/mitmproxy/contentviews/test_query.py |  4 +--
 5 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/mitmproxy/contentviews/base.py b/mitmproxy/contentviews/base.py
index e4ccf66f7..6072dfb76 100644
--- a/mitmproxy/contentviews/base.py
+++ b/mitmproxy/contentviews/base.py
@@ -35,6 +35,37 @@ class View:
         raise NotImplementedError()  # pragma: no cover
 
 
+def format_pairs(
+        items: typing.Iterable[typing.Tuple[TTextType, TTextType]]
+)-> typing.Iterator[TViewLine]:
+
+    """
+    Helper function that accepts a list of (k,v) pairs into a list of
+    [
+        ("key", key    )
+        ("value", value)
+    ]
+    where key is padded to a uniform width
+    """
+
+    max_key_len = max((len(k[0]) for k in items), default=0)
+    max_key_len = min((max_key_len, KEY_MAX), default=0)
+
+    for key, value in items:
+        if isinstance(key, bytes):
+
+            key += b":"
+        else:
+            key += ":"
+
+        key = key.ljust(max_key_len + 2)
+
+        yield [
+            ("header", key),
+            ("text", value)
+        ]
+
+
 def format_dict(
         d: typing.Mapping[TTextType, TTextType]
 ) -> typing.Iterator[TViewLine]:
@@ -47,18 +78,7 @@ def format_dict(
     entries, where key is padded to a uniform width.
     """
 
-    max_key_len = max((len(k) for k in d.keys()), default=0)
-    max_key_len = min((max_key_len, KEY_MAX), default=0)
-    for key, value in d.items():
-        if isinstance(key, bytes):
-            key += b":"
-        else:
-            key += ":"
-        key = key.ljust(max_key_len + 2)
-        yield [
-            ("header", key),
-            ("text", value)
-        ]
+    return format_pairs(d.items())
 
 
 def format_text(text: TTextType) -> typing.Iterator[TViewLine]:
diff --git a/mitmproxy/contentviews/query.py b/mitmproxy/contentviews/query.py
index 4896624bb..0f74ea2f4 100644
--- a/mitmproxy/contentviews/query.py
+++ b/mitmproxy/contentviews/query.py
@@ -9,6 +9,6 @@ class ViewQuery(base.View):
     def __call__(self, data, **metadata):
         query = metadata.get("query")
         if query:
-            return "Query", base.format_dict(query)
+            return "Query", base.format_pairs(query.items(multi=True))
         else:
             return "Query", base.format_text("")
diff --git a/mitmproxy/contentviews/urlencoded.py b/mitmproxy/contentviews/urlencoded.py
index 371a160e9..e35bbdb7a 100644
--- a/mitmproxy/contentviews/urlencoded.py
+++ b/mitmproxy/contentviews/urlencoded.py
@@ -1,5 +1,4 @@
 from mitmproxy.net.http import url
-from mitmproxy.coretypes import multidict
 from . import base
 
 
@@ -13,4 +12,4 @@ class ViewURLEncoded(base.View):
         except ValueError:
             return None
         d = url.decode(data)
-        return "URLEncoded form", base.format_dict(multidict.MultiDict(d))
+        return "URLEncoded form", base.format_pairs(d)
diff --git a/test/mitmproxy/contentviews/test_base.py b/test/mitmproxy/contentviews/test_base.py
index c94d8be2a..cd879bfda 100644
--- a/test/mitmproxy/contentviews/test_base.py
+++ b/test/mitmproxy/contentviews/test_base.py
@@ -15,3 +15,18 @@ def test_format_dict():
     f_d = base.format_dict(d)
     with pytest.raises(StopIteration):
         next(f_d)
+
+
+def test_format_pairs():
+    d = [("a", "c"), ("b", "d")]
+    f_d = base.format_pairs(d)
+    assert next(f_d)
+
+    d = [("abc", "")]
+    f_d = base.format_pairs(d)
+    assert next(f_d)
+
+    d = []
+    f_d = base.format_pairs(d)
+    with pytest.raises(StopIteration):
+        next(f_d)
diff --git a/test/mitmproxy/contentviews/test_query.py b/test/mitmproxy/contentviews/test_query.py
index 741b23f19..1ae1b3acb 100644
--- a/test/mitmproxy/contentviews/test_query.py
+++ b/test/mitmproxy/contentviews/test_query.py
@@ -6,8 +6,8 @@ from . import full_eval
 def test_view_query():
     d = ""
     v = full_eval(query.ViewQuery())
-    f = v(d, query=multidict.MultiDict([("foo", "bar")]))
+    f = v(d, query=multidict.MultiDict([("foo", "bar"), ("foo", "baz")]))
     assert f[0] == "Query"
-    assert f[1] == [[("header", "foo: "), ("text", "bar")]]
+    assert f[1] == [[("header", "foo: "), ("text", "bar")], [("header", "foo: "), ("text", "baz")]]
 
     assert v(d) == ("Query", [])