Create ODictCaseless for headers, use vanilla ODict for everything else.

5 files changed, 72 insertions, 44 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 6678771f..63f80dde 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -45,11 +45,14 @@ class ScriptContext:
 
 
 class ODict:
+    """
+        A dictionary-like object for managing ordered (key, value) data. 
+    """
     def __init__(self, lst=None):
         self.lst = lst or []
 
     def _kconv(self, s):
-        return s.lower()
+        return s
 
     def __eq__(self, other):
         return self.lst == other.lst
@@ -105,7 +108,7 @@ class ODict:
             Returns a copy of this object.
         """
         lst = copy.deepcopy(self.lst)
-        return ODict(lst)
+        return self.__class__(lst)
 
     def __repr__(self):
         elements = []
@@ -143,6 +146,15 @@ class ODict:
         return count
 
 
+class ODictCaseless(ODict):
+    """
+        A variant of ODict with "caseless" keys. This version _preserves_ key
+        case, but does not consider case when setting or getting items. 
+    """
+    def _kconv(self, s):
+        return s.lower()
+
+
 class HTTPMsg(controller.Msg):
     def decode(self):
         """
@@ -176,7 +188,7 @@ class Request(HTTPMsg):
         Exposes the following attributes:
 
             client_conn: ClientConnection object, or None if this is a replay.
-            headers: ODict object
+            headers: ODictCaseless object
             content: Content of the request, or None
 
             scheme: URL scheme (http/https)
@@ -188,6 +200,7 @@ class Request(HTTPMsg):
             method: HTTP method
     """
     def __init__(self, client_conn, host, port, scheme, method, path, headers, content, timestamp=None):
+        assert isinstance(headers, ODictCaseless)
         self.client_conn = client_conn
         self.host, self.port, self.scheme = host, port, scheme
         self.method, self.path, self.headers, self.content = method, path, headers, content
@@ -253,7 +266,7 @@ class Request(HTTPMsg):
         self.scheme = state["scheme"]
         self.method = state["method"]
         self.path = state["path"]
-        self.headers = ODict._from_state(state["headers"])
+        self.headers = ODictCaseless._from_state(state["headers"])
         self.content = state["content"]
         self.timestamp = state["timestamp"]
 
@@ -279,7 +292,7 @@ class Request(HTTPMsg):
             str(state["scheme"]),
             str(state["method"]),
             str(state["path"]),
-            ODict._from_state(state["headers"]),
+            ODictCaseless._from_state(state["headers"]),
             state["content"],
             state["timestamp"]
         )
@@ -415,6 +428,7 @@ class Response(HTTPMsg):
             timestamp: Seconds since the epoch
     """
     def __init__(self, request, code, msg, headers, content, timestamp=None):
+        assert isinstance(headers, ODictCaseless)
         self.request = request
         self.code, self.msg = code, msg
         self.headers, self.content = headers, content
@@ -484,7 +498,7 @@ class Response(HTTPMsg):
     def _load_state(self, state):
         self.code = state["code"]
         self.msg = state["msg"]
-        self.headers = ODict._from_state(state["headers"])
+        self.headers = ODictCaseless._from_state(state["headers"])
         self.content = state["content"]
         self.timestamp = state["timestamp"]
 
@@ -503,7 +517,7 @@ class Response(HTTPMsg):
             request,
             state["code"],
             str(state["msg"]),
-            ODict._from_state(state["headers"]),
+            ODictCaseless._from_state(state["headers"]),
             state["content"],
             state["timestamp"],
         )
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index ff5e3ec7..8e0aa9cd 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -52,7 +52,7 @@ def read_headers(fp):
                 name = line[:i]
                 value = line[i+1:].strip()
                 ret.append([name, value])
-    return flow.ODict(ret)
+    return flow.ODictCaseless(ret)
 
 
 def read_chunked(fp, limit):
