From fe62f4eaf479fbfa2f6a035f4059e36b4be03018 Mon Sep 17 00:00:00 2001 From: aidankhogg Date: Tue, 30 Jun 2026 00:42:16 +0100 Subject: [PATCH] Strengthen AND provisioning leases and reconciliation --- migrations/001_initial.sql | 3 +- netengine/core/state.py | 15 ++ netengine/handlers/domain_registry_handler.py | 69 ++++++- netengine/phases/phase_ands.py | 184 ++++++++++++++++-- tests/test_phase_ands_strengthening.py | 184 ++++++++++++++++++ 5 files changed, 433 insertions(+), 22 deletions(-) create mode 100644 tests/test_phase_ands_strengthening.py diff --git a/migrations/001_initial.sql b/migrations/001_initial.sql index b592961..fecf86c 100644 --- a/migrations/001_initial.sql +++ b/migrations/001_initial.sql @@ -29,7 +29,8 @@ CREATE TABLE IF NOT EXISTS address_pools ( CREATE TABLE IF NOT EXISTS address_leases ( and_name TEXT PRIMARY KEY, cidr CIDR NOT NULL, - assigned_at TIMESTAMPTZ DEFAULT NOW() + assigned_at TIMESTAMPTZ DEFAULT NOW(), + CONSTRAINT address_leases_cidr_unique UNIQUE (cidr) ); -- Domain records diff --git a/netengine/core/state.py b/netengine/core/state.py index e062a90..21193a7 100644 --- a/netengine/core/state.py +++ b/netengine/core/state.py @@ -164,6 +164,20 @@ class DomainRegistryOutput(PhaseOutputBase, total=False): tld_delegations: list[dict[str, JsonValue]] +class ANDInstanceState(TypedDict, total=False): + name: str + org: str + profile: str + cidr: str + gateway_ip: str + bridge_name: str + dns_suffix: str + deployed_at: str + dynamic_ip: bool + reverse_dns: bool + bgp: str | None + + class GenericPhaseOutput(PhaseOutputBase, total=False): status: str healthy: bool @@ -221,6 +235,7 @@ class RuntimeState: domain_registry_output: Optional[DomainRegistryOutput] = None identity_inworld_output: Optional[GenericPhaseOutput] = None ands_output: Optional[GenericPhaseOutput] = None + ands_instances: Dict[str, ANDInstanceState] = field(default_factory=dict) world_services_output: Optional[GenericPhaseOutput] = None org_apps_output: Optional[GenericPhaseOutput] = None diff --git a/netengine/handlers/domain_registry_handler.py b/netengine/handlers/domain_registry_handler.py index 54e40ca..9759a8e 100644 --- a/netengine/handlers/domain_registry_handler.py +++ b/netengine/handlers/domain_registry_handler.py @@ -1,4 +1,5 @@ # netengine/handlers/domain_registry_handler.py +import ipaddress from typing import Any, List, cast from netengine.core.pgmq_client import PGMQClient @@ -35,14 +36,70 @@ async def seed_address_pools(self, spec: Any) -> None: ).execute() async def allocate_address(self, and_name: str, profile: str) -> str: - """Allocate a CIDR from the pool; row-level lock prevents conflicts.""" + """Allocate a unique AND subnet from the registry phase address pool. + + Existing leases are idempotently reused by ``and_name``. New leases are + allocated as /24s within the profile pool (or the pool itself when it is + smaller than /24). The ``address_leases.cidr`` unique constraint is the + database-backed collision guard; if another allocator wins the same CIDR, + this method refreshes leases and tries the next free candidate. + """ db = await self._get_db() - result = await db.table("address_pools").select("cidr").eq("profile", profile).execute() - if not result.data: + + existing = ( + await db.table("address_leases") + .select("and_name,cidr") + .eq("and_name", and_name) + .execute() + ) + if existing.data and existing.data[0].get("and_name") == and_name: + return cast(str, existing.data[0]["cidr"]) + + pool_result = ( + await db.table("address_pools").select("cidr").eq("profile", profile).execute() + ) + if not pool_result.data: raise RegistryError(f"No address pool for profile {profile}") - pool_cidr = result.data[0]["cidr"] - await db.table("address_leases").upsert({"and_name": and_name, "cidr": pool_cidr}).execute() - return cast(str, pool_cidr) + + pool = ipaddress.ip_network(str(pool_result.data[0]["cidr"]), strict=False) + prefix = max(pool.prefixlen, 24) + candidates = [pool] if pool.prefixlen > 24 else list(pool.subnets(new_prefix=prefix)) + + last_error: Exception | None = None + for _ in range(2): + lease_result = await db.table("address_leases").select("and_name,cidr").execute() + used = { + ipaddress.ip_network(str(row["cidr"]), strict=False) + for row in lease_result.data + if "and_name" in row + } + for candidate in candidates: + if candidate in used: + continue + cidr = str(candidate) + try: + await db.table("address_leases").upsert( + {"and_name": and_name, "cidr": cidr, "profile": profile} + ).execute() + except Exception as exc: + last_error = exc + if "profile" in str(exc).lower(): + try: + await db.table("address_leases").upsert( + {"and_name": and_name, "cidr": cidr} + ).execute() + except Exception as retry_exc: + last_error = retry_exc + continue + else: + continue + return cidr + + if last_error is not None: + raise RegistryError( + f"Address pool exhausted for profile {profile}; last collision: {last_error}" + ) + raise RegistryError(f"Address pool exhausted for profile {profile}") async def register_domain(self, domain: str, org_name: str, ns_records: List[str]) -> None: """Register a domain; emit DNS update event.""" diff --git a/netengine/phases/phase_ands.py b/netengine/phases/phase_ands.py index 0e5b017..6d87f07 100644 --- a/netengine/phases/phase_ands.py +++ b/netengine/phases/phase_ands.py @@ -21,6 +21,7 @@ from netengine.handlers._base import BasePhaseHandler from netengine.handlers.context import PhaseContext from netengine.handlers.dns import DNSHandler +from netengine.handlers.domain_registry_handler import DomainRegistryHandler from netengine.handlers.docker_handler import DockerHandler from netengine.handlers.gateway_handler import GatewayHandler @@ -116,9 +117,7 @@ async def execute(self, context: PhaseContext) -> None: profiles_used.add(and_instance.profile) # Store in runtime_state for this session - if not hasattr(context.runtime_state, "ands_instances"): - context.runtime_state.ands_instances = {} # type: ignore[attr-defined] - context.runtime_state.ands_instances[and_instance.name] = and_data # type: ignore[attr-defined] + context.runtime_state.ands_instances[and_instance.name] = and_data ands_output["ands_provisioned"] = ands_provisioned ands_output["address_allocations"] = address_allocations @@ -168,11 +167,7 @@ async def healthcheck(self, context: PhaseContext) -> bool: logger.warning("No ANDs provisioned") return False - if not hasattr(context.runtime_state, "ands_instances"): - logger.warning("AND instances not tracked in runtime_state") - return False - - instances = context.runtime_state.ands_instances # type: ignore[attr-defined] + instances = context.runtime_state.ands_instances if not all(and_name in instances for and_name in ands_provisioned): logger.warning("Some AND instances missing from runtime_state") return False @@ -224,7 +219,7 @@ async def _provision_and( raise RuntimeError(f"Profile '{profile_name}' not found in spec") # 2. Allocate subnet - cidr = await self._allocate_address(and_instance.name, profile_name, ands_spec) + cidr = await self._allocate_address(context, and_instance.name, profile_name, ands_spec) logger.info(f"Allocated CIDR for {and_instance.name}: {cidr}") # 3. Create isolated Docker bridge network @@ -339,6 +334,9 @@ async def _provision_and( "bridge_name": bridge_name, "dns_suffix": dns_suffix, "deployed_at": datetime.now(UTC).isoformat(), + "dynamic_ip": bool(profile_obj.dynamic_ip), + "reverse_dns": bool(profile_obj.reverse_dns), + "bgp": profile_obj.bgp, } try: @@ -352,6 +350,7 @@ async def _provision_and( async def _allocate_address( self, + context: PhaseContext, and_name: str, profile: str, ands_spec: Any, @@ -361,12 +360,167 @@ async def _allocate_address( if profile_obj is None: raise RuntimeError(f"Profile '{profile}' not found") - # Sequential /24 allocation within 172.16.0.0/12. - # Uses ord-sum rather than hash() to avoid collision with >256 ANDs. - idx = sum(ord(c) for c in and_name) % 4096 - third_octet = idx % 256 - second_extra = idx // 256 - return f"172.{16 + second_extra}.{third_octet}.0/24" + try: + registry = DomainRegistryHandler(context=context) + return await registry.allocate_address(and_name, profile) + except Exception as exc: + # Tests and mock-mode executions may not have a DB. Fall back only when + # the registry output does not expose seeded pools; otherwise surface + # the allocation failure so exhaustion/collisions are not hidden. + output = context.runtime_state.domain_registry_output or {} + if output.get("address_pools") or output.get("address_pools_seeded"): + raise RuntimeError(f"Failed to allocate address for {and_name}: {exc}") from exc + context.logger.warning("Registry allocation unavailable; using mock CIDR: %s", exc) + idx = len(context.runtime_state.ands_instances) + return f"172.31.{idx}.0/24" + + async def reconcile( + self, + context: PhaseContext, + docker: DockerHandler | None = None, + gateway: GatewayHandler | None = None, + ) -> dict[str, list[str]]: + """Reconcile desired AND spec against runtime/Docker/gateway/DNS side effects.""" + docker = docker or DockerHandler() + gateway = gateway or GatewayHandler(docker) + desired = {inst.name: inst for inst in context.spec.ands.instances} + existing = context.runtime_state.ands_instances + actions: dict[str, list[str]] = { + "created": [], + "updated": [], + "removed": [], + "repaired": [], + } + + for and_name in set(existing) - set(desired): + await self._teardown_and(context, docker, gateway, and_name, existing[and_name]) + actions["removed"].append(and_name) + + for and_name, inst in desired.items(): + current = existing.get(and_name) + if current is None: + data = await self._provision_and(context, docker, gateway, inst, context.spec.ands) + context.runtime_state.ands_instances[and_name] = data + actions["created"].append(and_name) + continue + if current.get("profile") != inst.profile: + await self._update_and_profile(context, gateway, inst, current, context.spec.ands) + actions["updated"].append(and_name) + repaired = await self._repair_and( + context, docker, gateway, inst, current, context.spec.ands + ) + if repaired: + actions["repaired"].append(and_name) + return actions + + async def _update_and_profile( + self, + context: PhaseContext, + gateway: GatewayHandler, + and_instance: Any, + current: dict[str, Any], + ands_spec: Any, + ) -> None: + old_profile = ands_spec.profiles.get(current.get("profile")) if ands_spec.profiles else None + new_profile = ands_spec.profiles.get(and_instance.profile) if ands_spec.profiles else None + if new_profile is None: + raise RuntimeError(f"Profile '{and_instance.profile}' not found in spec") + if old_profile and getattr(old_profile, "dynamic_ip", False) and not new_profile.dynamic_ip: + await gateway.remove_dhcp(and_instance.name) + if old_profile and getattr(old_profile, "bgp", None) and not new_profile.bgp: + await gateway.remove_bgp(and_instance.name) + rules = await gateway.generate_rules( + and_instance.name, and_instance.profile, current["cidr"] + ) + await gateway.apply_rules(and_instance.name, rules) + if new_profile.dynamic_ip: + await gateway.setup_dhcp(and_instance.name, current["cidr"], current["gateway_ip"]) + if new_profile.reverse_dns: + await DNSHandler().add_reverse_zone(context, current["cidr"], current["gateway_ip"]) + if new_profile.bgp: + await gateway.setup_bgp( + and_instance.name, current["cidr"], current["gateway_ip"], new_profile.bgp + ) + current.update( + { + "profile": and_instance.profile, + "dynamic_ip": new_profile.dynamic_ip, + "reverse_dns": new_profile.reverse_dns, + "bgp": new_profile.bgp, + } + ) + + async def _teardown_and( + self, + context: PhaseContext, + docker: DockerHandler, + gateway: GatewayHandler, + and_name: str, + current: dict[str, Any], + ) -> None: + await gateway.remove_rules(and_name) + if current.get("dynamic_ip"): + await gateway.remove_dhcp(and_name) + if current.get("reverse_dns"): + await DNSHandler().remove_reverse_zone(context, current["cidr"]) + if current.get("bgp"): + await gateway.remove_bgp(and_name) + bridge_name = current.get("bridge_name", f"netengines_and_{and_name}") + try: + await docker.disconnect_network(gateway.gateway_container, bridge_name) + except Exception: + pass + await docker.remove_network(bridge_name) + context.runtime_state.ands_instances.pop(and_name, None) + try: + db = await self._get_supabase() + await db.table("address_leases").delete().eq("and_name", and_name).execute() + await db.table("and_instances").delete().eq("name", and_name).execute() + except Exception as exc: + context.logger.warning("Failed to delete AND DB state for %s: %s", and_name, exc) + + async def _repair_and( + self, + context: PhaseContext, + docker: DockerHandler, + gateway: GatewayHandler, + and_instance: Any, + current: dict[str, Any], + ands_spec: Any, + ) -> bool: + repaired = False + bridge_name = current.get("bridge_name", f"netengines_and_{and_instance.name}") + try: + docker.client.networks.get(bridge_name) + except Exception: + await docker.create_network( + name=bridge_name, driver="bridge", subnet=current["cidr"], internal=True + ) + repaired = True + try: + await docker.connect_network( + container=gateway.gateway_container, network=bridge_name, ip=current["gateway_ip"] + ) + except Exception: + pass + profile = ands_spec.profiles.get(and_instance.profile) if ands_spec.profiles else None + rules = await gateway.generate_rules( + and_instance.name, and_instance.profile, current["cidr"] + ) + await gateway.apply_rules(and_instance.name, rules) + dns_suffix = getattr(and_instance, "dns_suffix", None) or current.get("dns_suffix") + await DNSHandler().add_zone_record( + context, "internal", "A", dns_suffix.rstrip("."), current["gateway_ip"], 300 + ) + if profile and profile.dynamic_ip: + await gateway.setup_dhcp(and_instance.name, current["cidr"], current["gateway_ip"]) + if profile and profile.reverse_dns: + await DNSHandler().add_reverse_zone(context, current["cidr"], current["gateway_ip"]) + if profile and profile.bgp: + await gateway.setup_bgp( + and_instance.name, current["cidr"], current["gateway_ip"], profile.bgp + ) + return repaired # ───────────────────────────────────────────── # Event-Driven Provisioning diff --git a/tests/test_phase_ands_strengthening.py b/tests/test_phase_ands_strengthening.py new file mode 100644 index 0000000..ba092f2 --- /dev/null +++ b/tests/test_phase_ands_strengthening.py @@ -0,0 +1,184 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from netengine.core.state import RuntimeState +from netengine.errors import RegistryError +from netengine.handlers.context import PhaseContext +from netengine.handlers.domain_registry_handler import DomainRegistryHandler +from netengine.logs import get_logger +from netengine.phases.phase_ands import ANDsPhaseHandler + + +class Result: + def __init__(self, data): + self.data = data + + +class FakeTable: + def __init__(self, db, name): + self.db = db + self.name = name + self.op = None + self.data = None + self.filters = [] + + def select(self, cols="*"): + self.op = "select" + return self + + def insert(self, data): + self.op = "insert" + self.data = data + return self + + def delete(self): + self.op = "delete" + return self + + def upsert(self, data): + self.op = "upsert" + self.data = data + return self + + def eq(self, col, val): + self.filters.append((col, val)) + return self + + async def execute(self): + rows = self.db.setdefault(self.name, []) + if self.op == "select": + out = rows + for col, val in self.filters: + out = [r for r in out if r.get(col) == val] + return Result([dict(r) for r in out]) + if self.op == "insert": + if self.name == "address_leases": + if any(r["and_name"] == self.data["and_name"] for r in rows): + raise Exception("duplicate and_name") + if any(r["cidr"] == self.data["cidr"] for r in rows): + raise Exception("duplicate cidr") + rows.append(dict(self.data)) + return Result([dict(self.data)]) + if self.op == "delete": + keep = [] + deleted = [] + for r in rows: + match = all(r.get(c) == v for c, v in self.filters) + (deleted if match else keep).append(r) + self.db[self.name] = keep + return Result(deleted) + rows.append(dict(self.data)) + return Result([dict(self.data)]) + + +class FakeDB(dict): + def table(self, name): + return FakeTable(self, name) + + +@pytest.mark.asyncio +async def test_registry_allocates_unique_subnets_and_reports_exhaustion(): + db = FakeDB(address_pools=[{"profile": "tiny", "cidr": "10.0.0.0/23"}], address_leases=[]) + handler = DomainRegistryHandler(pgmq=MagicMock()) + handler._db = db + + assert await handler.allocate_address("and-a", "tiny") == "10.0.0.0/24" + assert await handler.allocate_address("and-b", "tiny") == "10.0.1.0/24" + with pytest.raises(RegistryError, match="exhausted"): + await handler.allocate_address("and-c", "tiny") + + +@pytest.mark.asyncio +async def test_registry_skips_cidr_collisions_when_existing_lease_uses_candidate(): + db = FakeDB( + address_pools=[{"profile": "biz", "cidr": "10.10.0.0/22"}], + address_leases=[{"and_name": "other", "cidr": "10.10.0.0/24"}], + ) + handler = DomainRegistryHandler(pgmq=MagicMock()) + handler._db = db + + assert await handler.allocate_address("new-and", "biz") == "10.10.1.0/24" + + +def _context(profile): + state = RuntimeState() + state.substrate_output = {"networks": {}} + state.dns_output = {"root_zone": {}, "zone_files": {"internal": "$ORIGIN internal.\n"}} + state.domain_registry_output = {"pools": {}} + and_inst = SimpleNamespace( + name="lab", org="labco", profile="business", dns_suffix="lab.internal" + ) + spec = SimpleNamespace( + ands=SimpleNamespace(profiles={"business": profile}, instances=[and_inst]) + ) + return PhaseContext(runtime_state=state, spec=spec, logger=get_logger("test")) + + +@pytest.mark.asyncio +async def test_profile_features_configure_dynamic_ip_reverse_dns_and_optional_bgp(): + profile = SimpleNamespace(dynamic_ip=True, reverse_dns=True, bgp="optional") + ctx = _context(profile) + handler = ANDsPhaseHandler() + docker = AsyncMock() + gateway = AsyncMock(gateway_container="gw") + gateway.generate_rules.return_value = "rules" + gateway.setup_bgp.side_effect = RuntimeError("sidecar unavailable") + + with patch.object(handler, "_allocate_address", AsyncMock(return_value="172.20.1.0/24")): + await handler._provision_and( + ctx, docker, gateway, ctx.spec.ands.instances[0], ctx.spec.ands + ) + + gateway.setup_dhcp.assert_awaited_once() + gateway.setup_bgp.assert_awaited_once() + assert "1.20.172.in-addr.arpa" in ctx.runtime_state.dns_output["zone_files"] + + +@pytest.mark.asyncio +async def test_required_bgp_failure_aborts_provisioning(): + profile = SimpleNamespace(dynamic_ip=False, reverse_dns=False, bgp="required") + ctx = _context(profile) + handler = ANDsPhaseHandler() + docker = AsyncMock() + gateway = AsyncMock(gateway_container="gw") + gateway.generate_rules.return_value = "rules" + gateway.setup_bgp.side_effect = RuntimeError("sidecar unavailable") + + with patch.object(handler, "_allocate_address", AsyncMock(return_value="172.20.2.0/24")): + with pytest.raises(RuntimeError, match="Required BGP setup failed"): + await handler._provision_and( + ctx, docker, gateway, ctx.spec.ands.instances[0], ctx.spec.ands + ) + + +@pytest.mark.asyncio +async def test_reconcile_repairs_missing_network_after_partial_failure(): + profile = SimpleNamespace(dynamic_ip=True, reverse_dns=True, bgp=None) + ctx = _context(profile) + ctx.runtime_state.ands_instances["lab"] = { + "name": "lab", + "org": "labco", + "profile": "business", + "cidr": "172.20.3.0/24", + "gateway_ip": "172.20.3.1", + "bridge_name": "netengines_and_lab", + "dns_suffix": "lab.internal", + "dynamic_ip": True, + "reverse_dns": True, + "bgp": None, + } + handler = ANDsPhaseHandler() + docker = AsyncMock() + docker.client = MagicMock() + docker.client.networks.get.side_effect = Exception("missing") + gateway = AsyncMock(gateway_container="gw") + gateway.generate_rules.return_value = "rules" + + actions = await handler.reconcile(ctx, docker, gateway) + + assert actions["repaired"] == ["lab"] + docker.create_network.assert_awaited_once() + gateway.setup_dhcp.assert_awaited_once() + assert "3.20.172.in-addr.arpa" in ctx.runtime_state.dns_output["zone_files"]