Skip to content

Commit d610a4e

Browse files
committed
Skip context creation for bypassed IPs
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
1 parent c77efdf commit d610a4e

14 files changed

Lines changed: 271 additions & 18 deletions

.github/workflows/qa-tests.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,4 +52,4 @@ jobs:
5252
app_port: 8080
5353
sleep_before_test: 30
5454
config_update_delay: 100
55-
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
55+
skip_tests: test_demo_apps_generic_tests,test_path_traversal,test_wave_attack,test_user_rate_limiting_1_minute
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
"""
2+
Helper that either sets the request context as current OR marks the request
3+
as bypassed (when the remote IP is in the bypass list).
4+
"""
5+
6+
from aikido_zen.helpers.logging import logger
7+
from aikido_zen.storage import bypassed_context_store
8+
from aikido_zen.thread.thread_cache import get_cache
9+
from . import current_context
10+
11+
12+
def apply_context_or_bypass(context):
13+
"""
14+
For bypassed IPs: clears the current context and sets the bypassed flag.
15+
For other IPs: sets the context as current and clears the bypassed flag.
16+
17+
Mirrors the firewall-java BypassedContextStore approach so almost every
18+
blocking site short-circuits naturally on `if not context: return`.
19+
"""
20+
try:
21+
cache = get_cache()
22+
if (
23+
cache
24+
and context
25+
and context.remote_address
26+
and cache.is_bypassed_ip(context.remote_address)
27+
):
28+
current_context.set(None)
29+
bypassed_context_store.set_bypassed(True)
30+
return
31+
32+
bypassed_context_store.set_bypassed(False)
33+
if context:
34+
context.set_as_current_context()
35+
except Exception as e:
36+
logger.debug("Exception in apply_context_or_bypass: %s", e)
37+
# On error, fall back to the previous behaviour (set context).
38+
bypassed_context_store.set_bypassed(False)
39+
if context:
40+
context.set_as_current_context()
Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
import pytest
2+
3+
from aikido_zen.context import current_context, get_current_context
4+
from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
5+
from aikido_zen.storage import bypassed_context_store
6+
from aikido_zen.test_utils.context_utils import generate_context
7+
from aikido_zen.thread.thread_cache import get_cache
8+
9+
10+
@pytest.fixture(autouse=True)
11+
def _reset_state():
12+
yield
13+
current_context.set(None)
14+
bypassed_context_store.clear()
15+
cache = get_cache()
16+
if cache:
17+
cache.reset()
18+
19+
20+
def _set_bypass_list(ips):
21+
cache = get_cache()
22+
cache.config.bypassed_ips = _IPMatcherStub(ips)
23+
24+
25+
class _IPMatcherStub:
26+
def __init__(self, ips):
27+
self._ips = set(ips)
28+
29+
def has(self, ip):
30+
return ip in self._ips
31+
32+
33+
def test_non_bypassed_ip_sets_context_and_clears_flag():
34+
_set_bypass_list({"9.9.9.9"})
35+
bypassed_context_store.set_bypassed(True) # stale value from previous request
36+
37+
ctx = generate_context(ip="1.2.3.4")
38+
apply_context_or_bypass(ctx)
39+
40+
assert get_current_context() is ctx
41+
assert bypassed_context_store.is_bypassed() is False
42+
43+
44+
def test_bypassed_ip_clears_context_and_sets_flag():
45+
_set_bypass_list({"1.2.3.4"})
46+
47+
ctx = generate_context(ip="1.2.3.4")
48+
apply_context_or_bypass(ctx)
49+
50+
assert get_current_context() is None
51+
assert bypassed_context_store.is_bypassed() is True
52+
53+
54+
def test_no_remote_address_falls_through_to_set_context():
55+
_set_bypass_list({"1.2.3.4"})
56+
57+
ctx = generate_context()
58+
ctx.remote_address = None
59+
apply_context_or_bypass(ctx)
60+
61+
assert get_current_context() is ctx
62+
assert bypassed_context_store.is_bypassed() is False
63+
64+
65+
def test_bypass_then_non_bypass_resets_flag():
66+
_set_bypass_list({"1.2.3.4"})
67+
68+
ctx_bypassed = generate_context(ip="1.2.3.4")
69+
apply_context_or_bypass(ctx_bypassed)
70+
assert bypassed_context_store.is_bypassed() is True
71+
72+
ctx_normal = generate_context(ip="9.9.9.9")
73+
apply_context_or_bypass(ctx_normal)
74+
assert get_current_context() is ctx_normal
75+
assert bypassed_context_store.is_bypassed() is False

aikido_zen/sinks/socket/should_block_outbound_domain.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from aikido_zen.storage import bypassed_context_store
12
from aikido_zen.thread.thread_cache import get_cache
23

34

@@ -6,6 +7,11 @@ def should_block_outbound_domain(hostname, port):
67
if not process_cache:
78
return False
89

