Options unification: sizes

Start dealing with corner cases:

- Sizes are always stored in options as strings
- Add a new core addon that's responsible for verifying settings that don't
belong to an addon
- Add a _processed scratch space on the Options object for processed core
values to be stored in. This is pretty dirty, but less dirty than re-parsing
values every time. We'll come up with something better down the track.
This commit is contained in:
Aldo Cortesi 2017-03-06 08:59:35 +13:00
parent 45d18ac8cb
commit 201c65960e
12 changed files with 82 additions and 46 deletions

View File

@ -3,6 +3,7 @@ from mitmproxy.addons import anticomp
from mitmproxy.addons import check_alpn from mitmproxy.addons import check_alpn
from mitmproxy.addons import check_ca from mitmproxy.addons import check_ca
from mitmproxy.addons import clientplayback from mitmproxy.addons import clientplayback
from mitmproxy.addons import core
from mitmproxy.addons import disable_h2c_upgrade from mitmproxy.addons import disable_h2c_upgrade
from mitmproxy.addons import onboarding from mitmproxy.addons import onboarding
from mitmproxy.addons import proxyauth from mitmproxy.addons import proxyauth
@ -19,6 +20,7 @@ from mitmproxy.addons import upstream_auth
def default_addons(): def default_addons():
return [ return [
core.Core(),
anticache.AntiCache(), anticache.AntiCache(),
anticomp.AntiComp(), anticomp.AntiComp(),
check_alpn.CheckALPN(), check_alpn.CheckALPN(),

20
mitmproxy/addons/core.py Normal file
View File

@ -0,0 +1,20 @@
"""
The core addon is responsible for verifying core settings that are not
checked by other addons.
"""
from mitmproxy import exceptions
from mitmproxy.utils import human
class Core:
def configure(self, options, updated):
if "body_size_limit" in updated and options.body_size_limit:
try:
options._processed["body_size_limit"] = human.parse_size(
options.body_size_limit
)
except ValueError as e:
raise exceptions.OptionsError(
"Invalid body size limit specification: %s" %
options.body_size_limit
)

View File

@ -1,6 +1,7 @@
from mitmproxy.net.http import http1 from mitmproxy.net.http import http1
from mitmproxy import exceptions from mitmproxy import exceptions
from mitmproxy import ctx from mitmproxy import ctx
from mitmproxy.utils import human
class StreamBodies: class StreamBodies:
@ -8,7 +9,11 @@ class StreamBodies:
self.max_size = None self.max_size = None
def configure(self, options, updated): def configure(self, options, updated):
self.max_size = options.stream_large_bodies if "stream_large_bodies" in updated and options.stream_large_bodies:
try:
self.max_size = human.parse_size(options.stream_large_bodies)
except ValueError as e:
raise exceptions.OptionsError(e)
def run(self, f, is_request): def run(self, f, is_request):
if self.max_size: if self.max_size:

View File

@ -77,7 +77,14 @@ class Options(optmanager.OptManager):
self.add_option("server_replay", [], Sequence[str]) self.add_option("server_replay", [], Sequence[str])
self.add_option("stickycookie", None, Optional[str]) self.add_option("stickycookie", None, Optional[str])
self.add_option("stickyauth", None, Optional[str]) self.add_option("stickyauth", None, Optional[str])
self.add_option("stream_large_bodies", None, Optional[int]) self.add_option(
"stream_large_bodies", None, Optional[str],
"""
Stream data to the client if response body exceeds the given
threshold. If streamed, the body will not be stored in any way.
Understands k/m/g suffixes, i.e. 3m for 3 megabytes.
"""
)
self.add_option( self.add_option(
"verbosity", 2, int, "verbosity", 2, int,
"Log verbosity." "Log verbosity."
@ -109,7 +116,11 @@ class Options(optmanager.OptManager):
"Add all certificates of the upstream server to the certificate chain " "Add all certificates of the upstream server to the certificate chain "
"that will be served to the proxy client, as extras." "that will be served to the proxy client, as extras."
) )
self.add_option("body_size_limit", None, Optional[int]) self.add_option(
"body_size_limit", None, Optional[str],
"Byte size limit of HTTP request and response bodies."
" Understands k/m/g suffixes, i.e. 3m for 3 megabytes."
)
self.add_option("cadir", CA_DIR, str) self.add_option("cadir", CA_DIR, str)
self.add_option("certs", [], Sequence[Tuple[str, str]]) self.add_option("certs", [], Sequence[Tuple[str, str]])
self.add_option("ciphers_client", DEFAULT_CLIENT_CIPHERS, str) self.add_option("ciphers_client", DEFAULT_CLIENT_CIPHERS, str)

View File

@ -12,7 +12,6 @@ import ruamel.yaml
from mitmproxy import exceptions from mitmproxy import exceptions
from mitmproxy.utils import typecheck from mitmproxy.utils import typecheck
""" """
The base implementation for Options. The base implementation for Options.
""" """
@ -92,6 +91,7 @@ class OptManager:
self.__dict__["_options"] = {} self.__dict__["_options"] = {}
self.__dict__["changed"] = blinker.Signal() self.__dict__["changed"] = blinker.Signal()
self.__dict__["errored"] = blinker.Signal() self.__dict__["errored"] = blinker.Signal()
self.__dict__["_processed"] = {}
def add_option( def add_option(
self, self,
@ -330,7 +330,7 @@ class OptManager:
help=o.help help=o.help
) )
parser.set_defaults(**{option: o.default}) parser.set_defaults(**{option: o.default})
elif o.typespec == int: elif o.typespec in (int, typing.Optional[int]):
parser.add_argument( parser.add_argument(
"--%s" % f, "--%s" % f,
action="store", action="store",
@ -339,5 +339,14 @@ class OptManager:
help=o.help, help=o.help,
metavar=metavar metavar=metavar
) )
elif o.typespec in (str, typing.Optional[str]):
parser.add_argument(
"--%s" % f,
action="store",
type=str,
dest=option,
help=o.help,
metavar=metavar
)
else: else:
raise ValueError("Unsupported option type: %s", o.typespec) raise ValueError("Unsupported option type: %s", o.typespec)

View File

@ -19,7 +19,7 @@ class Http1Layer(httpbase._HttpTransmissionLayer):
return http1.read_body( return http1.read_body(
self.client_conn.rfile, self.client_conn.rfile,
expected_size, expected_size,
self.config.options.body_size_limit self.config.options._processed.get("body_size_limit")
) )
def send_request(self, request): def send_request(self, request):
@ -35,7 +35,7 @@ class Http1Layer(httpbase._HttpTransmissionLayer):
return http1.read_body( return http1.read_body(
self.server_conn.rfile, self.server_conn.rfile,
expected_size, expected_size,
self.config.options.body_size_limit self.config.options._processed.get("body_size_limit")
) )
def send_response_headers(self, response): def send_response_headers(self, response):

View File

@ -183,7 +183,7 @@ class Http2Layer(base.Layer):
return True return True
def _handle_data_received(self, eid, event, source_conn): def _handle_data_received(self, eid, event, source_conn):
bsl = self.config.options.body_size_limit bsl = self.config.options._processed.get("body_size_limit")
if bsl and self.streams[eid].queued_data_length > bsl: if bsl and self.streams[eid].queued_data_length > bsl:
self.streams[eid].kill() self.streams[eid].kill()
self.connections[source_conn].safe_reset_stream( self.connections[source_conn].safe_reset_stream(

View File

@ -31,6 +31,7 @@ class RequestReplayThread(basethread.BaseThread):
def run(self): def run(self):
r = self.f.request r = self.f.request
bsl = self.config.options._processed.get("body_size_limit")
first_line_format_backup = r.first_line_format first_line_format_backup = r.first_line_format
server = None server = None
try: try:
@ -55,7 +56,7 @@ class RequestReplayThread(basethread.BaseThread):
resp = http1.read_response( resp = http1.read_response(
server.rfile, server.rfile,
connect_request, connect_request,
body_size_limit=self.config.options.body_size_limit body_size_limit=bsl
) )
if resp.status_code != 200: if resp.status_code != 200:
raise exceptions.ReplayException("Upstream server refuses CONNECT request") raise exceptions.ReplayException("Upstream server refuses CONNECT request")
@ -87,7 +88,7 @@ class RequestReplayThread(basethread.BaseThread):
http1.read_response( http1.read_response(
server.rfile, server.rfile,
r, r,
body_size_limit=self.config.options.body_size_limit body_size_limit=bsl
) )
) )
if self.channel: if self.channel:

View File

@ -4,7 +4,6 @@ import os
from mitmproxy import exceptions from mitmproxy import exceptions
from mitmproxy import options from mitmproxy import options
from mitmproxy import platform from mitmproxy import platform
from mitmproxy.utils import human
from mitmproxy.net import tcp from mitmproxy.net import tcp
from mitmproxy import version from mitmproxy import version
from mitmproxy.addons import view from mitmproxy.addons import view
@ -25,10 +24,6 @@ def get_common_options(args):
if args.stickyauth_filt: if args.stickyauth_filt:
stickyauth = args.stickyauth_filt stickyauth = args.stickyauth_filt
stream_large_bodies = args.stream_large_bodies
if stream_large_bodies:
stream_large_bodies = human.parse_size(stream_large_bodies)
if args.streamfile and args.streamfile[0] == args.rfile: if args.streamfile and args.streamfile[0] == args.rfile:
if args.streamfile[1] == "wb": if args.streamfile[1] == "wb":
raise exceptions.OptionsError( raise exceptions.OptionsError(
@ -49,15 +44,6 @@ def get_common_options(args):
parts = ["*", parts[0]] parts = ["*", parts[0]]
certs.append(parts) certs.append(parts)
body_size_limit = args.body_size_limit
if body_size_limit:
try:
body_size_limit = human.parse_size(body_size_limit)
except ValueError as e:
raise exceptions.OptionsError(
"Invalid body size limit specification: %s" % body_size_limit
)
# Establish proxy mode # Establish proxy mode
c = 0 c = 0
mode, upstream_server = "regular", None mode, upstream_server = "regular", None
@ -117,7 +103,7 @@ def get_common_options(args):
scripts=args.scripts, scripts=args.scripts,
stickycookie=stickycookie, stickycookie=stickycookie,
stickyauth=stickyauth, stickyauth=stickyauth,
stream_large_bodies=stream_large_bodies, stream_large_bodies=args.stream_large_bodies,
showhost=args.showhost, showhost=args.showhost,
streamfile=args.streamfile[0] if args.streamfile else None, streamfile=args.streamfile[0] if args.streamfile else None,
streamfile_append=True if args.streamfile and args.streamfile[1] == "a" else False, streamfile_append=True if args.streamfile and args.streamfile[1] == "a" else False,
@ -132,7 +118,7 @@ def get_common_options(args):
auth_singleuser = args.auth_singleuser, auth_singleuser = args.auth_singleuser,
auth_htpasswd = args.auth_htpasswd, auth_htpasswd = args.auth_htpasswd,
add_upstream_certs_to_client_chain = args.add_upstream_certs_to_client_chain, add_upstream_certs_to_client_chain = args.add_upstream_certs_to_client_chain,
body_size_limit = body_size_limit, body_size_limit = args.body_size_limit,
cadir = args.cadir, cadir = args.cadir,
certs = certs, certs = certs,
ciphers_client = args.ciphers_client, ciphers_client = args.ciphers_client,
@ -229,23 +215,8 @@ def basic_options(parser, opts):
help="Append flows to file." help="Append flows to file."
) )
opts.make_parser(parser, "anticomp") opts.make_parser(parser, "anticomp")
parser.add_argument( opts.make_parser(parser, "body_size_limit", metavar="SIZE")
"-Z", "--body-size-limit", opts.make_parser(parser, "stream_large_bodies")
action="store", dest="body_size_limit",
metavar="SIZE",
help="Byte size limit of HTTP request and response bodies."
" Understands k/m/g suffixes, i.e. 3m for 3 megabytes."
)
parser.add_argument(
"--stream",
action="store", dest="stream_large_bodies",
metavar="SIZE",
help="""
Stream data to the client if response body exceeds the given
threshold. If streamed, the body will not be stored in any way.
Understands k/m/g suffixes, i.e. 3m for 3 megabytes.
"""
)
def proxy_modes(parser, opts): def proxy_modes(parser, opts):

View File

@ -0,0 +1,13 @@
from mitmproxy import exceptions
from mitmproxy.addons import core
from mitmproxy.test import taddons
import pytest
def test_simple():
sa = core.Core()
with taddons.context() as tctx:
with pytest.raises(exceptions.OptionsError):
tctx.configure(sa, body_size_limit = "invalid")
tctx.configure(sa, body_size_limit = "1m")
assert tctx.options._processed["body_size_limit"]

View File

@ -1,13 +1,16 @@
from mitmproxy import exceptions
from mitmproxy.test import tflow from mitmproxy.test import tflow
from mitmproxy.test import taddons from mitmproxy.test import taddons
from mitmproxy.addons import streambodies from mitmproxy.addons import streambodies
import pytest
def test_simple(): def test_simple():
sa = streambodies.StreamBodies() sa = streambodies.StreamBodies()
with taddons.context() as tctx: with taddons.context() as tctx:
tctx.configure(sa, stream_large_bodies = 10) with pytest.raises(exceptions.OptionsError):
tctx.configure(sa, stream_large_bodies = "invalid")
tctx.configure(sa, stream_large_bodies = "10")
f = tflow.tflow() f = tflow.tflow()
f.request.content = b"" f.request.content = b""

View File

@ -499,7 +499,8 @@ class TestBodySizeLimit(_Http2Test):
return True return True
def test_body_size_limit(self): def test_body_size_limit(self):
self.config.options.body_size_limit = 20 self.config.options.body_size_limit = "20"
self.config.options._processed["body_size_limit"] = 20
client, h2_conn = self._setup_connection() client, h2_conn = self._setup_connection()