Fixing HTTP upgrade functionality

Author: fredthomsen

hyper/common/connection.py

@@ -55,7 +55,7 @@ def __init__(self,
         self._h1_kwargs = {'secure': secure, 'ssl_context': ssl_context}
         self._h2_kwargs = {
             'window_manager': window_manager, 'enable_push': enable_push,
-            'ssl_context': ssl_context
+            'secure': secure, 'ssl_context': ssl_context
         }
 
         # Add any unexpected kwargs to both dictionaries.
@@ -87,14 +87,11 @@ def request(self, method, url, body=None, headers={}):
             return self._conn.request(
                 method=method, url=url, body=body, headers=headers
             )
-        except (TLSUpgrade, HTTPUpgrade) as e:
-            # We upgraded in the NPN/ALPN handshake or via the HTTP Upgrade 
-            # mechanism. We can just go straight to the world of HTTP/2. 
-            #Replace the backing object and insert the socket into it.
-            if(type(e) is TLSUpgrade):
-                assert e.negotiated in H2_NPN_PROTOCOLS
-            else:
-                assert e.negotiated == H2C_PROTOCOL
+        except TLSUpgrade as e:
+            # We upgraded in the NPN/ALPN handshake. We can just go straight to
+            # the world of HTTP/2. Replace the backing object and insert the
+            # socket into it.
+            assert e.negotiated in H2_NPN_PROTOCOLS
 
             self._conn = HTTP20Connection(
                 self._host, self._port, **self._h2_kwargs
@@ -109,6 +106,31 @@ def request(self, method, url, body=None, headers={}):
                 method=method, url=url, body=body, headers=headers
             )
 
+    def get_response(self):
+        """
+        Returns a response object.
+        """
+        try:
+            return self._conn.get_response()
+        except HTTPUpgrade as e:
+            # We upgraded via the HTTP Upgrade mechanism. We can just 
+            #go straight to the world of HTTP/2. Replace the backing object 
+            #and insert the socket into it.
+            assert e.negotiated == H2C_PROTOCOL
+    
+            self._conn = HTTP20Connection(
+                self._host, self._port, **self._h2_kwargs
+            )
+
+            #stream id 1 is used by the upgrade request and response
+            self.next_stream_id += 2
+            self._conn._sock = e.sock
+            
+            # HTTP/2 preamble must be sent after receipt of a HTTP/1.1 101
+            self._conn._send_preamble()
+
+            return e.resp
+
     # Can anyone say 'proxy object pattern'?
     def __getattr__(self, name):
         return getattr(self._conn, name)

hyper/common/exceptions.py

@@ -56,7 +56,8 @@ class HTTPUpgrade(Exception):
     """
     We upgraded to a new protocol via the HTTP Upgrade response.
     """
-    def __init__(self, negotiated, sock):
+    def __init__(self, negotiated, sock, resp):
         super(HTTPUpgrade, self).__init__()
         self.negotiated = negotiated
         self.sock = sock
+        self.resp = resp

hyper/http11/connection.py

@@ -19,6 +19,7 @@
 from ..compat import bytes
 
 from ..http20.connection import H2C_PROTOCOL
+from ..http20.response import HTTP20Response
 from ..packages.hyperframe.frame import SettingsFrame
 
 # We prefer pycohttpparser to the pure-Python interpretation
@@ -138,9 +139,9 @@ def request(self, method, url, body=None, headers={}):
         headers[b'upgrade'] = H2C_PROTOCOL
 
         #need to encode SETTINGS frame payload in Base64 and put into the HTTP-2 Settings header 
-        http2Settings = SettingsFrame(0)
-        http2Settings.settings[SettingsFrame.INITIAL_WINDOW_SIZE] = 65535
-        headers[b'HTTP2-Settings'] = base64.b64encode(http2Settings.serialize_body())
+        http2_settings = SettingsFrame(0)
+        http2_settings.settings[SettingsFrame.INITIAL_WINDOW_SIZE] = 65535
+        headers[b'HTTP2-Settings'] = base64.b64encode(http2_settings.serialize_body())
 
         # We may need extra headers.
         if body:
@@ -179,9 +180,10 @@ def get_response(self):
         self._sock.advance_buffer(response.consumed)
 
         if(response.status == 101 and 
-           response.getheader('Connection') == 'Upgrade' and 
-           response.getheader('Upgrade') == H2C_PROTOCOL):
-            raise HTTPUpgrade(proto, sock)
+           b'upgrade' in headers['connection'] and bytes(H2C_PROTOCOL, 'utf-8') in headers['upgrade']):
+            # HTTP/1.1 requests that are upgrade to HTTP/2.0 are responded to with steam id of 1
+            headers[b':status'] = str(response.status)
+            raise HTTPUpgrade(H2C_PROTOCOL, self._sock, HTTP20Response(headers, ))
 
         return HTTP11Response(
             response.status,

hyper/http20/connection.py

@@ -287,7 +287,7 @@ def putrequest(self, method, selector, **kwargs):
         # HTTP/2 specific. These are: ":method", ":scheme", ":authority" and
         # ":path". We can set all of these now.
         s.add_header(":method", method)
-        s.add_header(":scheme", "https")  # We only support HTTPS at this time.
+        s.add_header(":scheme", "https" if self.secure else "http") 
         s.add_header(":authority", self.host)
         s.add_header(":path", selector)