Skip to content

Commit 4ab7b35

Browse files
committed
feat: Enhance CachedMongoTenantResolver with tombstone logic and update unit tests for verification and caching behavior
1 parent 1d6dd55 commit 4ab7b35

6 files changed

Lines changed: 313 additions & 6 deletions

File tree

services/tenant_resolver/cached_mongo.py

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@
3131
_CACHE_MISS = object()
3232
_NEG_CACHE_HIT = object()
3333

34+
# Sentinel value of a tombstone key. Any non-empty value works — readers
35+
# only check existence; the value is informational.
36+
_TOMB_VALUE = "1"
37+
3438

3539
class CachedMongoTenantResolver(TenantResolver):
3640
def __init__(
@@ -42,16 +46,25 @@ def __init__(
4246
# Negative TTL can be aggressive because the orchestrator calls
4347
# ``invalidate(fqdn)`` on every state transition
4448
negative_ttl_seconds: int = 300,
49+
# Tombstone window: how long after invalidate() new cache writes
50+
# for the same host are skipped. Defends against the read-through
51+
# race where a slow resolve() reads stale Mongo data, then races
52+
# invalidate() and writes the stale answer back.
53+
tombstone_ttl_seconds: int = 5,
4554
) -> None:
4655
self._repo = repo
4756
self._redis = redis_client
4857
self._system_default_domain = system_default_domain.lower().rstrip(".")
4958
self._positive_ttl = positive_ttl_seconds
5059
self._negative_ttl = negative_ttl_seconds
60+
self._tombstone_ttl = tombstone_ttl_seconds
5161

5262
def _key(self, host: str) -> str:
5363
return f"tenant:{host}"
5464

65+
def _tomb_key(self, host: str) -> str:
66+
return f"tenant_tomb:{host}"
67+
5568
@staticmethod
5669
def _normalise_host(host: str) -> str:
5770
"""Lowercased, dot-stripped, port-stripped host.
@@ -121,10 +134,32 @@ async def invalidate(self, host: str) -> None:
121134
# System default short-circuits resolve(); no cache slot to drop.
122135
return
123136
try:
124-
await self._redis.delete(self._key(normalised))
137+
# Atomic delete + tombstone via pipeline. Tombstone blocks any
138+
# in-flight resolve() from writing stale data back to the cache
139+
# for the next ``tombstone_ttl_seconds`` — see _cache_set /
140+
# _cache_set_negative for the read-side check.
141+
pipe = self._redis.pipeline()
142+
pipe.delete(self._key(normalised))
143+
pipe.setex(self._tomb_key(normalised), self._tombstone_ttl, _TOMB_VALUE)
144+
await pipe.execute()
125145
except Exception as exc:
126146
log.warning("tenant_cache_invalidate_error", host=host, error=str(exc))
127147

148+
async def _is_tombstoned(self, host: str) -> bool:
149+
"""Return True if invalidate() recently fired for *host*.
150+
151+
Used as the anti-stale guard before any cache write. Failures here
152+
degrade safely to ``False`` — the cache write proceeds (no worse
153+
than current behaviour).
154+
"""
155+
if self._redis is None:
156+
return False
157+
try:
158+
return await self._redis.get(self._tomb_key(host)) is not None
159+
except Exception as exc:
160+
log.warning("tenant_tomb_check_error", host=host, error=str(exc))
161+
return False
162+
128163
# ── Cache helpers ────────────────────────────────────────────────
129164

130165
async def _cache_get(self, host: str):
@@ -171,6 +206,11 @@ async def _cache_get(self, host: str):
171206
async def _cache_set(self, host: str, info: TenantInfo) -> None:
172207
if self._redis is None:
173208
return
209+
# Anti-stale guard: invalidate() may have fired while we were
210+
# querying Mongo — if so, skip the write so we don't poison the
211+
# cache with the now-stale answer.
212+
if await self._is_tombstoned(host):
213+
return
174214
try:
175215
payload = json.dumps(
176216
{
@@ -188,6 +228,8 @@ async def _cache_set(self, host: str, info: TenantInfo) -> None:
188228
async def _cache_set_negative(self, host: str) -> None:
189229
if self._redis is None:
190230
return
231+
if await self._is_tombstoned(host):
232+
return
191233
try:
192234
await self._redis.setex(self._key(host), self._negative_ttl, _NEG_SENTINEL)
193235
except Exception as exc:

tests/unit/repositories/test_custom_domain_repository.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,32 @@ async def test_set_eviction_pending_clears_error_on_success(self):
160160
assert ops["$set"]["eviction_pending"] is False
161161
assert ops["$set"]["last_eviction_error"] is None
162162

163+
@pytest.mark.asyncio
164+
async def test_list_by_owner_paginates_and_orders_by_created_at_desc(self):
165+
# Underpins the dashboard listing endpoint — must filter by owner,
166+
# sort newest first, and respect skip/limit so pagination works.
167+
col = AsyncMock()
168+
cursor = MagicMock()
169+
cursor.sort = MagicMock(return_value=cursor)
170+
cursor.skip = MagicMock(return_value=cursor)
171+
cursor.limit = MagicMock(return_value=cursor)
172+
cursor.to_list = AsyncMock(return_value=[])
173+
col.find = MagicMock(return_value=cursor)
174+
col.name = "custom_domains"
175+
repo = CustomDomainRepository(col)
176+
owner_id = ObjectId()
177+
178+
await repo.list_by_owner(owner_id, skip=20, limit=10)
179+
180+
# (1) filters by owner_id
181+
assert col.find.call_args.args[0] == {"owner_id": owner_id}
182+
# (2) sorts by created_at descending
183+
cursor.sort.assert_called_once_with("created_at", -1)
184+
# (3) respects skip + limit
185+
cursor.skip.assert_called_once_with(20)
186+
cursor.limit.assert_called_once_with(10)
187+
cursor.to_list.assert_awaited_once_with(length=10)
188+
163189
@pytest.mark.asyncio
164190
async def test_find_stale_active_excludes_system_default(self):
165191
col = AsyncMock()

tests/unit/services/test_custom_domain_service.py

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -426,3 +426,109 @@ async def test_status_transition_happens_regardless_of_caddy_outcome(self):
426426
# update_status was called with REVOKED first
427427
first_call = repo.update_status.call_args_list[0]
428428
assert first_call.args[1] == DomainStatus.REVOKED
429+
430+
431+
class TestSuspendNotFound:
432+
@pytest.mark.asyncio
433+
async def test_suspend_missing_domain_is_noop(self):
434+
# Worker may race with a concurrent delete — domain disappears
435+
# between scan and suspend. Must silently noop, not crash.
436+
svc, repo, _, edge, _ = _build_service()
437+
repo.find_by_id = AsyncMock(return_value=None)
438+
439+
# Must not raise
440+
await svc.suspend(DOMAIN_OID, reason="missing_domain_ok")
441+
442+
repo.update_status.assert_not_called()
443+
edge.announce_revoked.assert_not_called()
444+
445+
446+
class TestVerifyAttemptsQuota:
447+
@pytest.mark.asyncio
448+
async def test_quota_increments_redis_counter(self):
449+
redis = AsyncMock()
450+
redis.incr = AsyncMock(return_value=1)
451+
redis.expire = AsyncMock()
452+
svc, repo, _, _, _ = _build_service(redis=redis)
453+
starting = _doc(status=DomainStatus.PENDING)
454+
repo.find_by_id = AsyncMock(side_effect=[starting, starting])
455+
456+
await svc.verify(DOMAIN_OID, _user())
457+
458+
redis.incr.assert_awaited_once()
459+
# First incr → expire is set so the counter rolls over after the window
460+
redis.expire.assert_awaited_once()
461+
462+
@pytest.mark.asyncio
463+
async def test_quota_exceeded_raises_quota_error(self):
464+
redis = AsyncMock()
465+
redis.incr = AsyncMock(return_value=99) # well over default cap of 5
466+
redis.expire = AsyncMock()
467+
svc, repo, _, _, _ = _build_service(redis=redis)
468+
repo.find_by_id = AsyncMock(return_value=_doc(status=DomainStatus.PENDING))
469+
470+
with pytest.raises(DomainQuotaExceededError):
471+
await svc.verify(DOMAIN_OID, _user())
472+
473+
@pytest.mark.asyncio
474+
async def test_quota_fails_open_when_redis_errors(self):
475+
# If Redis is down we degrade to "no quota enforcement" rather than
476+
# blocking all verifies — staff allowlist + per-user-per-domain
477+
# natural rate-limit cover the abuse vector.
478+
redis = AsyncMock()
479+
redis.incr = AsyncMock(side_effect=Exception("redis down"))
480+
svc, repo, _, _, _ = _build_service(redis=redis)
481+
starting = _doc(status=DomainStatus.PENDING)
482+
repo.find_by_id = AsyncMock(side_effect=[starting, starting])
483+
484+
# Must not raise — verify proceeds normally despite Redis fault.
485+
await svc.verify(DOMAIN_OID, _user())
486+
487+
488+
class TestReverifyActive:
489+
@pytest.mark.asyncio
490+
async def test_success_bumps_last_verified_and_clears_error(self):
491+
svc, repo, _, _, _ = _build_service()
492+
d = _doc(status=DomainStatus.ACTIVE)
493+
repo.find_stale_active = AsyncMock(return_value=[d])
494+
495+
result_pairs = await svc.reverify_active(batch_size=10)
496+
497+
assert len(result_pairs) == 1
498+
ok_call = repo.update_status.call_args
499+
# ACTIVE remains, last_verified_at bumped, error cleared
500+
assert ok_call.args[1] == DomainStatus.ACTIVE
501+
assert ok_call.kwargs["bump_last_verified_at"] is True
502+
assert ok_call.kwargs["last_verification_error"] is None
503+
504+
@pytest.mark.asyncio
505+
async def test_failure_records_reason_keeps_status(self):
506+
svc, repo, verifiers, _, _ = _build_service()
507+
verifiers[VerificationMethod.CNAME].verify = AsyncMock(
508+
return_value=VerificationResult(False, reason="DNS NXDOMAIN")
509+
)
510+
d = _doc(status=DomainStatus.ACTIVE)
511+
repo.find_stale_active = AsyncMock(return_value=[d])
512+
513+
await svc.reverify_active(batch_size=10)
514+
515+
# Status unchanged (worker doesn't auto-suspend on a single fail —
516+
# that's the consecutive-failure counter's job, lives in the worker)
517+
call = repo.update_status.call_args
518+
assert call.args[1] == DomainStatus.ACTIVE
519+
# bump_last_verified_at omitted on failure → default False applies.
520+
assert call.kwargs.get("bump_last_verified_at", False) is False
521+
assert call.kwargs["last_verification_error"] == "DNS NXDOMAIN"
522+
523+
@pytest.mark.asyncio
524+
async def test_skips_doc_when_verifier_missing(self):
525+
# Defensive: if a doc references a verification_method that isn't
526+
# wired (legacy data), the loop must skip it without crashing.
527+
svc, repo, _, _, _ = _build_service(verifiers={}) # no verifiers
528+
d = _doc(status=DomainStatus.ACTIVE)
529+
repo.find_stale_active = AsyncMock(return_value=[d])
530+
531+
result_pairs = await svc.reverify_active(batch_size=10)
532+
533+
assert result_pairs == []
534+
repo.update_status.assert_not_called()

tests/unit/services/test_tenant_resolver.py

Lines changed: 66 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import json
66
from datetime import datetime, timezone
7-
from unittest.mock import AsyncMock
7+
from unittest.mock import AsyncMock, MagicMock
88

99
import pytest
1010
from bson import ObjectId
@@ -140,8 +140,9 @@ async def test_negative_cache_short_circuits_mongo(self):
140140
repo = AsyncMock()
141141
repo.find_active_by_fqdn = AsyncMock(return_value=None)
142142
redis = AsyncMock()
143-
# Second call simulates Redis returning the negative sentinel.
144-
redis.get = AsyncMock(side_effect=[None, "__none__"])
143+
# Cold flow: cache_get → None, tomb check → None (allow write).
144+
# Hot flow: cache_get → "__none__" (negative-hit, no further reads).
145+
redis.get = AsyncMock(side_effect=[None, None, "__none__"])
145146
r = CachedMongoTenantResolver(repo, redis, system_default_domain="spoo.me")
146147

147148
# Cold call: miss → Mongo lookup → negative-cache write
@@ -165,13 +166,24 @@ async def test_negative_sentinel_handled_when_redis_returns_bytes(self):
165166
repo.find_active_by_fqdn.assert_not_called()
166167

167168
@pytest.mark.asyncio
168-
async def test_invalidate_drops_cache_entry(self):
169+
async def test_invalidate_drops_cache_entry_and_sets_tombstone(self):
170+
# Pipelined: delete the cache slot AND set a tombstone so any
171+
# in-flight resolve() can't write the now-stale answer back.
169172
repo = AsyncMock()
170173
redis = AsyncMock()
174+
pipe = MagicMock()
175+
pipe.execute = AsyncMock()
176+
redis.pipeline = MagicMock(return_value=pipe)
171177
r = CachedMongoTenantResolver(repo, redis, system_default_domain="spoo.me")
172178

173179
await r.invalidate("links.acme.com")
174-
redis.delete.assert_awaited_once_with("tenant:links.acme.com")
180+
pipe.delete.assert_called_once_with("tenant:links.acme.com")
181+
pipe.setex.assert_called_once()
182+
# Tombstone key + value
183+
args = pipe.setex.call_args.args
184+
assert args[0] == "tenant_tomb:links.acme.com"
185+
assert args[2] == "1"
186+
pipe.execute.assert_awaited_once()
175187

176188
@pytest.mark.asyncio
177189
async def test_invalidate_skips_system_default(self):
@@ -193,6 +205,55 @@ async def test_invalidate_tolerates_redis_none(self):
193205
# Must not raise.
194206
await r.invalidate("links.acme.com")
195207

208+
@pytest.mark.asyncio
209+
async def test_tombstone_blocks_stale_positive_writeback(self):
210+
# Race: resolve() reads stale Mongo, then invalidate() fires,
211+
# then resolve() tries to write the stale doc back. The tombstone
212+
# set by invalidate() must cause _cache_set to skip the write.
213+
d = _doc()
214+
repo = AsyncMock()
215+
repo.find_active_by_fqdn = AsyncMock(return_value=d)
216+
redis = AsyncMock()
217+
# Cache miss on initial read, then tombstone present at write time.
218+
redis.get = AsyncMock(side_effect=[None, "1"])
219+
r = CachedMongoTenantResolver(repo, redis, system_default_domain="spoo.me")
220+
221+
info = await r.resolve("links.acme.com")
222+
# The current request still gets the answer it computed
223+
assert info is not None
224+
# But the stale answer must NOT have been written back to cache
225+
redis.setex.assert_not_called()
226+
227+
@pytest.mark.asyncio
228+
async def test_tombstone_blocks_stale_negative_writeback(self):
229+
# Same race for the negative path: resolve() reads None from Mongo,
230+
# invalidate() flips the doc to ACTIVE, resolve() must not write
231+
# the stale negative back.
232+
repo = AsyncMock()
233+
repo.find_active_by_fqdn = AsyncMock(return_value=None)
234+
redis = AsyncMock()
235+
redis.get = AsyncMock(side_effect=[None, "1"])
236+
r = CachedMongoTenantResolver(repo, redis, system_default_domain="spoo.me")
237+
238+
result = await r.resolve("nonexistent.example.com")
239+
assert result is None
240+
# No negative-cache write — tombstone forced the skip
241+
redis.setex.assert_not_called()
242+
243+
@pytest.mark.asyncio
244+
async def test_no_tombstone_writes_proceed_normally(self):
245+
# Sanity: when no race fires (no tombstone), cache writes happen.
246+
d = _doc()
247+
repo = AsyncMock()
248+
repo.find_active_by_fqdn = AsyncMock(return_value=d)
249+
redis = AsyncMock()
250+
# Cache miss + no tombstone → write should fire.
251+
redis.get = AsyncMock(side_effect=[None, None])
252+
r = CachedMongoTenantResolver(repo, redis, system_default_domain="spoo.me")
253+
254+
await r.resolve("links.acme.com")
255+
redis.setex.assert_awaited_once()
256+
196257

197258
class TestNormaliseHost:
198259
@pytest.mark.parametrize(

tests/unit/services/verifiers/test_a_record_verifier.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@
22

33
from __future__ import annotations
44

5+
import asyncio
56
from unittest.mock import AsyncMock, MagicMock, patch
67

8+
import dns.exception
79
import dns.resolver
810
import pytest
911

@@ -67,3 +69,31 @@ async def test_no_answer(self):
6769
r = await v.verify("acme.com")
6870
assert r.verified is False
6971
assert "1.2.3.4" in r.reason
72+
73+
@pytest.mark.asyncio
74+
async def test_timeout_returns_failure(self):
75+
v = ARecordVerifier(["1.2.3.4"])
76+
with patch(
77+
"services.verifiers.a_record_verifier.dns.asyncresolver.resolve",
78+
new=AsyncMock(side_effect=asyncio.TimeoutError()),
79+
):
80+
r = await v.verify("slow.com")
81+
assert r.verified is False
82+
assert "timed out" in r.reason
83+
84+
@pytest.mark.asyncio
85+
async def test_generic_dns_exception_swallowed(self):
86+
# Verifiers MUST NOT raise on DNS errors — every DNSException
87+
# subclass must become a VerificationResult(verified=False).
88+
v = ARecordVerifier(["1.2.3.4"])
89+
90+
class WeirdDnsError(dns.exception.DNSException):
91+
pass
92+
93+
with patch(
94+
"services.verifiers.a_record_verifier.dns.asyncresolver.resolve",
95+
new=AsyncMock(side_effect=WeirdDnsError("upstream broke")),
96+
):
97+
r = await v.verify("acme.com")
98+
assert r.verified is False
99+
assert "DNS error" in r.reason

0 commit comments

Comments
 (0)