"""Unit tests for the orchestrator bottle registry + attribution (PRD 0070).""" from __future__ import annotations import sqlite3 import tempfile import time import unittest from contextlib import closing from pathlib import Path from bot_bottle.orchestrator.registry import ( BottleRecord, RegistryStore, new_identity_token, ) class TestRegistryStore(unittest.TestCase): def setUp(self) -> None: self._tmp = tempfile.TemporaryDirectory() self.db = Path(self._tmp.name) / "registry.db" self.store = RegistryStore(self.db) self.store.migrate() def tearDown(self) -> None: self._tmp.cleanup() def _force_active_row(self, source_ip: str, bottle_id: str) -> None: """Insert a raw active row, bypassing `register`'s same-IP supersede — the only way to manufacture the ambiguity the fail-closed guard exists for (crash-orphaned / externally-injected duplicate).""" conn = sqlite3.connect(self.db) conn.execute( "INSERT INTO orchestrator_bottles " "(bottle_id, source_ip, identity_token, state, created_at, metadata, policy) " "VALUES (?, ?, ?, 'active', 0.0, '', '')", (bottle_id, source_ip, "tok-" + bottle_id), ) conn.commit() conn.close() def test_register_mints_id_and_token(self) -> None: rec = self.store.register("10.243.0.1") self.assertTrue(rec.bottle_id) self.assertTrue(rec.identity_token) self.assertEqual("active", rec.state) self.assertEqual(rec, self.store.get(rec.bottle_id)) def test_identity_tokens_are_unique(self) -> None: tokens = {new_identity_token() for _ in range(200)} self.assertEqual(200, len(tokens)) a = self.store.register("10.243.0.1") b = self.store.register("10.243.0.3") self.assertNotEqual(a.identity_token, b.identity_token) def test_all_and_deregister(self) -> None: a = self.store.register("10.243.0.1") b = self.store.register("10.243.0.3") self.assertEqual({a.bottle_id, b.bottle_id}, {r.bottle_id for r in self.store.all()}) self.assertTrue(self.store.deregister(a.bottle_id)) self.assertEqual([b.bottle_id], [r.bottle_id for r in self.store.all()]) def test_deregister_missing_is_false(self) -> None: self.assertFalse(self.store.deregister("nope")) def test_explicit_id_replaces(self) -> None: self.store.register("10.243.0.1", bottle_id="fixed", metadata="a") self.store.register("10.243.0.9", bottle_id="fixed", metadata="b") rec = self.store.get("fixed") assert rec is not None self.assertEqual("10.243.0.9", rec.source_ip) self.assertEqual("b", rec.metadata) self.assertEqual(1, len(self.store.all())) def test_attribute_success_requires_ip_and_token(self) -> None: rec = self.store.register("10.243.0.1") got = self.store.attribute("10.243.0.1", rec.identity_token) assert got is not None self.assertEqual(rec.bottle_id, got.bottle_id) def test_attribute_wrong_token_denied(self) -> None: self.store.register("10.243.0.1") self.assertIsNone(self.store.attribute("10.243.0.1", "wrong-token")) def test_attribute_unknown_ip_denied(self) -> None: rec = self.store.register("10.243.0.1") self.assertIsNone(self.store.attribute("10.243.9.9", rec.identity_token)) def test_attribute_empty_token_denied(self) -> None: self.store.register("10.243.0.1") self.assertIsNone(self.store.attribute("10.243.0.1", "")) def test_attribute_ambiguous_ip_denied(self) -> None: # Two active bottles on one source IP is a misconfiguration — deny # rather than guess (fail-closed), even with a valid token. (Forced # directly: `register` now supersedes, so this can only arise from an # orphaned/injected row.) a = self.store.register("10.243.0.1", bottle_id="a") self._force_active_row("10.243.0.1", "b") self.assertIsNone(self.store.attribute("10.243.0.1", a.identity_token)) def test_reregister_same_source_ip_supersedes(self) -> None: # Re-registering a source IP (reused IP / crash recovery / persisted # DB) must supersede the prior active bottle, not add a second active # row that would fail-closed the whole IP to 403. a = self.store.register("10.243.0.1", bottle_id="a", policy="A") b = self.store.register("10.243.0.1", bottle_id="b", policy="B") # exactly one active row at that IP, and it's the new bottle got = self.store.by_source_ip("10.243.0.1") assert got is not None self.assertEqual("b", got.bottle_id) self.assertEqual("B", got.policy) # the superseded bottle is gone and its token no longer attributes self.assertIsNone(self.store.get("a")) self.assertIsNone(self.store.attribute("10.243.0.1", a.identity_token)) self.assertIsNotNone(self.store.attribute("10.243.0.1", b.identity_token)) def test_state_persists_across_reopen(self) -> None: rec = self.store.register("10.243.0.1") reopened = RegistryStore(self.db) got = reopened.get(rec.bottle_id) assert got is not None self.assertEqual(rec, got) # Attribution works against the reopened (durable) store too. self.assertIsNotNone(reopened.attribute("10.243.0.1", rec.identity_token)) def test_redacted_hides_token(self) -> None: rec = BottleRecord( bottle_id="x", source_ip="10.243.0.1", identity_token="secret" ) self.assertNotIn("identity_token", rec.redacted()) self.assertEqual("10.243.0.1", rec.redacted()["source_ip"]) def test_register_with_policy_resolves_via_attribution(self) -> None: rec = self.store.register("10.243.0.1", policy='{"allow":["x"]}') self.assertEqual('{"allow":["x"]}', rec.policy) got = self.store.attribute("10.243.0.1", rec.identity_token) assert got is not None self.assertEqual('{"allow":["x"]}', got.policy) def test_set_policy_updates_live(self) -> None: rec = self.store.register("10.243.0.1") self.assertEqual("", rec.policy) self.assertTrue(self.store.set_policy(rec.bottle_id, '{"v":2}')) got = self.store.get(rec.bottle_id) assert got is not None self.assertEqual('{"v":2}', got.policy) def test_set_policy_missing_is_false(self) -> None: self.assertFalse(self.store.set_policy("ghost", "{}")) def test_policy_defaults_empty_and_persists(self) -> None: rec = self.store.register("10.243.0.1") got = RegistryStore(self.db).get(rec.bottle_id) assert got is not None self.assertEqual("", got.policy) def test_by_source_ip_returns_single_active(self) -> None: rec = self.store.register("10.243.0.1") got = self.store.by_source_ip("10.243.0.1") assert got is not None self.assertEqual(rec.bottle_id, got.bottle_id) def test_by_source_ip_unknown_or_ambiguous_denied(self) -> None: self.assertIsNone(self.store.by_source_ip("10.243.9.9")) # unknown self.store.register("10.243.0.1", bottle_id="a") self._force_active_row("10.243.0.1", "b") # orphaned duplicate self.assertIsNone(self.store.by_source_ip("10.243.0.1")) # ambiguous if __name__ == "__main__": unittest.main() class TestReapAbsent(unittest.TestCase): """`reap_absent` — the self-heal for rows whose bottle is gone. An orphan is not merely untidy: source IPs get recycled, and `by_source_ip` fail-closes on ambiguity, so a leftover row at a reused address resolves *no* policy for the next bottle that lands there and every host it asks for is denied. """ def setUp(self) -> None: self._tmp = tempfile.TemporaryDirectory() self.db = Path(self._tmp.name) / "registry.db" self.store = RegistryStore(self.db) self.store.migrate() def tearDown(self) -> None: self._tmp.cleanup() def _aged(self, source_ip: str, *, age: float) -> BottleRecord: """Register a bottle and backdate it past the grace window.""" rec = self.store.register(source_ip) with closing(sqlite3.connect(self.db)) as conn: conn.execute( "UPDATE orchestrator_bottles SET created_at = ? WHERE bottle_id = ?", (time.time() - age, rec.bottle_id), ) conn.commit() return rec def test_reaps_row_with_no_live_container(self) -> None: gone = self._aged("10.243.0.9", age=600) reaped = self.store.reap_absent([]) self.assertEqual([gone.bottle_id], [r.bottle_id for r in reaped]) self.assertIsNone(self.store.get(gone.bottle_id)) def test_keeps_row_whose_ip_is_live(self) -> None: alive = self._aged("10.243.0.9", age=600) self.assertEqual([], self.store.reap_absent(["10.243.0.9"])) self.assertIsNotNone(self.store.get(alive.bottle_id)) def test_grace_window_protects_an_in_flight_launch(self) -> None: """A bottle registered moments ago is never reaped, even though the caller's enumeration didn't see its address yet.""" fresh = self.store.register("10.243.0.10") self.assertEqual([], self.store.reap_absent([])) self.assertIsNotNone(self.store.get(fresh.bottle_id)) def test_reaping_the_orphan_unbricks_the_reused_address(self) -> None: """The regression this exists for: an orphan at an address that vmnet later hands to a new bottle makes `by_source_ip` ambiguous, so the new bottle resolves no policy at all.""" orphan = self._aged("10.243.0.11", age=600) # A new bottle lands on the recycled address. Force the row in directly # so `register`'s own supersede sweep doesn't mask the ambiguity. with closing(sqlite3.connect(self.db)) as conn: conn.execute( "INSERT INTO orchestrator_bottles " "(bottle_id, source_ip, identity_token, state, created_at, metadata, policy) " "VALUES ('newbottle', '10.243.0.11', 'tok-new', 'active', ?, '', 'routes: []')", (time.time(),), ) conn.commit() self.assertIsNone(self.store.by_source_ip("10.243.0.11")) # bricked reaped = self.store.reap_absent(["10.243.0.11"], grace_seconds=60) self.assertEqual([orphan.bottle_id], [r.bottle_id for r in reaped]) rec = self.store.by_source_ip("10.243.0.11") assert rec is not None self.assertEqual("newbottle", rec.bottle_id) def test_ignores_empty_ips_in_the_live_set(self) -> None: gone = self._aged("10.243.0.12", age=600) self.assertEqual( [gone.bottle_id], [r.bottle_id for r in self.store.reap_absent(["", "10.243.0.99"])], )