Add HTTP body size limit specification to command-line tools.

5 files changed, 53 insertions, 17 deletions

diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py
index 5741970b..e368a30b 100644
--- a/libmproxy/cmdline.py
+++ b/libmproxy/cmdline.py
@@ -55,6 +55,12 @@ def common_options(parser):
         help="Show event log."
     )
     parser.add_option(
+        "-l",
+        action="store", dest="body_size_limit", default=None,
+        help="Byte size limit of HTTP request and response bodies."\
+             " Understands k/m/g suffixes, i.e. 3m for 3 megabytes."
+    )
+    parser.add_option(
         "-n",
         action="store_true", dest="no_server",
         help="Don't start a proxy server."
diff --git a/libmproxy/console.py b/libmproxy/console.py
index c979b3ad..d090314f 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -855,7 +855,6 @@ class Options(object):
         "keepserving",
         "kill",
         "intercept",
-        "limit",
         "no_server",
         "refresh_server_playback",
         "rfile",
@@ -943,11 +942,6 @@ class ConsoleMaster(flow.FlowMaster):
         self.conn_list_view = None
         self.set_palette()
 
-        r = self.set_limit(options.limit)
-        if r:
-            print >> sys.stderr, "Limit error:", r
-            sys.exit(1)
-
         r = self.set_intercept(options.intercept)
         if r:
             print >> sys.stderr, "Intercept error:", r
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index c09352c2..e6ead002 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -1260,7 +1260,7 @@ class FlowMaster(controller.Master):
             f.response = None
             f.error = None
             self.process_new_request(f)
-            rt = proxy.RequestReplayThread(f, self.masterq)
+            rt = proxy.RequestReplayThread(f, self.masterq, None)
             rt.start()
         #end nocover
 
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 6a1d8f6d..d1596b0a 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -22,12 +22,13 @@ class ProxyError(Exception):
 
 
 class ProxyConfig:
-    def __init__(self, certfile = None, ciphers = None, cacert = None, cert_wait_time=0):
+    def __init__(self, certfile = None, ciphers = None, cacert = None, cert_wait_time=0, body_size_limit = None):
         self.certfile = certfile
         self.ciphers = ciphers
         self.cacert = cacert
         self.certdir = None
         self.cert_wait_time = cert_wait_time
+        self.body_size_limit = body_size_limit
 
 
 def read_chunked(fp, limit):
@@ -160,13 +161,13 @@ class FileLike:
 
 #begin nocover
 class RequestReplayThread(threading.Thread):
-    def __init__(self, flow, masterq):
-        self.flow, self.masterq = flow, masterq
+    def __init__(self, flow, masterq, body_size_limit):
+        self.flow, self.masterq, self.body_size_limit = flow, masterq, body_size_limit
         threading.Thread.__init__(self)
 
     def run(self):
         try:
-            server = ServerConnection(self.flow.request)
+            server = ServerConnection(self.flow.request, self.body_size_limit)
             server.send_request(self.flow.request)
             response = server.read_response()
             response._send(self.masterq)
@@ -176,7 +177,8 @@ class RequestReplayThread(threading.Thread):
 
 
 class ServerConnection:
-    def __init__(self, request):
+    def __init__(self, request, body_size_limit):
+        self.body_size_limit = body_size_limit
         self.host = request.host
         self.port = request.port
         self.scheme = request.scheme
@@ -226,7 +228,7 @@ class ServerConnection:
         if self.request.method == "HEAD" or code == 204 or code == 304:
             content = ""
         else:
-            content = read_http_body(self.rfile, self, headers, True, None)
+            content = read_http_body(self.rfile, self, headers, True, self.body_size_limit)
         return flow.Response(self.request, code, msg, headers, content)
 
     def terminate(self):
@@ -274,7 +276,7 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
                 request = False
                 response = response._send(self.mqueue)
             else:
-                server = ServerConnection(request)
+                server = ServerConnection(request, self.config.body_size_limit)
                 server.send_request(request)
                 try:
                     response = server.read_response()
@@ -381,7 +383,7 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
                     client_conn.close = True
                 if value == "keep-alive":
                     client_conn.close = False
-        content = read_http_body(self.rfile, client_conn, headers, False, None)
+        content = read_http_body(self.rfile, client_conn, headers, False, self.config.body_size_limit)
         return flow.Request(client_conn, host, port, scheme, method, path, headers, content)
 
     def send_response(self, response):
@@ -463,7 +465,7 @@ def certificate_option_group(parser):
     parser.add_option_group(group)
 
 
-def process_certificate_option_group(parser, options):
+def process_proxy_options(parser, options):
     if options.cert:
         options.cert = os.path.expanduser(options.cert)
         if not os.path.exists(options.cert):
@@ -475,9 +477,13 @@ def process_certificate_option_group(parser, options):
         utils.dummy_ca(cacert)
     if getattr(options, "cache", None) is not None:
         options.cache = os.path.expanduser(options.cache)
+    body_size_limit = utils.parse_size(options.body_size_limit)
     return ProxyConfig(
         certfile = options.cert,
         cacert = cacert,
         ciphers = options.ciphers,
-        cert_wait_time = options.cert_wait_time
+        cert_wait_time = options.cert_wait_time,
+        body_size_limit = body_size_limit
     )
+
+
diff --git a/libmproxy/utils.py b/libmproxy/utils.py
index 37b751dc..a7e41752 100644
--- a/libmproxy/utils.py
+++ b/libmproxy/utils.py
@@ -410,3 +410,33 @@ def parse_url(url):
     return scheme, host, port, path
 
 
+def parse_size(s):
+    """
+        Parses a size specification. Valid specifications are:
+            
+            123: bytes
+            123k: kilobytes
+            123m: megabytes
+            123g: gigabytes
+    """
+    if not s:
+        return None
+    mult = None
+    if s[-1].lower() == "k":
+        mult = 1024**1
+    elif s[-1].lower() == "m":
+        mult = 1024**2
+    elif s[-1].lower() == "g":
+        mult = 1024**3
+
+    if mult:
+        s = s[:-1]
+    else:
+        mult = 1
+    try:
+        return int(s) * mult
+    except ValueError:
+        raise ValueError("Invalid size specification: %s"%s)
+
+
+