diff --git a/test/test_filt.py b/test/test_filt.py
index 6c9df232..d9d92ffe 100644
--- a/test/test_filt.py
+++ b/test/test_filt.py
@@ -74,7 +74,7 @@ class uParsing(libpry.AutoTree):
 class uMatching(libpry.AutoTree):
     def req(self):
         conn = flow.ClientConnect(("one", 2222))
-        headers = flow.ODict()
+        headers = flow.ODictCaseless()
         headers["header"] = ["qvalue"]
         return flow.Request(
                     conn,
@@ -89,7 +89,7 @@ class uMatching(libpry.AutoTree):
 
     def resp(self):
         q = self.req()
-        headers = flow.ODict()
+        headers = flow.ODictCaseless()
         headers["header_response"] = ["svalue"]
         return flow.Response(
                     q,
diff --git a/test/test_flow.py b/test/test_flow.py
index 7ca489c6..52e6781f 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -617,7 +617,7 @@ class uFlowMaster(libpry.AutoTree):
 
 class uRequest(libpry.AutoTree):
     def test_simple(self):
-        h = flow.ODict()
+        h = flow.ODictCaseless()
         h["test"] = ["test"]
         c = flow.ClientConnect(("addr", 2222))
         r = flow.Request(c, "host", 22, "https", "GET", "/", h, "content")
@@ -639,7 +639,7 @@ class uRequest(libpry.AutoTree):
         assert r._assemble(True)
 
     def test_getset_form_urlencoded(self):
-        h = flow.ODict()
+        h = flow.ODictCaseless()
         h["content-type"] = [flow.HDR_FORM_URLENCODED]
         d = flow.ODict([("one", "two"), ("three", "four")])
         r = flow.Request(None, "host", 22, "https", "GET", "/", h, utils.urlencode(d.lst))
@@ -653,7 +653,7 @@ class uRequest(libpry.AutoTree):
         assert not r.get_form_urlencoded()
 
     def test_getset_query(self):
-        h = flow.ODict()
+        h = flow.ODictCaseless()
 
         r = flow.Request(None, "host", 22, "https", "GET", "/foo?x=y&a=b", h, "content")
         q = r.get_query()
@@ -676,7 +676,7 @@ class uRequest(libpry.AutoTree):
         assert r.get_query() == qv
 
     def test_anticache(self):
-        h = flow.ODict()
+        h = flow.ODictCaseless()
         r = flow.Request(None, "host", 22, "https", "GET", "/", h, "content")
         h["if-modified-since"] = ["test"]
         h["if-none-match"] = ["test"]
@@ -685,7 +685,7 @@ class uRequest(libpry.AutoTree):
         assert not "if-none-match" in r.headers
 
     def test_getset_state(self):
-        h = flow.ODict()
+        h = flow.ODictCaseless()
         h["test"] = ["test"]
         c = flow.ClientConnect(("addr", 2222))
         r = flow.Request(c, "host", 22, "https", "GET", "/", h, "content")
@@ -753,7 +753,7 @@ class uRequest(libpry.AutoTree):
 
 class uResponse(libpry.AutoTree):
     def test_simple(self):
-        h = flow.ODict()
+        h = flow.ODictCaseless()
         h["test"] = ["test"]
         c = flow.ClientConnect(("addr", 2222))
         req = flow.Request(c, "host", 22, "https", "GET", "/", h, "content")
@@ -798,7 +798,7 @@ class uResponse(libpry.AutoTree):
 
 
     def test_getset_state(self):
-        h = flow.ODict()
+        h = flow.ODictCaseless()
         h["test"] = ["test"]
         c = flow.ClientConnect(("addr", 2222))
         req = flow.Request(c, "host", 22, "https", "GET", "/", h, "content")
@@ -886,31 +886,31 @@ class uClientConnect(libpry.AutoTree):
 
 class uODict(libpry.AutoTree):
     def setUp(self):
-        self.hd = flow.ODict()
+        self.od = flow.ODict()
 
     def test_str_err(self):
         h = flow.ODict()
         libpry.raises(ValueError, h.__setitem__, "key", "foo")
 
     def test_dictToHeader1(self):
-        self.hd.add("one", "uno")
-        self.hd.add("two", "due")
-        self.hd.add("two", "tre")
+        self.od.add("one", "uno")
+        self.od.add("two", "due")
+        self.od.add("two", "tre")
         expected = [
             "one: uno\r\n",
             "two: due\r\n",
             "two: tre\r\n",
             "\r\n"
         ]
-        out = repr(self.hd)
+        out = repr(self.od)
         for i in expected:
             assert out.find(i) >= 0
 
     def test_dictToHeader2(self):
-        self.hd["one"] = ["uno"]
+        self.od["one"] = ["uno"]
         expected1 = "one: uno\r\n"
         expected2 = "\r\n"
-        out = repr(self.hd)
+        out = repr(self.od)
         assert out.find(expected1) >= 0
         assert out.find(expected2) >= 0
 
@@ -924,36 +924,49 @@ class uODict(libpry.AutoTree):
         assert not h.match_re("nonono")
 
     def test_getset_state(self):
-        self.hd.add("foo", 1)
-        self.hd.add("foo", 2)
-        self.hd.add("bar", 3)
-        state = self.hd._get_state()
+        self.od.add("foo", 1)
+        self.od.add("foo", 2)
+        self.od.add("bar", 3)
+        state = self.od._get_state()
         nd = flow.ODict._from_state(state)
-        assert nd == self.hd
+        assert nd == self.od
 
     def test_copy(self):
-        self.hd.add("foo", 1)
-        self.hd.add("foo", 2)
-        self.hd.add("bar", 3)
-        assert self.hd == self.hd.copy()
+        self.od.add("foo", 1)
+        self.od.add("foo", 2)
+        self.od.add("bar", 3)
+        assert self.od == self.od.copy()
 
     def test_del(self):
-        self.hd.add("foo", 1)
-        self.hd.add("Foo", 2)
-        self.hd.add("bar", 3)
-        del self.hd["foo"]
-        assert len(self.hd.lst) == 1
+        self.od.add("foo", 1)
+        self.od.add("Foo", 2)
+        self.od.add("bar", 3)
+        del self.od["foo"]
+        assert len(self.od.lst) == 2
 
     def test_replace(self):
-        self.hd.add("one", "two")
-        self.hd.add("two", "one")
-        assert self.hd.replace("one", "vun") == 2
-        assert self.hd.lst == [
+        self.od.add("one", "two")
+        self.od.add("two", "one")
+        assert self.od.replace("one", "vun") == 2
+        assert self.od.lst == [
             ["vun", "two"],
             ["two", "vun"],
         ]
 
 
+class uODictCaseless(libpry.AutoTree):
+    def setUp(self):
+        self.od = flow.ODictCaseless()
+
+    def test_del(self):
+        self.od.add("foo", 1)
+        self.od.add("Foo", 2)
+        self.od.add("bar", 3)
+        del self.od["foo"]
+        assert len(self.od) == 1
+
+
+
 tests = [
     uStickyCookieState(),
     uStickyAuthState(),
@@ -968,4 +981,5 @@ tests = [
     uError(),
     uClientConnect(),
     uODict(),
+    uODictCaseless(),
 ]
diff --git a/test/tutils.py b/test/tutils.py
index c22eb1a0..b11c997a 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -7,7 +7,7 @@ import random
 def treq(conn=None):
     if not conn:
         conn = flow.ClientConnect(("address", 22))
-    headers = flow.ODict()
+    headers = flow.ODictCaseless()
     headers["header"] = ["qvalue"]
     return flow.Request(conn, "host", 80, "http", "GET", "/path", headers, "content")
 
@@ -15,7 +15,7 @@ def treq(conn=None):
 def tresp(req=None):
     if not req:
         req = treq()
-    headers = flow.ODict()
+    headers = flow.ODictCaseless()
     headers["header_response"] = ["svalue"]
     return flow.Response(req, 200, "message", headers, "content_response")