Attempted fix for urllib3 crashes

projects/urllib3/fuzz_requests.py

@@ -16,16 +16,11 @@
 
 import atheris
 from http.server import BaseHTTPRequestHandler, HTTPServer
-import os
 import random
 import sys
 import threading
-import time
 import urllib3
 
-timeout = urllib3.util.Timeout(connect=1.0, read=1.0)
-urllib_pool = urllib3.PoolManager(timeout=timeout)
-
 PORT = 8011
 
 GLOBAL_RESPONSE_MESSAGE = ""
@@ -69,67 +64,79 @@ def log_request(self, code="-", size="-"):
         return
 
 
-def run_webserver():
-    with HTTPServer(("", PORT), handler) as server:
-        server.serve_forever()
+class StoppableHTTPServer(HTTPServer):
+    def run(self):
+        try:
+            self.serve_forever()
+        finally:
+            self.server_close()
 
 
 REQUEST_METHODS = ["POST", "GET", "HEAD", "PUT", "DELETE", "OPTIONS", "PATCH"]
 CONTENT_ENCODING_TYPES = [None, "gzip", "deflate"]
 
 
 def TestOneInput(input_bytes):
-    global GLOBAL_RESPONSE_MESSAGE, GLOBAL_RESPONSE_CODE, GLOBAL_CONTENT_ENCODING
-    fdp = atheris.FuzzedDataProvider(input_bytes)
+    global GLOBAL_RESPONSE_MESSAGE, GLOBAL_RESPONSE_CODE, GLOBAL_CONTENT_ENCODING, PORT
 
-    # Fuzz Http Response
-    GLOBAL_RESPONSE_MESSAGE = fdp.ConsumeUnicodeNoSurrogates(sys.maxsize)
-    GLOBAL_RESPONSE_CODE = fdp.ConsumeIntInRange(200, 599)
-    GLOBAL_CONTENT_ENCODING = fdp.PickValueInList(CONTENT_ENCODING_TYPES)
-
-    # Fuzz Http Request
-    requestType = fdp.PickValueInList(REQUEST_METHODS)
-    # Optionally provide request headers
-    requestHeaders = urllib3._collections.HTTPHeaderDict({})
-    for i in range(0, fdp.ConsumeIntInRange(0, 10)):
-        requestHeaders.add(
-            fdp.ConsumeString(sys.maxsize), fdp.ConsumeString(sys.maxsize)
-        )
-    requestHeaders = None if fdp.ConsumeBool() else requestHeaders
-
-    # Optionally generate form data for request
-    formData = {}
-    for i in range(0, fdp.ConsumeIntInRange(0, 100)):
-        formData[fdp.ConsumeString(sys.maxsize)] = fdp.ConsumeString(sys.maxsize)
-    formData = None if fdp.ConsumeBool() else formData
-
-    # Optionally generate request body
-    requestBody = None if fdp.ConsumeBool() else fdp.ConsumeString(sys.maxsize)
-
-    r = urllib_pool.request(
-        requestType,
-        f"http://localhost:{PORT}/",
-        headers=requestHeaders,
-        fields=formData,
-        body=requestBody
-    )
-    r.status
-    r.data
-    r.headers
+    timeout = urllib3.util.Timeout(connect=1.0, read=1.0)
+    urllib_pool = urllib3.PoolManager(timeout=timeout)
 
-def main():
     # Try and get an open port to run our test web server
     for attempt in range(10):
         try:
-            PORT = random.randint(8000,9999)
-            x = threading.Thread(target=run_webserver, daemon=True)
-            x.start()
+            PORT = random.randint(8000, 9999)
+            server = StoppableHTTPServer(("127.0.0.1", PORT), handler)
+            t1 = threading.Thread(None, server.run)
+            t1.start()
             break
-        except OSError as e:
-            pass 
+        except OSError:
+            pass
 
-    time.sleep(0.5)  # Short delay to start test server
+    fdp = atheris.FuzzedDataProvider(input_bytes)
 
+    BATCH_SIZE = 100
+    for iteration in range(BATCH_SIZE):
+        # Fuzz Http Response
+        GLOBAL_RESPONSE_MESSAGE = fdp.ConsumeUnicodeNoSurrogates(sys.maxsize)
+        GLOBAL_RESPONSE_CODE = fdp.ConsumeIntInRange(200, 599)
+        GLOBAL_CONTENT_ENCODING = fdp.PickValueInList(CONTENT_ENCODING_TYPES)
+
+        # Fuzz Http Request
+        requestType = fdp.PickValueInList(REQUEST_METHODS)
+        # Optionally provide request headers
+        requestHeaders = urllib3._collections.HTTPHeaderDict({})
+        for i in range(0, fdp.ConsumeIntInRange(0, 10)):
+            requestHeaders.add(
+                fdp.ConsumeString(sys.maxsize), fdp.ConsumeString(sys.maxsize)
+            )
+        requestHeaders = None if fdp.ConsumeBool() else requestHeaders
+
+        # Optionally generate form data for request
+        formData = {}
+        for i in range(0, fdp.ConsumeIntInRange(0, 100)):
+            formData[fdp.ConsumeString(sys.maxsize)] = fdp.ConsumeString(sys.maxsize)
+        formData = None if fdp.ConsumeBool() else formData
+
+        # Optionally generate request body
+        requestBody = None if fdp.ConsumeBool() else fdp.ConsumeString(sys.maxsize)
+
+        r = urllib_pool.request(
+            requestType,
+            f"http://localhost:{PORT}/",
+            headers=requestHeaders,
+            fields=formData,
+            body=requestBody,
+        )
+        r.status
+        r.data
+        r.headers
+
+    server.shutdown()
+    t1.join()
+
+
+def main():
     atheris.instrument_all()
     atheris.Setup(sys.argv, TestOneInput)
     atheris.Fuzz()