http2: implement more language features

3 files changed, 151 insertions, 53 deletions

diff --git a/libpathod/language/http.py b/libpathod/language/http.py
index 3979a1ee..3c9df484 100644
--- a/libpathod/language/http.py
+++ b/libpathod/language/http.py
@@ -82,12 +82,10 @@ class ShortcutUserAgent(_HeaderMixin, base.OptionsOrValue):
     key = base.TokValueLiteral("User-Agent")
 
     def values(self, settings):
+        value = self.value.val
         if self.option_used:
-            value = http_uastrings.get_by_shortcut(
-                self.value.val.lower()
-            )[2]
-        else:
-            value = self.value.val
+            value = http_uastrings.get_by_shortcut(value.lower())[2]
+
         return self.format_header(
             self.key.get_generator(settings),
             value
@@ -143,12 +141,12 @@ class _HTTPMessage(message.Message):
 class Response(_HTTPMessage):
     unique_name = None
     comps = (
-        Body,
         Header,
         ShortcutContentType,
         ShortcutLocation,
         Raw,
         Reason,
+        Body,
 
         actions.PauseAt,
         actions.DisconnectAt,
@@ -256,12 +254,12 @@ class NestedResponse(base.NestedMessage):
 
 class Request(_HTTPMessage):
     comps = (
-        Body,
         Header,
         ShortcutContentType,
         ShortcutUserAgent,
         Raw,
         NestedResponse,
+        Body,
         Times,
 
         actions.PauseAt,
diff --git a/libpathod/language/http2.py b/libpathod/language/http2.py
index 86e04056..b206c0cf 100644
--- a/libpathod/language/http2.py
+++ b/libpathod/language/http2.py
@@ -1,4 +1,6 @@
 import pyparsing as pp
+
+from netlib import http_status, http_uastrings
 from . import base, message
 
 """
@@ -24,6 +26,65 @@ from . import base, message
         h2f:42:DATA:END_STREAM,PADDED:0x1234567:'content body payload'
 """
 
+def get_header(val, headers):
+    """
+        Header keys may be Values, so we have to "generate" them as we try the
+        match.
+    """
+    for h in headers:
+        k = h.key.get_generator({})
+        if len(k) == len(val) and k[:].lower() == val.lower():
+            return h
+    return None
+
+
+class _HeaderMixin(object):
+    unique_name = None
+
+    def values(self, settings):
+        return (
+            self.key.get_generator(settings),
+            self.value.get_generator(settings),
+        )
+
+class _HTTP2Message(message.Message):
+    @property
+    def actions(self):
+        return []  # self.toks(actions._Action)
+
+    @property
+    def headers(self):
+        headers = self.toks(_HeaderMixin)
+
+        if not self.raw:
+            if not get_header("content-length", headers):
+                if not self.body:
+                    length = 0
+                else:
+                    length = len(self.body.string())
+                headers.append(
+                    Header(
+                        base.TokValueLiteral("content-length"),
+                        base.TokValueLiteral(str(length)),
+                    )
+                )
+        return headers
+
+    @property
+    def raw(self):
+        return bool(self.tok(Raw))
+
+    @property
+    def body(self):
+        return self.tok(Body)
+
+    def resolve(self, settings):
+        return self
+
+
+class Code(base.Integer):
+    pass
+
 
 class Method(base.OptionsOrValue):
     options = [
@@ -39,17 +100,41 @@ class Path(base.Value):
     pass
 
 
-class Header(base.KeyValue):
-    unique_name = None
+class Header(_HeaderMixin, base.KeyValue):
     preamble = "h"
 
+
+
+class ShortcutContentType(_HeaderMixin, base.Value):
+    preamble = "c"
+    key = base.TokValueLiteral("content-type")
+
+
+class ShortcutLocation(_HeaderMixin, base.Value):
+    preamble = "l"
+    key = base.TokValueLiteral("location")
+
+
+class ShortcutUserAgent(_HeaderMixin, base.OptionsOrValue):
+    preamble = "u"
+    options = [i[1] for i in http_uastrings.UASTRINGS]
+    key = base.TokValueLiteral("user-agent")
+
     def values(self, settings):
+        value = self.value.val
+        if self.option_used:
+            value = http_uastrings.get_by_shortcut(value.lower())[2]
+
         return (
             self.key.get_generator(settings),
-            self.value.get_generator(settings),
+            value
         )
 
 
+class Raw(base.CaselessLiteral):
+    TOK = "r"
+
+
 class Body(base.Value):
     preamble = "b"
 
@@ -58,13 +143,13 @@ class Times(base.Integer):
     preamble = "x"
 
 
-class Code(base.Integer):
-    pass
-
-
-class Request(message.Message):
+class Request(_HTTP2Message):
     comps = (
         Header,
+        ShortcutContentType,
+        ShortcutUserAgent,
+        Raw,
+        # NestedResponse,
         Body,
         Times,
     )
@@ -83,21 +168,9 @@ class Request(message.Message):
         return self.tok(Path)
 
     @property
-    def headers(self):
-        return self.toks(Header)
-
-    @property
-    def body(self):
-        return self.tok(Body)
-
-    @property
     def times(self):
         return self.tok(Times)
 
-    @property
-    def actions(self):
-        return []
-
     @classmethod
     def expr(cls):
         parts = [i.expr() for i in cls.comps]
@@ -113,9 +186,6 @@ class Request(message.Message):
         resp = resp.setParseAction(cls)
         return resp
 
-    def resolve(self, settings, msg=None):
-        return self
-
     def values(self, settings):
         if self.rendered_values:
             return self.rendered_values
@@ -129,7 +199,7 @@ class Request(message.Message):
             self.rendered_values = settings.protocol.create_request(
                 self.method.string(),
                 self.path.string(),
-                headers,  # TODO: parse that into a dict?!
+                headers,
                 body)
             return self.rendered_values
 
@@ -137,11 +207,14 @@ class Request(message.Message):
         return ":".join([i.spec() for i in self.tokens])
 
 
-class Response(message.Message):
+class Response(_HTTP2Message):
     unique_name = None
     comps = (
         Header,
         Body,
+        ShortcutContentType,
+        ShortcutLocation,
+        Raw,
     )
 
     def __init__(self, tokens):
@@ -153,21 +226,6 @@ class Response(message.Message):
     def code(self):
         return self.tok(Code)
 
-    @property
-    def headers(self):
-        return self.toks(Header)
-
-    @property
-    def body(self):
-        return self.tok(Body)
-
-    @property
-    def actions(self):
-        return []
-
-    def resolve(self, settings, msg=None):
-        return self
-
     @classmethod
     def expr(cls):
         parts = [i.expr() for i in cls.comps]
@@ -194,7 +252,7 @@ class Response(message.Message):
             self.rendered_values = settings.protocol.create_response(
                 self.code.string(),
                 self.stream_id,
-                headers,  # TODO: parse that into a dict?!
+                headers,
                 body)
             return self.rendered_values
 
diff --git a/test/test_language_http2.py b/test/test_language_http2.py
index a78c2bee..b25ad43c 100644
--- a/test/test_language_http2.py
+++ b/test/test_language_http2.py
@@ -90,6 +90,39 @@ class TestRequest:
             default_settings(),
         )
 
+    def test_raw_content_length(self):
+        r = parse_request('GET:/:r')
+        assert len(r.headers) == 0
+
+        r = parse_request('GET:/:r:b"foobar"')
+        assert len(r.headers) == 0
+
+        r = parse_request('GET:/')
+        assert len(r.headers) == 1
+        assert r.headers[0].values(default_settings()) == ("content-length", "0")
+
+        r = parse_request('GET:/:b"foobar"')
+        assert len(r.headers) == 1
+        assert r.headers[0].values(default_settings()) == ("content-length", "6")
+
+        r = parse_request('GET:/:b"foobar":h"content-length"="42"')
+        assert len(r.headers) == 1
+        assert r.headers[0].values(default_settings()) == ("content-length", "42")
+
+        r = parse_request('GET:/:r:b"foobar":h"content-length"="42"')
+        assert len(r.headers) == 1
+        assert r.headers[0].values(default_settings()) == ("content-length", "42")
+
+    def test_content_type(self):
+        r = parse_request('GET:/:r:c"foobar"')
+        assert len(r.headers) == 1
+        assert r.headers[0].values(default_settings()) == ("content-type", "foobar")
+
+    def test_user_agent(self):
+        r = parse_request('GET:/:r:ua')
+        assert len(r.headers) == 1
+        assert r.headers[0].values(default_settings()) == ("user-agent", netlib.http_uastrings.get_by_shortcut('a')[2])
+
     def test_render_with_headers(self):
         s = cStringIO.StringIO()
         r = parse_request('GET:/foo:h"foo"="bar"')
@@ -129,18 +162,27 @@ class TestResponse:
     def test_err(self):
         tutils.raises(language.ParseException, parse_response, 'GET:/')
 
-    def test_simple(self):
-        r = parse_response('200')
-        assert r.code.string() == "200"
+    def test_raw_content_length(self):
+        r = parse_response('200:r')
         assert len(r.headers) == 0
 
-        r = parse_response('200:h"foo"="bar"')
+        r = parse_response('200')
+        assert len(r.headers) == 1
+        assert r.headers[0].values(default_settings()) == ("content-length", "0")
+
+    def test_content_type(self):
+        r = parse_response('200:r:c"foobar"')
+        assert len(r.headers) == 1
+        assert r.headers[0].values(default_settings()) == ("content-type", "foobar")
+
+    def test_simple(self):
+        r = parse_response('200:r:h"foo"="bar"')
         assert r.code.string() == "200"
         assert len(r.headers) == 1
         assert r.headers[0].values(default_settings()) == ("foo", "bar")
         assert r.body is None
 
-        r = parse_response('200:h"foo"="bar":bfoobar:h"bla"="fasel"')
+        r = parse_response('200:r:h"foo"="bar":bfoobar:h"bla"="fasel"')
         assert r.code.string() == "200"
         assert len(r.headers) == 2
         assert r.headers[0].values(default_settings()) == ("foo", "bar")