Cap nesting depth when extracting strings from user input

aikido_zen/helpers/extract_strings_from_user_input.py

@@ -7,6 +7,9 @@
 from aikido_zen.helpers.build_path_to_payload import build_path_to_payload
 import aikido_zen.context as ctx
 
+# Deeper input would overflow the stack and let the request through unchecked.
+MAX_TRAVERSAL_DEPTH = 64
+
 
 def extract_strings_from_user_input_cached(obj, source):
     """Use the cache to speed up getting user input"""
@@ -32,40 +35,58 @@ def extract_strings_from_user_input(obj, path_to_payload=None):
     if path_to_payload is None:
         path_to_payload = []
 
+    results, _ = extract_strings_and_nesting(obj, path_to_payload)
+    return results
+
+
+def extract_strings_and_nesting(obj, path_to_payload):
+    """Extracts strings from an object and returns how deep its containers nest"""
     results = {}
 
+    if len(path_to_payload) >= MAX_TRAVERSAL_DEPTH:
+        return results, MAX_TRAVERSAL_DEPTH + 1
+
+    nesting = 0
+
     if is_mapping(obj):
         #  Stringifying the dict and adding it as user input is resource intensive
         #  And in most cases shouldn't be necessary.
         for key, value in obj.items():
             results[key] = build_path_to_payload(path_to_payload)
-            for k, v in extract_strings_from_user_input(
+            child_results, child_nesting = extract_strings_and_nesting(
                 value, path_to_payload + [{"type": "object", "key": key}]
-            ).items():
+            )
+            for k, v in child_results.items():
                 results[k] = v
+            nesting = max(nesting, child_nesting + 1)
 
     if isinstance(obj, (set, list, tuple)):
         #  Add the stringified array as well to the results, there might
         #  be accidental concatenation if the client expects a string but gets the array
         #  E.g. HTTP Parameter pollution
-        results[str(obj)] = build_path_to_payload(path_to_payload)
         for i, value in enumerate(obj):
-            for k, v in extract_strings_from_user_input(
+            child_results, child_nesting = extract_strings_and_nesting(
                 value, path_to_payload + [{"type": "array", "index": i}]
-            ).items():
+            )
+            for k, v in child_results.items():
                 results[k] = v
+            nesting = max(nesting, child_nesting + 1)
+
+        if nesting <= MAX_TRAVERSAL_DEPTH:
+            results[str(obj)] = build_path_to_payload(path_to_payload)
 
     if isinstance(obj, str):
         results[obj] = build_path_to_payload(path_to_payload)
         jwt = try_decode_as_jwt(obj)
         if jwt[0]:
-            for k, v in extract_strings_from_user_input(
+            child_results, _ = extract_strings_and_nesting(
                 jwt[1], path_to_payload + [{"type": "jwt"}]
-            ).items():
+            )
+            for k, v in child_results.items():
                 if k == "iss" or v.endswith("<jwt>.iss"):
                     # Do not add the issuer of the JWT as a string because it can contain a
                     # domain / url and produce false positives
                     continue
                 results[k] = v
 
-    return results
+    return results, nesting

aikido_zen/helpers/extract_strings_from_user_input_test.py

@@ -1,6 +1,8 @@
+import base64
 import pytest
 from unittest.mock import MagicMock, patch
 from aikido_zen.helpers.extract_strings_from_user_input import (
+    MAX_TRAVERSAL_DEPTH,
     extract_strings_from_user_input,
     extract_strings_from_user_input_cached,
 )
@@ -288,4 +290,90 @@ def test_extract_strings_from_user_input_cached_multiple_sources(mock_context):
             assert mock_context.parsed_userinput["source2"] == result2
 
 
+def nested_list(depth, leaf="asd"):
+    current = [leaf]
+    for _ in range(depth):
+        current = [current]
+    return current
+
+
+def nested_dict(depth, leaf="asd"):
+    current = leaf
+    for _ in range(depth):
+        current = {"a": current}
+    return current
+
+
+def fake_jwt(payload_json):
+    payload = base64.urlsafe_b64encode(payload_json.encode()).rstrip(b"=").decode()
+    return f".{payload}."
+
+
+def test_extracts_strings_up_to_max_depth_only():
+    assert extract_strings_from_user_input(
+        nested_dict(MAX_TRAVERSAL_DEPTH - 1, "leaf")
+    ) == from_obj(
+        {
+            "a": ".a" * (MAX_TRAVERSAL_DEPTH - 2),
+            "leaf": ".a" * (MAX_TRAVERSAL_DEPTH - 1),
+        }
+    )
+    assert extract_strings_from_user_input(
+        nested_dict(MAX_TRAVERSAL_DEPTH, "leaf")
+    ) == from_obj({"a": ".a" * (MAX_TRAVERSAL_DEPTH - 1)})
+
+
+def test_stringifies_arrays_unless_nested_deeper_than_max_depth():
+    assert extract_strings_from_user_input({"arr": [["p"], {"k": ["q"]}]}) == from_obj(
+        {
+            "arr": ".",
+            "[['p'], {'k': ['q']}]": ".arr",
+            "['p']": ".arr.[0]",
+            "p": ".arr.[0].[0]",
+            "k": ".arr.[1]",
+            "['q']": ".arr.[1].k",
+            "q": ".arr.[1].k.[0]",
+        }
+    )
+    assert extract_strings_from_user_input(
+        {"arr": nested_list(MAX_TRAVERSAL_DEPTH + 1)}
+    ) == from_obj({"arr": "."})
+
+
+def test_deeply_nested_list_keeps_other_strings():
+    # deep enough that any recursion through the value overflows the stack
+    assert extract_strings_from_user_input(
+        {"deep": nested_list(100_000), "user_input": "/etc/passwd"}
+    ) == from_obj({"deep": ".", "user_input": ".", "/etc/passwd": ".user_input"})
+
+
+def test_deeply_nested_mapping_keeps_other_strings():
+    assert extract_strings_from_user_input(
+        {"deep": nested_dict(5000), "user_input": "/etc/passwd"}
+    ) == from_obj(
+        {
+            "deep": ".",
+            "a": ".deep" + ".a" * (MAX_TRAVERSAL_DEPTH - 2),
+            "user_input": ".",
+            "/etc/passwd": ".user_input",
+        }
+    )
+
+
+def test_deeply_nested_jwt_payload_keeps_other_strings():
+    depth = 100_000
+    jwt = fake_jwt("[" * depth + '"asd"' + "]" * depth)
+
+    assert extract_strings_from_user_input(
+        {"user_input": "/etc/passwd", "not_used": jwt}
+    ) == from_obj(
+        {
+            "user_input": ".",
+            "/etc/passwd": ".user_input",
+            "not_used": ".",
+            jwt: ".not_used",
+        }
+    )
+
+
 # To run the tests, use the command: pytest <filename>.py