From d610a4eae2b4a132be94e54aa50669a22ad97936 Mon Sep 17 00:00:00 2001
From: bitterpanda <bitterpanda@proton.me>
Date: Thu, 16 Apr 2026 16:14:00 +0200
Subject: [PATCH] Skip context creation for bypassed IPs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Bypassed IPs now produce no request context and a thread-local
BypassedContextStore flag instead. This mirrors the firewall-java approach
(PR #284): every per-request blocking site already short-circuits on
"if not context: return", so geo blocking, IP blocklists, bot blocking,
route allowlists, blocked user IDs, rate limiting, attack detection, and
heartbeat stats all skip naturally — without scattering bypass checks
across each call site.

Context-less checks (outbound DNS reporting, stored SSRF) consult the
BypassedContextStore flag directly so they too become no-ops for traffic
from a bypassed IP.

Re-enables in the QA suite:
- test_bypassed_ip
- test_bypassed_ip_for_geo_blocking
- test_block_traffic_by_countries
- test_outbound_domain_blocking
---
 .github/workflows/qa-tests.yml                |  2 +-
 aikido_zen/context/apply_or_bypass.py         | 40 ++++++++++
 aikido_zen/context/apply_or_bypass_test.py    | 75 +++++++++++++++++++
 .../socket/should_block_outbound_domain.py    |  6 ++
 .../should_block_outbound_domain_test.py      | 56 ++++++++++++++
 aikido_zen/sources/django/run_init_stage.py   |  3 +-
 aikido_zen/sources/flask/__init__.py          |  3 +-
 aikido_zen/sources/quart.py                   |  3 +-
 .../starlette/starlette_applications.py       |  3 +-
 aikido_zen/storage/bypassed_context_store.py  | 27 +++++++
 .../storage/bypassed_context_store_test.py    | 37 +++++++++
 aikido_zen/vulnerabilities/__init__.py        | 10 ++-
 aikido_zen/vulnerabilities/init_test.py       | 14 +++-
 .../ssrf/inspect_getaddrinfo_result.py        | 10 +--
 14 files changed, 271 insertions(+), 18 deletions(-)
 create mode 100644 aikido_zen/context/apply_or_bypass.py
 create mode 100644 aikido_zen/context/apply_or_bypass_test.py
 create mode 100644 aikido_zen/sinks/socket/should_block_outbound_domain_test.py
 create mode 100644 aikido_zen/storage/bypassed_context_store.py
 create mode 100644 aikido_zen/storage/bypassed_context_store_test.py

diff --git a/.github/workflows/qa-tests.yml b/.github/workflows/qa-tests.yml
index 075a75b54..15dc3b803 100644
--- a/.github/workflows/qa-tests.yml
+++ b/.github/workflows/qa-tests.yml
@@ -52,4 +52,4 @@ jobs:
           app_port: 8080
           sleep_before_test: 30
           config_update_delay: 100
-          skip_tests: test_bypassed_ip_for_geo_blocking,test_demo_apps_generic_tests,test_path_traversal,test_outbound_domain_blocking,test_bypassed_ip,test_wave_attack,test_block_traffic_by_countries,test_user_rate_limiting_1_minute
+          skip_tests: test_demo_apps_generic_tests,test_path_traversal,test_wave_attack,test_user_rate_limiting_1_minute
diff --git a/aikido_zen/context/apply_or_bypass.py b/aikido_zen/context/apply_or_bypass.py
new file mode 100644
index 000000000..3d6ace1df
--- /dev/null
+++ b/aikido_zen/context/apply_or_bypass.py
@@ -0,0 +1,40 @@
+"""
+Helper that either sets the request context as current OR marks the request
+as bypassed (when the remote IP is in the bypass list).
+"""
+
+from aikido_zen.helpers.logging import logger
+from aikido_zen.storage import bypassed_context_store
+from aikido_zen.thread.thread_cache import get_cache
+from . import current_context
+
+
+def apply_context_or_bypass(context):
+    """
+    For bypassed IPs: clears the current context and sets the bypassed flag.
+    For other IPs:    sets the context as current and clears the bypassed flag.
+
+    Mirrors the firewall-java BypassedContextStore approach so almost every
+    blocking site short-circuits naturally on `if not context: return`.
+    """
+    try:
+        cache = get_cache()
+        if (
+            cache
+            and context
+            and context.remote_address
+            and cache.is_bypassed_ip(context.remote_address)
+        ):
+            current_context.set(None)
+            bypassed_context_store.set_bypassed(True)
+            return
+
+        bypassed_context_store.set_bypassed(False)
+        if context:
+            context.set_as_current_context()
+    except Exception as e:
+        logger.debug("Exception in apply_context_or_bypass: %s", e)
+        # On error, fall back to the previous behaviour (set context).
+        bypassed_context_store.set_bypassed(False)
+        if context:
+            context.set_as_current_context()
diff --git a/aikido_zen/context/apply_or_bypass_test.py b/aikido_zen/context/apply_or_bypass_test.py
new file mode 100644
index 000000000..29c04ad85
--- /dev/null
+++ b/aikido_zen/context/apply_or_bypass_test.py
@@ -0,0 +1,75 @@
+import pytest
+
+from aikido_zen.context import current_context, get_current_context
+from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
+from aikido_zen.storage import bypassed_context_store
+from aikido_zen.test_utils.context_utils import generate_context
+from aikido_zen.thread.thread_cache import get_cache
+
+
+@pytest.fixture(autouse=True)
+def _reset_state():
+    yield
+    current_context.set(None)
+    bypassed_context_store.clear()
+    cache = get_cache()
+    if cache:
+        cache.reset()
+
+
+def _set_bypass_list(ips):
+    cache = get_cache()
+    cache.config.bypassed_ips = _IPMatcherStub(ips)
+
+
+class _IPMatcherStub:
+    def __init__(self, ips):
+        self._ips = set(ips)
+
+    def has(self, ip):
+        return ip in self._ips
+
+
+def test_non_bypassed_ip_sets_context_and_clears_flag():
+    _set_bypass_list({"9.9.9.9"})
+    bypassed_context_store.set_bypassed(True)  # stale value from previous request
+
+    ctx = generate_context(ip="1.2.3.4")
+    apply_context_or_bypass(ctx)
+
+    assert get_current_context() is ctx
+    assert bypassed_context_store.is_bypassed() is False
+
+
+def test_bypassed_ip_clears_context_and_sets_flag():
+    _set_bypass_list({"1.2.3.4"})
+
+    ctx = generate_context(ip="1.2.3.4")
+    apply_context_or_bypass(ctx)
+
+    assert get_current_context() is None
+    assert bypassed_context_store.is_bypassed() is True
+
+
+def test_no_remote_address_falls_through_to_set_context():
+    _set_bypass_list({"1.2.3.4"})
+
+    ctx = generate_context()
+    ctx.remote_address = None
+    apply_context_or_bypass(ctx)
+
+    assert get_current_context() is ctx
+    assert bypassed_context_store.is_bypassed() is False
+
+
+def test_bypass_then_non_bypass_resets_flag():
+    _set_bypass_list({"1.2.3.4"})
+
+    ctx_bypassed = generate_context(ip="1.2.3.4")
+    apply_context_or_bypass(ctx_bypassed)
+    assert bypassed_context_store.is_bypassed() is True
+
+    ctx_normal = generate_context(ip="9.9.9.9")
+    apply_context_or_bypass(ctx_normal)
+    assert get_current_context() is ctx_normal
+    assert bypassed_context_store.is_bypassed() is False
diff --git a/aikido_zen/sinks/socket/should_block_outbound_domain.py b/aikido_zen/sinks/socket/should_block_outbound_domain.py
index aefdd14b4..e2a4bd4ce 100644
--- a/aikido_zen/sinks/socket/should_block_outbound_domain.py
+++ b/aikido_zen/sinks/socket/should_block_outbound_domain.py
@@ -1,3 +1,4 @@
+from aikido_zen.storage import bypassed_context_store
 from aikido_zen.thread.thread_cache import get_cache
 
 
@@ -6,6 +7,11 @@ def should_block_outbound_domain(hostname, port):
     if not process_cache:
         return False
 
+    if bypassed_context_store.is_bypassed():
+        # Bypassed IPs are trusted — don't report their outbound hostnames
+        # in heartbeats and don't block their outbound traffic.
+        return False
+
     # We store the hostname before checking the blocking status
     # This is because if we are in lockdown mode and blocking all new hostnames, it should still
     # show up in the dashboard. This allows the user to allow traffic to newly detected hostnames.
diff --git a/aikido_zen/sinks/socket/should_block_outbound_domain_test.py b/aikido_zen/sinks/socket/should_block_outbound_domain_test.py
new file mode 100644
index 000000000..674c61fb9
--- /dev/null
+++ b/aikido_zen/sinks/socket/should_block_outbound_domain_test.py
@@ -0,0 +1,56 @@
+import pytest
+
+from aikido_zen.sinks.socket.should_block_outbound_domain import (
+    should_block_outbound_domain,
+)
+from aikido_zen.storage import bypassed_context_store
+from aikido_zen.thread.thread_cache import get_cache
+
+
+@pytest.fixture(autouse=True)
+def _reset_state():
+    cache = get_cache()
+    if cache:
+        cache.reset()
+    bypassed_context_store.clear()
+    yield
+    if cache:
+        cache.reset()
+    bypassed_context_store.clear()
+
+
+def test_unknown_domain_not_blocked():
+    assert should_block_outbound_domain("safe.example.com", 80) is False
+    assert any(
+        h["hostname"] == "safe.example.com"
+        for h in get_cache().hostnames.as_array()
+    )
+
+
+def test_blocked_domain_is_blocked_and_recorded():
+    cache = get_cache()
+    cache.config.update_outbound_domains([{"hostname": "evil.example.com", "mode": "block"}])
+
+    assert should_block_outbound_domain("evil.example.com", 80) is True
+    # Blocked domains are still recorded so they show up in the dashboard.
+    assert any(
+        h["hostname"] == "evil.example.com"
+        for h in cache.hostnames.as_array()
+    )
+
+
+def test_bypassed_request_does_not_block_or_record():
+    cache = get_cache()
+    cache.config.update_outbound_domains([{"hostname": "evil.example.com", "mode": "block"}])
+    bypassed_context_store.set_bypassed(True)
+
+    assert should_block_outbound_domain("evil.example.com", 80) is False
+    # No hostname pollution from bypassed-IP requests.
+    assert cache.hostnames.as_array() == []
+
+
+def test_bypassed_request_does_not_record_unknown_domain():
+    bypassed_context_store.set_bypassed(True)
+
+    assert should_block_outbound_domain("anything.example.com", 80) is False
+    assert get_cache().hostnames.as_array() == []
diff --git a/aikido_zen/sources/django/run_init_stage.py b/aikido_zen/sources/django/run_init_stage.py
index 372ee129e..09ba5ff42 100644
--- a/aikido_zen/sources/django/run_init_stage.py
+++ b/aikido_zen/sources/django/run_init_stage.py
@@ -1,6 +1,7 @@
 """Exports run_init_stage function"""
 
 from aikido_zen.context import Context
+from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
 from aikido_zen.helpers.logging import logger
 from .extract_body import extract_body_from_django_request
 from .extract_cookies import extract_cookies_from_django_request
@@ -24,7 +25,7 @@ def run_init_stage(request):
         else:
             return
         context.set_cookies(cookies)
-        context.set_as_current_context()
+        apply_context_or_bypass(context)
 
         # Init stage needs to be run with context already set :
         request_handler(stage="init")
diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py
index bc73df788..7eec0a385 100644
--- a/aikido_zen/sources/flask/__init__.py
+++ b/aikido_zen/sources/flask/__init__.py
@@ -1,4 +1,5 @@
 from aikido_zen.context import Context
+from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
 from aikido_zen.helpers.get_argument import get_argument
 from aikido_zen.sinks import on_import, patch_function, before_modify_return, after
 from .extract_cookies import extract_cookies_from_flask_request_and_save_data
@@ -33,7 +34,7 @@ def _call(func, instance, args, kwargs):
     start_response = get_argument(args, kwargs, 1, "start_response")
 
     context1 = Context(req=environ, source="flask")
-    context1.set_as_current_context()
+    apply_context_or_bypass(context1)
     funcs.request_handler(stage="init")
 
     # Checks for blocked IPs, blocked UAs, ...
diff --git a/aikido_zen/sources/quart.py b/aikido_zen/sources/quart.py
index 9e0fa6f63..eea83a88b 100644
--- a/aikido_zen/sources/quart.py
+++ b/aikido_zen/sources/quart.py
@@ -1,4 +1,5 @@
 from aikido_zen.context import Context, get_current_context
+from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
 from .functions.request_handler import request_handler
 from ..helpers.get_argument import get_argument
 from ..sinks import on_import, patch_function, before, before_async
@@ -11,7 +12,7 @@ def _call(func, instance, args, kwargs):
         return
 
     new_context = Context(req=scope, source="quart")
-    new_context.set_as_current_context()
+    apply_context_or_bypass(new_context)
     request_handler(stage="init")
 
 
diff --git a/aikido_zen/sources/starlette/starlette_applications.py b/aikido_zen/sources/starlette/starlette_applications.py
index 980d85f20..b45e42a76 100644
--- a/aikido_zen/sources/starlette/starlette_applications.py
+++ b/aikido_zen/sources/starlette/starlette_applications.py
@@ -1,6 +1,7 @@
 """Wraps starlette.applications for initial request_handler"""
 
 from aikido_zen.context import Context
+from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
 from ..functions.request_handler import request_handler
 from ...helpers.get_argument import get_argument
 from ...sinks import on_import, patch_function, before
@@ -13,7 +14,7 @@ def _call(func, instance, args, kwargs):
         return
 
     new_context = Context(req=scope, source="starlette")
-    new_context.set_as_current_context()
+    apply_context_or_bypass(new_context)
     request_handler(stage="init")
 
 
diff --git a/aikido_zen/storage/bypassed_context_store.py b/aikido_zen/storage/bypassed_context_store.py
new file mode 100644
index 000000000..40ddacbf3
--- /dev/null
+++ b/aikido_zen/storage/bypassed_context_store.py
@@ -0,0 +1,27 @@
+"""
+Records whether the current request's remote IP is in the bypass list.
+
+Bypassed requests intentionally do not set a Context (so all per-request
+protection short-circuits), but checks that run without a Context — for
+example outbound DNS reporting — still need a way to detect "this work
+was triggered by a bypassed request".
+"""
+
+import contextvars
+
+_bypassed = contextvars.ContextVar("aikido_bypassed_ip", default=False)
+
+
+def set_bypassed(value: bool) -> None:
+    _bypassed.set(bool(value))
+
+
+def is_bypassed() -> bool:
+    try:
+        return _bypassed.get()
+    except Exception:
+        return False
+
+
+def clear() -> None:
+    _bypassed.set(False)
diff --git a/aikido_zen/storage/bypassed_context_store_test.py b/aikido_zen/storage/bypassed_context_store_test.py
new file mode 100644
index 000000000..b041dd4fb
--- /dev/null
+++ b/aikido_zen/storage/bypassed_context_store_test.py
@@ -0,0 +1,37 @@
+import pytest
+
+from aikido_zen.storage import bypassed_context_store
+
+
+@pytest.fixture(autouse=True)
+def _reset_after_each_test():
+    yield
+    bypassed_context_store.clear()
+
+
+def test_default_is_false():
+    assert bypassed_context_store.is_bypassed() is False
+
+
+def test_set_bypassed_true_then_false():
+    bypassed_context_store.set_bypassed(True)
+    assert bypassed_context_store.is_bypassed() is True
+
+    bypassed_context_store.set_bypassed(False)
+    assert bypassed_context_store.is_bypassed() is False
+
+
+def test_clear_resets_to_false():
+    bypassed_context_store.set_bypassed(True)
+    assert bypassed_context_store.is_bypassed() is True
+
+    bypassed_context_store.clear()
+    assert bypassed_context_store.is_bypassed() is False
+
+
+def test_truthy_values_coerced_to_bool():
+    bypassed_context_store.set_bypassed("yes")
+    assert bypassed_context_store.is_bypassed() is True
+
+    bypassed_context_store.set_bypassed(0)
+    assert bypassed_context_store.is_bypassed() is False
diff --git a/aikido_zen/vulnerabilities/__init__.py b/aikido_zen/vulnerabilities/__init__.py
index 436e63094..5d175a8f9 100644
--- a/aikido_zen/vulnerabilities/__init__.py
+++ b/aikido_zen/vulnerabilities/__init__.py
@@ -19,6 +19,7 @@
 from aikido_zen.helpers.is_protection_forced_off_cached import (
     is_protection_forced_off_cached,
 )
+from aikido_zen.storage import bypassed_context_store
 from aikido_zen.thread.thread_cache import get_cache
 from .sql_injection.context_contains_sql_injection import context_contains_sql_injection
 from .nosql_injection.check_context import check_context_for_nosql_injection
@@ -44,6 +45,11 @@ def run_vulnerability_scan(kind, op, args):
     if is_protection_forced_off_cached(context):
         return
 
+    if bypassed_context_store.is_bypassed():
+        # Bypassed IPs are trusted across all vulnerability kinds, including
+        # the context-less SSRF path (e.g. stored SSRF).
+        return
+
     comms = comm.get_comms()
     thread_cache = get_cache()
     if not context and kind != "ssrf":
@@ -55,10 +61,6 @@ def run_vulnerability_scan(kind, op, args):
         # Make a special exception for SSRF, which checks itself if thread cache is set.
         # This is because some scans/tests for SSRF do not require a thread cache to be set.
         return
-    if thread_cache and context:
-        if thread_cache.is_bypassed_ip(context.remote_address):
-            #  This IP is on the bypass list, not scanning
-            return
 
     error_type = AikidoException  # Default error
     error_args = tuple()
diff --git a/aikido_zen/vulnerabilities/init_test.py b/aikido_zen/vulnerabilities/init_test.py
index 54e10e89f..9c50debca 100644
--- a/aikido_zen/vulnerabilities/init_test.py
+++ b/aikido_zen/vulnerabilities/init_test.py
@@ -3,6 +3,7 @@
 from . import run_vulnerability_scan
 from aikido_zen.context import current_context, Context
 from aikido_zen.errors import AikidoSQLInjection
+from aikido_zen.storage import bypassed_context_store
 from aikido_zen.thread.thread_cache import get_cache
 from aikido_zen.helpers.ip_matcher import IPMatcher
 
@@ -10,10 +11,12 @@
 @pytest.fixture(autouse=True)
 def run_around_tests():
     get_cache().reset()
+    bypassed_context_store.clear()
     yield
     # Make sure to reset context and cache after every test so it does not
     # interfere with other tests
     current_context.set(None)
+    bypassed_context_store.clear()
     get_cache().reset()
 
 
@@ -81,10 +84,11 @@ def test_ssrf(caplog, get_context):
 
 
 def test_lifecycle_cache_bypassed_ip(caplog, get_context):
+    # Bypassed requests are signalled via BypassedContextStore (set by the
+    # framework entry point's apply_context_or_bypass helper). When the flag
+    # is set, run_vulnerability_scan must short-circuit before any scan logic.
     get_context.set_as_current_context()
-    cache = get_cache()
-    cache.config.bypassed_ips = IPMatcher(["198.51.100.23"])
-    assert cache.is_bypassed_ip("198.51.100.23")
+    bypassed_context_store.set_bypassed(True)
     run_vulnerability_scan(kind="test", op="test", args=tuple())
     assert len(caplog.text) == 0
 
@@ -184,8 +188,10 @@ def test_ssrf_vulnerability_scan_no_port(get_context):
 
 
 def test_ssrf_vulnerability_scan_bypassed_ip(get_context):
+    # Bypassed IP: BypassedContextStore flag is set by the framework entry point.
+    # run_vulnerability_scan must short-circuit before any SSRF inspection.
     get_context.set_as_current_context()
-    get_cache().config.bypassed_ips = IPMatcher(["198.51.100.23"])
+    bypassed_context_store.set_bypassed(True)
 
     dns_results = MagicMock()
     hostname = "example.com"
diff --git a/aikido_zen/vulnerabilities/ssrf/inspect_getaddrinfo_result.py b/aikido_zen/vulnerabilities/ssrf/inspect_getaddrinfo_result.py
index 43fa993f1..3faaa0637 100644
--- a/aikido_zen/vulnerabilities/ssrf/inspect_getaddrinfo_result.py
+++ b/aikido_zen/vulnerabilities/ssrf/inspect_getaddrinfo_result.py
@@ -4,7 +4,7 @@
 
 from aikido_zen.context import get_current_context
 from aikido_zen.helpers.logging import logger
-from aikido_zen.thread.thread_cache import get_cache
+from aikido_zen.storage import bypassed_context_store
 from .imds import resolves_to_imds_ip
 from aikido_zen.helpers.net.is_private_ip import is_private_ip
 from .find_hostname_in_context import find_hostname_in_context
@@ -18,6 +18,10 @@ def inspect_getaddrinfo_result(dns_results, hostname, port):
     if not hostname or not dns_results:
         return  # Ensure that the data we get isnt empty
 
+    if bypassed_context_store.is_bypassed():
+        # Bypassed IPs are trusted: skip both stored-SSRF and request-SSRF checks.
+        return
+
     ip_addresses = extract_ip_array_from_results(dns_results)
     imds_ip = resolves_to_imds_ip(ip_addresses, hostname)
     if imds_ip:
@@ -37,10 +41,6 @@ def inspect_getaddrinfo_result(dns_results, hostname, port):
     context = get_current_context()
     if not context:
         return  # Context should be set to check user input.
-    if get_cache() and get_cache().is_bypassed_ip(context.remote_address):
-        # We check for bypassed ip's here since it is not checked for us
-        # in run_vulnerability_scan due to the exception for SSRF (see above code)
-        return
 
     # attack_findings is an object containing source, pathToPayload and payload.
     attack_findings = find_hostname_in_context(hostname, context, port)