from __future__ import annotations import time import pytest from yuxi.channels.adapters.nostr.guard import ( AuditRecord, GuardPolicy, NostrGuard, RateLimiter, SeenTracker, ) class TestGuardPolicy: def test_default_values(self): policy = GuardPolicy() assert policy.allowed_kinds == {1, 4, 5, 7, 1059} assert policy.max_ciphertext_bytes == 50_000 assert policy.max_plaintext_bytes == 10_000 assert policy.max_future_skew_sec == 30 assert policy.rate_limit_window_ms == 10_000 assert policy.rate_limit_max_per_sender_per_window == 20 assert policy.rate_limit_max_global_per_window == 200 assert policy.dm_policy == "pairing" assert policy.allow_from == [] def test_custom_values(self): policy = GuardPolicy( allowed_kinds={1, 4}, max_ciphertext_bytes=100, max_plaintext_bytes=50, max_future_skew_sec=60, rate_limit_window_ms=5000, rate_limit_max_per_sender_per_window=5, rate_limit_max_global_per_window=50, dm_policy="open", allow_from=["pub_a"], ) assert policy.allowed_kinds == {1, 4} assert policy.dm_policy == "open" assert policy.allow_from == ["pub_a"] class TestSeenTrackerExtended: def test_mark_seen_twice(self): tracker = SeenTracker(max_size=100) tracker.mark_seen("event_001") tracker.mark_seen("event_001") assert len(tracker) == 1 def test_is_seen_updates_position(self): tracker = SeenTracker(max_size=100, ttl_sec=10) tracker.mark_seen("event_a") tracker.mark_seen("event_b") tracker.mark_seen("event_c") tracker.is_seen("event_a") assert len(tracker) == 3 def test_ttl_expiry_removes_on_check(self): tracker = SeenTracker(max_size=100, ttl_sec=0) tracker.mark_seen("event_001") time.sleep(0.01) assert tracker.is_seen("event_001") is False assert len(tracker) == 0 def test_periodic_eviction(self): tracker = SeenTracker(max_size=5, ttl_sec=0) for i in range(5): tracker.mark_seen(f"event_{i}") time.sleep(0.01) tracker.mark_seen("event_5") assert len(tracker) <= 5 def test_len_property(self): tracker = SeenTracker(max_size=100) tracker.mark_seen("e1") tracker.mark_seen("e2") tracker.mark_seen("e3") assert len(tracker) == 3 class TestRateLimiterExtended: def test_window_expiry(self): limiter = RateLimiter(window_ms=1, max_per_sender=3, max_global=10) limiter.check_and_record("pubkey_a") limiter.check_and_record("pubkey_a") assert limiter.check_and_record("pubkey_a") is True time.sleep(0.002) assert limiter.check_and_record("pubkey_a") is True def test_clear_resets(self): limiter = RateLimiter(max_per_sender=3, max_global=10) limiter.check_and_record("pubkey_a") limiter.check_and_record("pubkey_a") limiter.clear() assert limiter.check_and_record("pubkey_a") is True def test_separate_senders_dont_interfere(self): limiter = RateLimiter(max_per_sender=2, max_global=10) assert limiter.check_and_record("sender_a") is True assert limiter.check_and_record("sender_a") is True assert limiter.check_and_record("sender_a") is False assert limiter.check_and_record("sender_b") is True class TestNostrGuardExtended: @pytest.fixture def guard(self): return NostrGuard(own_pubkey="a" * 64) def test_reject_missing_event_id(self, guard): event = { "kind": 1, "content": "no id", "pubkey": "b" * 64, "created_at": int(time.time()), } reason = guard.check(event) assert reason == "missing event id" def test_reject_empty_event_id(self, guard): event = { "id": "", "kind": 1, "content": "empty id", "pubkey": "b" * 64, "created_at": int(time.time()), } reason = guard.check(event) assert reason == "missing event id" def test_reject_stale_event(self, guard): event = { "id": "stale_event", "kind": 1, "content": "old", "pubkey": "b" * 64, "created_at": 1000, } reason = guard.check(event, since_ts=999999) assert reason is not None assert "stale" in reason def test_reject_dm_disabled(self, guard): guard.policy.dm_policy = "disabled" event = { "id": "disabled_dm", "kind": 1, "content": "hello", "pubkey": "b" * 64, "created_at": int(time.time()), } reason = guard.check(event) assert "disabled" in reason def test_reject_whitelist_empty_allow_from(self, guard): guard.policy.dm_policy = "allowlist" guard.policy.allow_from = [] event = { "id": "whitelist_reject", "kind": 1, "content": "not allowed", "pubkey": "b" * 64, "created_at": int(time.time()), } reason = guard.check(event) assert "whitelist" in reason or "白名单" in reason def test_reject_not_in_allow_from(self): from yuxi.channels.adapters.nostr.guard import NostrGuard, GuardPolicy policy = GuardPolicy(allow_from=["c" * 64]) guard = NostrGuard(own_pubkey="a" * 64, policy=policy) event = { "id": "not_allowed", "kind": 1, "content": "hello", "pubkey": "b" * 64, "created_at": int(time.time()), } reason = guard.check(event) assert "allow_from" in reason def test_pass_in_allow_from(self, guard): guard.policy.allow_from = ["b" * 64] event = { "id": "allowed_event", "kind": 1, "content": "hello", "pubkey": "b" * 64, "created_at": int(time.time()), } reason = guard.check(event) assert reason is None def test_check_plaintext_size_pass(self, guard): reason = guard.check_plaintext_size("short text") assert reason is None def test_check_plaintext_size_reject(self, guard): guard.policy.max_plaintext_bytes = 10 reason = guard.check_plaintext_size("x" * 100) assert reason is not None assert "plaintext too large" in reason def test_done_processing_removes_from_inflight(self, guard): event = { "id": "inflight_test", "kind": 1, "content": "hello", "pubkey": "b" * 64, "created_at": int(time.time()), } guard.check(event) assert len(guard.inflight) == 1 guard.done_processing("inflight_test") assert len(guard.inflight) == 0 def test_inflight_blocks_concurrent(self, guard): event = { "id": "concurrent_event", "kind": 1, "content": "hello", "pubkey": "b" * 64, "created_at": int(time.time()), } assert guard.check(event) is None reason = guard.check(event) assert reason in ("event in-flight (concurrent processing)", "duplicate event") def test_own_pubkey_property(self, guard): assert guard.own_pubkey == "a" * 64 def test_inflight_is_frozenset(self, guard): assert isinstance(guard.inflight, frozenset) def test_audit_log_records_rejections(self, guard): event = { "id": "audit_event", "kind": 999, "content": "bad", "pubkey": "b" * 64, "created_at": int(time.time()), } guard.check(event) audit = guard.get_audit_log() assert len(audit) >= 1 record = audit[0] assert "disallowed" in record["reason"] def test_audit_log_max_size(self, guard): for i in range(1000): event = { "id": f"audit_{i}", "kind": 999, "content": "bad", "pubkey": "b" * 64, "created_at": int(time.time()), } guard.check(event) audit = guard.get_audit_log() assert 500 <= len(audit) <= 1000 def test_clear_audit_log(self, guard): event = { "id": "clear_test", "kind": 999, "content": "bad", "pubkey": "b" * 64, "created_at": int(time.time()), } guard.check(event) assert len(guard.get_audit_log()) > 0 guard.clear_audit_log() assert len(guard.get_audit_log()) == 0 def test_rate_limit_rejection(self): from yuxi.channels.adapters.nostr.guard import ( NostrGuard, GuardPolicy, RateLimiter, ) policy = GuardPolicy(rate_limit_max_per_sender_per_window=2, rate_limit_max_global_per_window=100) guard = NostrGuard(own_pubkey="a" * 64, policy=policy) for i in range(3): event = { "id": f"rate_event_{i}", "kind": 1, "content": "hello", "pubkey": "b" * 64, "created_at": int(time.time()), } guard.check(event) guard.done_processing(f"rate_event_{i}") event = { "id": "rate_over_limit", "kind": 1, "content": "hello", "pubkey": "b" * 64, "created_at": int(time.time()), } reason = guard.check(event) assert reason == "rate limited"