mitmproxy/test/test_protocol_http2.py
2016-02-04 09:52:29 +01:00

432 lines
14 KiB
Python

from __future__ import (absolute_import, print_function, division)
import OpenSSL
import pytest
import traceback
import os
import tempfile
import sys
from libmproxy.proxy.config import ProxyConfig
from libmproxy.proxy.server import ProxyServer
from libmproxy.cmdline import APP_HOST, APP_PORT
import logging
logging.getLogger("hyper.packages.hpack.hpack").setLevel(logging.WARNING)
logging.getLogger("requests.packages.urllib3.connectionpool").setLevel(logging.WARNING)
logging.getLogger("passlib.utils.compat").setLevel(logging.WARNING)
logging.getLogger("passlib.registry").setLevel(logging.WARNING)
logging.getLogger("PIL.Image").setLevel(logging.WARNING)
logging.getLogger("PIL.PngImagePlugin").setLevel(logging.WARNING)
import netlib
from netlib import tservers as netlib_tservers
from netlib.utils import http2_read_raw_frame
import h2
from hyperframe.frame import Frame
from libmproxy import utils
from . import tservers
requires_alpn = pytest.mark.skipif(
not OpenSSL._util.lib.Cryptography_HAS_ALPN,
reason="requires OpenSSL with ALPN support")
class _Http2ServerBase(netlib_tservers.ServerTestBase):
ssl = dict(alpn_select=b'h2')
class handler(netlib.tcp.BaseHandler):
def handle(self):
h2_conn = h2.connection.H2Connection(client_side=False)
preamble = self.rfile.read(24)
h2_conn.initiate_connection()
h2_conn.receive_data(preamble)
self.wfile.write(h2_conn.data_to_send())
self.wfile.flush()
done = False
while not done:
try:
raw = b''.join(http2_read_raw_frame(self.rfile))
events = h2_conn.receive_data(raw)
except:
break
self.wfile.write(h2_conn.data_to_send())
self.wfile.flush()
for event in events:
try:
if not self.server.handle_server_event(event, h2_conn, self.rfile, self.wfile):
done = True
break
except Exception as e:
print(repr(e))
print(traceback.format_exc())
done = True
break
def handle_server_event(self, h2_conn, rfile, wfile):
raise NotImplementedError()
class _Http2TestBase(object):
@classmethod
def setup_class(self):
self.config = ProxyConfig(**self.get_proxy_config())
tmaster = tservers.TestMaster(self.config)
tmaster.start_app(APP_HOST, APP_PORT)
self.proxy = tservers.ProxyThread(tmaster)
self.proxy.start()
@classmethod
def teardown_class(cls):
cls.proxy.shutdown()
@property
def master(self):
return self.proxy.tmaster
@classmethod
def get_proxy_config(cls):
cls.cadir = os.path.join(tempfile.gettempdir(), "mitmproxy")
return dict(
no_upstream_cert = False,
cadir = cls.cadir,
authenticator = None,
)
def setup(self):
self.master.clear_log()
self.master.state.clear()
self.server.server.handle_server_event = self.handle_server_event
def _setup_connection(self):
self.config.http2 = True
client = netlib.tcp.TCPClient(("127.0.0.1", self.proxy.port))
client.connect()
# send CONNECT request
client.wfile.write(
b"CONNECT localhost:%d HTTP/1.1\r\n"
b"Host: localhost:%d\r\n"
b"\r\n" % (self.server.server.address.port, self.server.server.address.port)
)
client.wfile.flush()
# read CONNECT response
while client.rfile.readline() != "\r\n":
pass
client.convert_to_ssl(alpn_protos=[b'h2'])
h2_conn = h2.connection.H2Connection(client_side=True)
h2_conn.initiate_connection()
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
return client, h2_conn
def _send_request(self, wfile, h2_conn, stream_id=1, headers=[], body=b''):
h2_conn.send_headers(
stream_id=stream_id,
headers=headers,
end_stream=(len(body) == 0),
)
if body:
h2_conn.send_data(stream_id, body)
h2_conn.end_stream(stream_id)
wfile.write(h2_conn.data_to_send())
wfile.flush()
@requires_alpn
class TestSimple(_Http2TestBase, _Http2ServerBase):
@classmethod
def setup_class(self):
_Http2TestBase.setup_class()
_Http2ServerBase.setup_class()
@classmethod
def teardown_class(self):
_Http2TestBase.teardown_class()
_Http2ServerBase.teardown_class()
@classmethod
def handle_server_event(self, event, h2_conn, rfile, wfile):
if isinstance(event, h2.events.ConnectionTerminated):
return False
elif isinstance(event, h2.events.RequestReceived):
h2_conn.send_headers(1, [
(':status', '200'),
('foo', 'bar'),
])
h2_conn.send_data(1, b'foobar')
h2_conn.end_stream(1)
wfile.write(h2_conn.data_to_send())
wfile.flush()
return True
def test_simple(self):
client, h2_conn = self._setup_connection()
self._send_request(client.wfile, h2_conn, headers=[
(':authority', "127.0.0.1:%s" % self.server.server.address.port),
(':method', 'GET'),
(':scheme', 'https'),
(':path', '/'),
], body='my request body echoed back to me')
done = False
while not done:
try:
events = h2_conn.receive_data(b''.join(http2_read_raw_frame(client.rfile)))
except:
break
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
for event in events:
if isinstance(event, h2.events.StreamEnded):
done = True
h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
assert len(self.master.state.flows) == 1
assert self.master.state.flows[0].response.status_code == 200
assert self.master.state.flows[0].response.headers['foo'] == 'bar'
assert self.master.state.flows[0].response.body == b'foobar'
@requires_alpn
class TestWithBodies(_Http2TestBase, _Http2ServerBase):
tmp_data_buffer_foobar = b''
@classmethod
def setup_class(self):
_Http2TestBase.setup_class()
_Http2ServerBase.setup_class()
@classmethod
def teardown_class(self):
_Http2TestBase.teardown_class()
_Http2ServerBase.teardown_class()
@classmethod
def handle_server_event(self, event, h2_conn, rfile, wfile):
if isinstance(event, h2.events.ConnectionTerminated):
return False
if isinstance(event, h2.events.DataReceived):
self.tmp_data_buffer_foobar += event.data
elif isinstance(event, h2.events.StreamEnded):
h2_conn.send_headers(1, [
(':status', '200'),
])
h2_conn.send_data(1, self.tmp_data_buffer_foobar)
h2_conn.end_stream(1)
wfile.write(h2_conn.data_to_send())
wfile.flush()
return True
def test_with_bodies(self):
client, h2_conn = self._setup_connection()
self._send_request(
client.wfile,
h2_conn,
headers=[
(':authority', "127.0.0.1:%s" % self.server.server.address.port),
(':method', 'GET'),
(':scheme', 'https'),
(':path', '/'),
],
body='foobar with request body',
)
done = False
while not done:
try:
events = h2_conn.receive_data(b''.join(http2_read_raw_frame(client.rfile)))
except:
break
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
for event in events:
if isinstance(event, h2.events.StreamEnded):
done = True
h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
assert self.master.state.flows[0].response.body == b'foobar with request body'
@requires_alpn
class TestPushPromise(_Http2TestBase, _Http2ServerBase):
@classmethod
def setup_class(self):
_Http2TestBase.setup_class()
_Http2ServerBase.setup_class()
@classmethod
def teardown_class(self):
_Http2TestBase.teardown_class()
_Http2ServerBase.teardown_class()
@classmethod
def handle_server_event(self, event, h2_conn, rfile, wfile):
if isinstance(event, h2.events.ConnectionTerminated):
return False
elif isinstance(event, h2.events.RequestReceived):
if event.stream_id != 1:
# ignore requests initiated by push promises
return True
h2_conn.send_headers(1, [(':status', '200')])
h2_conn.push_stream(1, 2, [
(':authority', "127.0.0.1:%s" % self.port),
(':method', 'GET'),
(':scheme', 'https'),
(':path', '/pushed_stream_foo'),
('foo', 'bar')
])
h2_conn.push_stream(1, 4, [
(':authority', "127.0.0.1:%s" % self.port),
(':method', 'GET'),
(':scheme', 'https'),
(':path', '/pushed_stream_bar'),
('foo', 'bar')
])
wfile.write(h2_conn.data_to_send())
wfile.flush()
h2_conn.send_headers(2, [(':status', '200')])
h2_conn.send_headers(4, [(':status', '200')])
wfile.write(h2_conn.data_to_send())
wfile.flush()
h2_conn.send_data(1, b'regular_stream')
h2_conn.send_data(2, b'pushed_stream_foo')
h2_conn.send_data(4, b'pushed_stream_bar')
wfile.write(h2_conn.data_to_send())
wfile.flush()
h2_conn.end_stream(1)
h2_conn.end_stream(2)
h2_conn.end_stream(4)
wfile.write(h2_conn.data_to_send())
wfile.flush()
return True
def test_push_promise(self):
client, h2_conn = self._setup_connection()
self._send_request(client.wfile, h2_conn, stream_id=1, headers=[
(':authority', "127.0.0.1:%s" % self.server.server.address.port),
(':method', 'GET'),
(':scheme', 'https'),
(':path', '/'),
('foo', 'bar')
])
done = False
ended_streams = 0
pushed_streams = 0
responses = 0
while not done:
try:
raw = b''.join(http2_read_raw_frame(client.rfile))
events = h2_conn.receive_data(raw)
except:
break
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
for event in events:
if isinstance(event, h2.events.StreamEnded):
ended_streams += 1
elif isinstance(event, h2.events.PushedStreamReceived):
pushed_streams += 1
elif isinstance(event, h2.events.ResponseReceived):
responses += 1
if isinstance(event, h2.events.ConnectionTerminated):
done = True
if responses == 3 and ended_streams == 3 and pushed_streams == 2:
done = True
h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
assert ended_streams == 3
assert pushed_streams == 2
bodies = [flow.response.body for flow in self.master.state.flows]
assert len(bodies) == 3
assert b'regular_stream' in bodies
assert b'pushed_stream_foo' in bodies
assert b'pushed_stream_bar' in bodies
def test_push_promise_reset(self):
client, h2_conn = self._setup_connection()
self._send_request(client.wfile, h2_conn, stream_id=1, headers=[
(':authority', "127.0.0.1:%s" % self.server.server.address.port),
(':method', 'GET'),
(':scheme', 'https'),
(':path', '/'),
('foo', 'bar')
])
done = False
ended_streams = 0
pushed_streams = 0
responses = 0
while not done:
try:
events = h2_conn.receive_data(b''.join(http2_read_raw_frame(client.rfile)))
except:
break
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
for event in events:
if isinstance(event, h2.events.StreamEnded) and event.stream_id == 1:
ended_streams += 1
elif isinstance(event, h2.events.PushedStreamReceived):
pushed_streams += 1
h2_conn.reset_stream(event.pushed_stream_id, error_code=0x8)
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
elif isinstance(event, h2.events.ResponseReceived):
responses += 1
if isinstance(event, h2.events.ConnectionTerminated):
done = True
if responses >= 1 and ended_streams >= 1 and pushed_streams == 2:
done = True
h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send())
client.wfile.flush()
bodies = [flow.response.body for flow in self.master.state.flows if flow.response]
assert len(bodies) >= 1
assert b'regular_stream' in bodies
# the other two bodies might not be transmitted before the reset