From 1b09002edc984d4ead3cce118ed583c3ceca0b99 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 31 Jul 2017 00:04:46 +0200 Subject: [PATCH] remove OptManager._processed Instead of having the core addon do postprocessing on body_size_limit, we add a cache to the parsing function. First, this avoids any potential issues with options and _processed getting out of sync. As anecdotal evidence, the previous implementation did not clear _processed when body_size_limit was reset to None. Second, it achieves the same end result without introducing a new concept of a "_processed" scratch space. Third, it works even if addons aren't present, and does not require workarounds as previously present in test_http2.py. refs https://github.com/mitmproxy/mitmproxy/pull/2484#pullrequestreview-53101507 --- mitmproxy/addons/core_option_validation.py | 6 ++---- mitmproxy/optmanager.py | 1 - mitmproxy/proxy/protocol/http1.py | 5 +++-- mitmproxy/proxy/protocol/http2.py | 3 ++- mitmproxy/proxy/protocol/http_replay.py | 3 ++- mitmproxy/utils/human.py | 11 ++++++++++- test/mitmproxy/addons/test_core_option_validation.py | 1 - test/mitmproxy/proxy/protocol/test_http2.py | 3 --- test/mitmproxy/utils/test_human.py | 1 + 9 files changed, 20 insertions(+), 14 deletions(-) diff --git a/mitmproxy/addons/core_option_validation.py b/mitmproxy/addons/core_option_validation.py index baeee7646..42da0b748 100644 --- a/mitmproxy/addons/core_option_validation.py +++ b/mitmproxy/addons/core_option_validation.py @@ -19,11 +19,9 @@ class CoreOptionValidation: "then the upstream certificate is not retrieved before generating " "the client certificate chain." ) - if "body_size_limit" in updated and opts.body_size_limit: + if "body_size_limit" in updated: try: - opts._processed["body_size_limit"] = human.parse_size( - opts.body_size_limit - ) + human.parse_size(opts.body_size_limit) except ValueError as e: raise exceptions.OptionsError( "Invalid body size limit specification: %s" % diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py index c28ec685f..601301753 100644 --- a/mitmproxy/optmanager.py +++ b/mitmproxy/optmanager.py @@ -94,7 +94,6 @@ class OptManager: self.__dict__["_options"] = {} self.__dict__["changed"] = blinker.Signal() self.__dict__["errored"] = blinker.Signal() - self.__dict__["_processed"] = {} def add_option( self, diff --git a/mitmproxy/proxy/protocol/http1.py b/mitmproxy/proxy/protocol/http1.py index 84cd63249..91f1e9b7d 100644 --- a/mitmproxy/proxy/protocol/http1.py +++ b/mitmproxy/proxy/protocol/http1.py @@ -1,6 +1,7 @@ from mitmproxy import http from mitmproxy.proxy.protocol import http as httpbase from mitmproxy.net.http import http1 +from mitmproxy.utils import human class Http1Layer(httpbase._HttpTransmissionLayer): @@ -19,7 +20,7 @@ class Http1Layer(httpbase._HttpTransmissionLayer): return http1.read_body( self.client_conn.rfile, expected_size, - self.config.options._processed.get("body_size_limit") + human.parse_size(self.config.options.body_size_limit) ) def send_request_headers(self, request): @@ -45,7 +46,7 @@ class Http1Layer(httpbase._HttpTransmissionLayer): return http1.read_body( self.server_conn.rfile, expected_size, - self.config.options._processed.get("body_size_limit") + human.parse_size(self.config.options.body_size_limit) ) def send_response_headers(self, response): diff --git a/mitmproxy/proxy/protocol/http2.py b/mitmproxy/proxy/protocol/http2.py index eab5292f7..cf021291f 100644 --- a/mitmproxy/proxy/protocol/http2.py +++ b/mitmproxy/proxy/protocol/http2.py @@ -17,6 +17,7 @@ import mitmproxy.net.http from mitmproxy.net import tcp from mitmproxy.types import basethread from mitmproxy.net.http import http2, headers +from mitmproxy.utils import human class SafeH2Connection(connection.H2Connection): @@ -183,7 +184,7 @@ class Http2Layer(base.Layer): return True def _handle_data_received(self, eid, event, source_conn): - bsl = self.config.options._processed.get("body_size_limit") + bsl = human.parse_size(self.config.options.body_size_limit) if bsl and self.streams[eid].queued_data_length > bsl: self.streams[eid].kill() self.connections[source_conn].safe_reset_stream( diff --git a/mitmproxy/proxy/protocol/http_replay.py b/mitmproxy/proxy/protocol/http_replay.py index 915c71c40..fd673a6f5 100644 --- a/mitmproxy/proxy/protocol/http_replay.py +++ b/mitmproxy/proxy/protocol/http_replay.py @@ -12,6 +12,7 @@ from mitmproxy import connections from mitmproxy.net import server_spec from mitmproxy.net.http import http1 from mitmproxy.types import basethread +from mitmproxy.utils import human # TODO: Doesn't really belong into mitmproxy.proxy.protocol... @@ -44,7 +45,7 @@ class RequestReplayThread(basethread.BaseThread): def run(self): r = self.f.request - bsl = self.options._processed.get("body_size_limit") + bsl = human.parse_size(self.options.body_size_limit) first_line_format_backup = r.first_line_format server = None try: diff --git a/mitmproxy/utils/human.py b/mitmproxy/utils/human.py index d67fb3104..e2e3142ac 100644 --- a/mitmproxy/utils/human.py +++ b/mitmproxy/utils/human.py @@ -1,6 +1,8 @@ import datetime import ipaddress import time +import functools +import typing SIZE_TABLE = [ ("b", 1024 ** 0), @@ -25,7 +27,14 @@ def pretty_size(size): return "%s%s" % (size, SIZE_TABLE[0][0]) -def parse_size(s): +@functools.lru_cache() +def parse_size(s: typing.Optional[str]) -> typing.Optional[int]: + """ + Parse a size with an optional k/m/... suffix. + Invalid values raise a ValueError. For added convenience, passing `None` returns `None`. + """ + if s is None: + return None try: return int(s) except ValueError: diff --git a/test/mitmproxy/addons/test_core_option_validation.py b/test/mitmproxy/addons/test_core_option_validation.py index 6d6d5ba41..cd5d4dfa7 100644 --- a/test/mitmproxy/addons/test_core_option_validation.py +++ b/test/mitmproxy/addons/test_core_option_validation.py @@ -11,7 +11,6 @@ def test_simple(): with pytest.raises(exceptions.OptionsError): tctx.configure(sa, body_size_limit = "invalid") tctx.configure(sa, body_size_limit = "1m") - assert tctx.master.options._processed["body_size_limit"] with pytest.raises(exceptions.OptionsError, match="mutually exclusive"): tctx.configure( diff --git a/test/mitmproxy/proxy/protocol/test_http2.py b/test/mitmproxy/proxy/protocol/test_http2.py index 5e6fa701c..4f161ef57 100644 --- a/test/mitmproxy/proxy/protocol/test_http2.py +++ b/test/mitmproxy/proxy/protocol/test_http2.py @@ -507,9 +507,6 @@ class TestBodySizeLimit(_Http2Test): def test_body_size_limit(self): self.options.body_size_limit = "20" - # FIXME: This should not be required? - self.options._processed["body_size_limit"] = 20 - h2_conn = self.setup_connection() self._send_request( diff --git a/test/mitmproxy/utils/test_human.py b/test/mitmproxy/utils/test_human.py index 76dc2f887..e8ffaad4a 100644 --- a/test/mitmproxy/utils/test_human.py +++ b/test/mitmproxy/utils/test_human.py @@ -22,6 +22,7 @@ def test_parse_size(): human.parse_size("1f") with pytest.raises(ValueError): human.parse_size("ak") + assert human.parse_size(None) is None def test_pretty_size():