10+
if bypassed_context_store.is_bypassed():
11+
# Bypassed IPs are trusted — don't report their outbound hostnames
12+
# in heartbeats and don't block their outbound traffic.
13+
return False
14+
915
# We store the hostname before checking the blocking status
1016
# This is because if we are in lockdown mode and blocking all new hostnames, it should still
1117
# show up in the dashboard. This allows the user to allow traffic to newly detected hostnames.
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import pytest
2+
3+
from aikido_zen.sinks.socket.should_block_outbound_domain import (
4+
should_block_outbound_domain,
5+
)
6+
from aikido_zen.storage import bypassed_context_store
7+
from aikido_zen.thread.thread_cache import get_cache
8+
9+
10+
@pytest.fixture(autouse=True)
11+
def _reset_state():
12+
cache = get_cache()
13+
if cache:
14+
cache.reset()
15+
bypassed_context_store.clear()
16+
yield
17+
if cache:
18+
cache.reset()
19+
bypassed_context_store.clear()
20+
21+
22+
def test_unknown_domain_not_blocked():
23+
assert should_block_outbound_domain("safe.example.com", 80) is False
24+
assert any(
25+
h["hostname"] == "safe.example.com"
26+
for h in get_cache().hostnames.as_array()
27+
)
28+
29+
30+
def test_blocked_domain_is_blocked_and_recorded():
31+
cache = get_cache()
32+
cache.config.update_outbound_domains([{"hostname": "evil.example.com", "mode": "block"}])
33+
34+
assert should_block_outbound_domain("evil.example.com", 80) is True
35+
# Blocked domains are still recorded so they show up in the dashboard.
36+
assert any(
37+
h["hostname"] == "evil.example.com"
38+
for h in cache.hostnames.as_array()
39+
)
40+
41+
42+
def test_bypassed_request_does_not_block_or_record():
43+
cache = get_cache()
44+
cache.config.update_outbound_domains([{"hostname": "evil.example.com", "mode": "block"}])
45+
bypassed_context_store.set_bypassed(True)
46+
47+
assert should_block_outbound_domain("evil.example.com", 80) is False
48+
# No hostname pollution from bypassed-IP requests.
49+
assert cache.hostnames.as_array() == []
50+
51+
52+
def test_bypassed_request_does_not_record_unknown_domain():
53+
bypassed_context_store.set_bypassed(True)
54+
55+
assert should_block_outbound_domain("anything.example.com", 80) is False
56+
assert get_cache().hostnames.as_array() == []

aikido_zen/sources/django/run_init_stage.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Exports run_init_stage function"""
22

33
from aikido_zen.context import Context
4+
from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
45
from aikido_zen.helpers.logging import logger
56
from .extract_body import extract_body_from_django_request
67
from .extract_cookies import extract_cookies_from_django_request
@@ -24,7 +25,7 @@ def run_init_stage(request):
2425
else:
2526
return
2627
context.set_cookies(cookies)
27-
context.set_as_current_context()
28+
apply_context_or_bypass(context)
2829

2930
# Init stage needs to be run with context already set :
3031
request_handler(stage="init")

aikido_zen/sources/flask/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from aikido_zen.context import Context
2+
from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
23
from aikido_zen.helpers.get_argument import get_argument
34
from aikido_zen.sinks import on_import, patch_function, before_modify_return, after
45
from .extract_cookies import extract_cookies_from_flask_request_and_save_data
@@ -33,7 +34,7 @@ def _call(func, instance, args, kwargs):
3334
start_response = get_argument(args, kwargs, 1, "start_response")
3435

3536
context1 = Context(req=environ, source="flask")
36-
context1.set_as_current_context()
37+
apply_context_or_bypass(context1)
3738
funcs.request_handler(stage="init")
3839

3940
# Checks for blocked IPs, blocked UAs, ...

aikido_zen/sources/quart.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from aikido_zen.context import Context, get_current_context
2+
from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
23
from .functions.request_handler import request_handler
34
from ..helpers.get_argument import get_argument
45
from ..sinks import on_import, patch_function, before, before_async
@@ -11,7 +12,7 @@ def _call(func, instance, args, kwargs):
1112
return
1213

1314
new_context = Context(req=scope, source="quart")
14-
new_context.set_as_current_context()
15+
apply_context_or_bypass(new_context)
1516
request_handler(stage="init")
1617

1718

aikido_zen/sources/starlette/starlette_applications.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Wraps starlette.applications for initial request_handler"""
22

33
from aikido_zen.context import Context
4+
from aikido_zen.context.apply_or_bypass import apply_context_or_bypass
45
from ..functions.request_handler import request_handler
56
from ...helpers.get_argument import get_argument
67
from ...sinks import on_import, patch_function, before
@@ -13,7 +14,7 @@ def _call(func, instance, args, kwargs):
1314
return
1415

1516
new_context = Context(req=scope, source="starlette")
16-
new_context.set_as_current_context()
17+
apply_context_or_bypass(new_context)
1718
request_handler(stage="init")
1819

1920

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
"""
2+
Records whether the current request's remote IP is in the bypass list.
3+
4+
Bypassed requests intentionally do not set a Context (so all per-request
5+
protection short-circuits), but checks that run without a Context — for
6+
example outbound DNS reporting — still need a way to detect "this work
7+
was triggered by a bypassed request".
8+
"""
9+
10+
import contextvars
11+
12+
_bypassed = contextvars.ContextVar("aikido_bypassed_ip", default=False)
13+
14+
15+
def set_bypassed(value: bool) -> None:
16+
_bypassed.set(bool(value))
17+
18+
19+
def is_bypassed() -> bool:
20+
try:
21+
return _bypassed.get()
22+
except Exception:
23+
return False
24+
25+
26+
def clear() -> None:
27+
_bypassed.set(False)

0 commit comments

Comments
 (0)