From d2268ddb1ebfca2b3314c3bde12d3a356fd6a114 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Mon, 18 Jul 2016 16:58:53 +1200 Subject: [PATCH] ProxyConfig: --body-size-limit to Options --- mitmproxy/cmdline.py | 6 +++++- mitmproxy/flow/options.py | 2 ++ mitmproxy/protocol/http1.py | 16 +++++++++++++--- mitmproxy/protocol/http2.py | 15 +++++++++++---- mitmproxy/protocol/http_replay.py | 4 ++-- mitmproxy/proxy/config.py | 15 +++------------ test/mitmproxy/test_protocol_http2.py | 2 +- 7 files changed, 37 insertions(+), 23 deletions(-) diff --git a/mitmproxy/cmdline.py b/mitmproxy/cmdline.py index 9cf8be3cb..3b54f0baf 100644 --- a/mitmproxy/cmdline.py +++ b/mitmproxy/cmdline.py @@ -184,7 +184,6 @@ def get_common_options(args): "That would trigger an infinite loop." ) - # Proxy config certs = [] for i in args.certs: @@ -193,6 +192,10 @@ def get_common_options(args): parts = ["*", parts[0]] certs.append(parts) + body_size_limit = args.body_size_limit + if body_size_limit: + body_size_limit = human.parse_size(body_size_limit) + return dict( app=args.app, app_host=args.app_host, @@ -222,6 +225,7 @@ def get_common_options(args): replay_ignore_payload_params=args.replay_ignore_payload_params, replay_ignore_host=args.replay_ignore_host, + body_size_limit = body_size_limit, cadir = args.cadir, certs = certs, clientcerts = args.clientcerts, diff --git a/mitmproxy/flow/options.py b/mitmproxy/flow/options.py index 51672f9c0..c8839399b 100644 --- a/mitmproxy/flow/options.py +++ b/mitmproxy/flow/options.py @@ -39,6 +39,7 @@ class Options(options.Options): replay_ignore_host=False, # type: bool # Proxy options + body_size_limit=None, # type: Optional[int] cadir = cmdline.CA_DIR, # type: str certs = (), # type: Sequence[Tuple[str, str]] clientcerts = None, # type: Optional[str] @@ -76,6 +77,7 @@ class Options(options.Options): self.replay_ignore_host = replay_ignore_host # Proxy options + self.body_size_limit = body_size_limit self.cadir = cadir self.certs = certs self.clientcerts = clientcerts diff --git a/mitmproxy/protocol/http1.py b/mitmproxy/protocol/http1.py index 7055a7fd8..8698fe312 100644 --- a/mitmproxy/protocol/http1.py +++ b/mitmproxy/protocol/http1.py @@ -12,12 +12,18 @@ class Http1Layer(http._HttpTransmissionLayer): self.mode = mode def read_request(self): - req = http1.read_request(self.client_conn.rfile, body_size_limit=self.config.body_size_limit) + req = http1.read_request( + self.client_conn.rfile, body_size_limit=self.config.options.body_size_limit + ) return models.HTTPRequest.wrap(req) def read_request_body(self, request): expected_size = http1.expected_http_body_size(request) - return http1.read_body(self.client_conn.rfile, expected_size, self.config.body_size_limit) + return http1.read_body( + self.client_conn.rfile, + expected_size, + self.config.options.body_size_limit + ) def send_request(self, request): self.server_conn.wfile.write(http1.assemble_request(request)) @@ -29,7 +35,11 @@ class Http1Layer(http._HttpTransmissionLayer): def read_response_body(self, request, response): expected_size = http1.expected_http_body_size(request, response) - return http1.read_body(self.server_conn.rfile, expected_size, self.config.body_size_limit) + return http1.read_body( + self.server_conn.rfile, + expected_size, + self.config.options.body_size_limit + ) def send_response_headers(self, response): raw = http1.assemble_response_head(response) diff --git a/mitmproxy/protocol/http2.py b/mitmproxy/protocol/http2.py index ee66393fa..1285e10ee 100644 --- a/mitmproxy/protocol/http2.py +++ b/mitmproxy/protocol/http2.py @@ -183,14 +183,21 @@ class Http2Layer(base.Layer): return True def _handle_data_received(self, eid, event, source_conn): - if self.config.body_size_limit and self.streams[eid].queued_data_length > self.config.body_size_limit: + bsl = self.config.options.body_size_limit + if bsl and self.streams[eid].queued_data_length > bsl: self.streams[eid].zombie = time.time() - source_conn.h2.safe_reset_stream(event.stream_id, h2.errors.REFUSED_STREAM) - self.log("HTTP body too large. Limit is {}.".format(self.config.body_size_limit), "info") + source_conn.h2.safe_reset_stream( + event.stream_id, + h2.errors.REFUSED_STREAM + ) + self.log("HTTP body too large. Limit is {}.".format(bsl), "info") else: self.streams[eid].data_queue.put(event.data) self.streams[eid].queued_data_length += len(event.data) - source_conn.h2.safe_increment_flow_control(event.stream_id, event.flow_controlled_length) + source_conn.h2.safe_increment_flow_control( + event.stream_id, + event.flow_controlled_length + ) return True def _handle_stream_ended(self, eid): diff --git a/mitmproxy/protocol/http_replay.py b/mitmproxy/protocol/http_replay.py index 27dfb741c..ba58075e1 100644 --- a/mitmproxy/protocol/http_replay.py +++ b/mitmproxy/protocol/http_replay.py @@ -55,7 +55,7 @@ class RequestReplayThread(basethread.BaseThread): resp = http1.read_response( server.rfile, connect_request, - body_size_limit=self.config.body_size_limit + body_size_limit=self.config.options.body_size_limit ) if resp.status_code != 200: raise exceptions.ReplayException("Upstream server refuses CONNECT request") @@ -83,7 +83,7 @@ class RequestReplayThread(basethread.BaseThread): self.flow.response = models.HTTPResponse.wrap(http1.read_response( server.rfile, r, - body_size_limit=self.config.body_size_limit + body_size_limit=self.config.options.body_size_limit )) if self.channel: response_reply = self.channel.ask("response", self.flow) diff --git a/mitmproxy/proxy/config.py b/mitmproxy/proxy/config.py index ae15a3472..e0994d344 100644 --- a/mitmproxy/proxy/config.py +++ b/mitmproxy/proxy/config.py @@ -10,7 +10,6 @@ from OpenSSL import SSL, crypto from mitmproxy import platform from mitmproxy import exceptions from netlib import certutils -from netlib import human from netlib import tcp from netlib.http import authentication @@ -61,7 +60,6 @@ class ProxyConfig: self, options, no_upstream_cert=False, - body_size_limit=None, mode="regular", upstream_server=None, upstream_auth=None, @@ -84,7 +82,6 @@ class ProxyConfig: self.ciphers_client = ciphers_client self.ciphers_server = ciphers_server self.no_upstream_cert = no_upstream_cert - self.body_size_limit = body_size_limit self.mode = mode if upstream_server: self.upstream_server = ServerSpec(upstream_server[0], tcp.Address.wrap(upstream_server[1])) @@ -114,10 +111,10 @@ class ProxyConfig: self.certstore = None self.clientcerts = None - self.config(options) - options.changed.connect(self) + self.configure(options) + options.changed.connect(self.configure) - def config(self, options): + def configure(self, options): certstore_path = os.path.expanduser(options.cadir) if not os.path.exists(certstore_path): raise exceptions.OptionsError( @@ -151,12 +148,7 @@ class ProxyConfig: ) - def process_proxy_options(parser, options, args): - body_size_limit = args.body_size_limit - if body_size_limit: - body_size_limit = human.parse_size(body_size_limit) - c = 0 mode, upstream_server, upstream_auth = "regular", None, None if args.transparent_proxy: @@ -228,7 +220,6 @@ def process_proxy_options(parser, options, args): return ProxyConfig( options, no_upstream_cert=args.no_upstream_cert, - body_size_limit=body_size_limit, mode=mode, upstream_server=upstream_server, upstream_auth=upstream_auth, diff --git a/test/mitmproxy/test_protocol_http2.py b/test/mitmproxy/test_protocol_http2.py index 34e6656a5..d910ecaec 100644 --- a/test/mitmproxy/test_protocol_http2.py +++ b/test/mitmproxy/test_protocol_http2.py @@ -591,7 +591,7 @@ class TestBodySizeLimit(_Http2Test): return True def test_body_size_limit(self): - self.config.body_size_limit = 20 + self.config.options.body_size_limit = 20 client, h2_conn = self._setup_connection()