adjust http1-related code

7 files changed, 45 insertions, 20 deletions

diff --git a/netlib/http/__init__.py b/netlib/http/__init__.py
index 9303de09..d72884b3 100644
--- a/netlib/http/__init__.py
+++ b/netlib/http/__init__.py
@@ -1,9 +1,12 @@
+from __future__ import absolute_import, print_function, division
 from .models import Request, Response, Headers
+from .models import ALPN_PROTO_HTTP1, ALPN_PROTO_H2
 from .models import HDR_FORM_MULTIPART, HDR_FORM_URLENCODED, CONTENT_MISSING
 from . import http1, http2
 
 __all__ = [
     "Request", "Response", "Headers",
+    "ALPN_PROTO_HTTP1", "ALPN_PROTO_H2",
     "HDR_FORM_MULTIPART", "HDR_FORM_URLENCODED", "CONTENT_MISSING",
-    "http1", "http2"
+    "http1", "http2",
 ]
diff --git a/netlib/http/http1/__init__.py b/netlib/http/http1/__init__.py
index a72c2e05..2d33ff8a 100644
--- a/netlib/http/http1/__init__.py
+++ b/netlib/http/http1/__init__.py
@@ -1,3 +1,4 @@
+from __future__ import absolute_import, print_function, division
 from .read import (
     read_request, read_request_head,
     read_response, read_response_head,
diff --git a/netlib/http/http1/assemble.py b/netlib/http/http1/assemble.py
index 47c7e95a..ace25d79 100644
--- a/netlib/http/http1/assemble.py
+++ b/netlib/http/http1/assemble.py
@@ -25,9 +25,9 @@ def assemble_response(response):
     return head + response.body
 
 
-def assemble_response_head(response):
+def assemble_response_head(response, preserve_transfer_encoding=False):
     first_line = _assemble_response_line(response)
-    headers = _assemble_response_headers(response)
+    headers = _assemble_response_headers(response, preserve_transfer_encoding)
     return b"%s\r\n%s\r\n" % (first_line, headers)
 
 
diff --git a/netlib/http/http1/read.py b/netlib/http/http1/read.py
index 4c423c4c..62025d15 100644
--- a/netlib/http/http1/read.py
+++ b/netlib/http/http1/read.py
@@ -6,8 +6,7 @@ import re
 from ... import utils
 from ...exceptions import HttpReadDisconnect, HttpSyntaxException, HttpException
 from .. import Request, Response, Headers
-
-ALPN_PROTO_HTTP1 = b'http/1.1'
+from netlib.tcp import NetLibDisconnect
 
 
 def read_request(rfile, body_size_limit=None):
@@ -157,10 +156,10 @@ def connection_close(http_version, headers):
 
     # If we don't have a Connection header, HTTP 1.1 connections are assumed to
     # be persistent
-    return http_version != (1, 1)
+    return http_version != b"HTTP/1.1"
 
 
-def expected_http_body_size(request, response=False):
+def expected_http_body_size(request, response=None):
     """
         Returns:
             The expected body length:
@@ -211,10 +210,13 @@ def expected_http_body_size(request, response=False):
 
 
 def _get_first_line(rfile):
-    line = rfile.readline()
-    if line == b"\r\n" or line == b"\n":
-        # Possible leftover from previous message
+    try:
         line = rfile.readline()
+        if line == b"\r\n" or line == b"\n":
+            # Possible leftover from previous message
+            line = rfile.readline()
+    except NetLibDisconnect:
+        raise HttpReadDisconnect()
     if not line:
         raise HttpReadDisconnect()
     line = line.strip()
@@ -317,6 +319,8 @@ def _read_headers(rfile):
             try:
                 name, value = line.split(b":", 1)
                 value = value.strip()
+                if not name or not value:
+                    raise ValueError()
                 ret.append([name, value])
             except ValueError:
                 raise HttpSyntaxException("Invalid headers")
diff --git a/netlib/http/http2/__init__.py b/netlib/http/http2/__init__.py
index e69de29b..7043d36f 100644
--- a/netlib/http/http2/__init__.py
+++ b/netlib/http/http2/__init__.py
@@ -0,0 +1,6 @@
+from __future__ import absolute_import, print_function, division
+from .connections import HTTP2Protocol
+
+__all__ = [
+    "HTTP2Protocol"
+]
diff --git a/netlib/http/http2/connections.py b/netlib/http/http2/connections.py
index 036bf68f..5220d5d2 100644
--- a/netlib/http/http2/connections.py
+++ b/netlib/http/http2/connections.py
@@ -3,8 +3,8 @@ import itertools
 import time
 
 from hpack.hpack import Encoder, Decoder
-from netlib import http, utils
-from netlib.http import models as semantics
+from ... import utils
+from .. import Headers, Response, Request, ALPN_PROTO_H2
 from . import frame
 
 
@@ -36,8 +36,6 @@ class HTTP2Protocol(object):
 
     CLIENT_CONNECTION_PREFACE = "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"
 
-    ALPN_PROTO_H2 = 'h2'
-
     def __init__(
         self,
         tcp_handler=None,
@@ -62,6 +60,7 @@ class HTTP2Protocol(object):
 
     def read_request(
         self,
+        __rfile,
         include_body=True,
         body_size_limit=None,
         allow_empty=False,
@@ -111,7 +110,7 @@ class HTTP2Protocol(object):
             port = 80 if scheme == 'http' else 443
         port = int(port)
 
-        request = http.Request(
+        request = Request(
             form_in,
             method,
             scheme,
@@ -131,6 +130,7 @@ class HTTP2Protocol(object):
 
     def read_response(
         self,
+        __rfile,
         request_method='',
         body_size_limit=None,
         include_body=True,
@@ -159,7 +159,7 @@ class HTTP2Protocol(object):
         else:
             timestamp_end = None
 
-        response = http.Response(
+        response = Response(
             (2, 0),
             int(headers.get(':status', 502)),
             "",
@@ -172,8 +172,16 @@ class HTTP2Protocol(object):
 
         return response
 
+    def assemble(self, message):
+        if isinstance(message, Request):
+            return self.assemble_request(message)
+        elif isinstance(message, Response):
+            return self.assemble_response(message)
+        else:
+            raise ValueError("HTTP message not supported.")
+
     def assemble_request(self, request):
-        assert isinstance(request, semantics.Request)
+        assert isinstance(request, Request)
 
         authority = self.tcp_handler.sni if self.tcp_handler.sni else self.tcp_handler.address.host
         if self.tcp_handler.address.port != 443:
@@ -200,7 +208,7 @@ class HTTP2Protocol(object):
             self._create_body(request.body, stream_id)))
 
     def assemble_response(self, response):
-        assert isinstance(response, semantics.Response)
+        assert isinstance(response, Response)
 
         headers = response.headers.copy()
 
@@ -275,7 +283,7 @@ class HTTP2Protocol(object):
 
     def check_alpn(self):
         alp = self.tcp_handler.get_alpn_proto_negotiated()
-        if alp != self.ALPN_PROTO_H2:
+        if alp != ALPN_PROTO_H2:
             raise NotImplementedError(
                 "HTTP2Protocol can not handle unknown ALP: %s" % alp)
         return True
@@ -405,7 +413,7 @@ class HTTP2Protocol(object):
             else:
                 self._handle_unexpected_frame(frm)
 
-        headers = http.Headers(
+        headers = Headers(
             [[str(k), str(v)] for k, v in self.decoder.decode(header_block_fragment)]
         )
 
diff --git a/netlib/http/models.py b/netlib/http/models.py
index 572d66c9..2d09535c 100644
--- a/netlib/http/models.py
+++ b/netlib/http/models.py
@@ -13,6 +13,9 @@ try:
 except ImportError:
     from collections.abc import MutableMapping
 
+# TODO: Move somewhere else?
+ALPN_PROTO_HTTP1 = b'http/1.1'
+ALPN_PROTO_H2 = b'h2'
 HDR_FORM_URLENCODED = b"application/x-www-form-urlencoded"
 HDR_FORM_MULTIPART = b"multipart/form-data"