From 8ea834e10818ba89a36cc2974c79a95fbb031401 Mon Sep 17 00:00:00 2001 From: Dmitrii Iurco Date: Fri, 1 May 2026 16:23:31 -0400 Subject: [PATCH] feat: Phase 3 - per-peer internet routing via exit cell Adds the ability to route a specific peer's internet traffic through a connected cell acting as an exit relay. Cell A side: - PUT /api/peers//route-via {"via_cell": "cellB"} sets route_via - Updates WG AllowedIPs to include 0.0.0.0/0 for the exit cell peer - Adds ip rule + ip route in policy table inside cell-wireguard so the specific peer's traffic egresses via cellB's WG IP - Sets exit_relay_active on the cell link and pushes use_as_exit_relay=True to cellB via peer-sync Cell B side: - Receives use_as_exit_relay in the peer-sync payload - Calls apply_cell_rules(..., exit_relay=True) to add FORWARD -o eth0 ACCEPT - Stores remote_exit_relay_active flag for startup recovery Startup recovery: - apply_all_cell_rules passes exit_relay=remote_exit_relay_active (cellB) - _apply_startup_enforcement reapplies ip rule for each peer with route_via (cellA) since policy routing rules don't survive container restart peer_registry gets route_via field with lazy migration. 22 new tests across test_cell_link_manager, test_peer_registry, test_peer_route_via. Co-Authored-By: Claude Sonnet 4.6 --- api/app.py | 13 +++ api/cell_link_manager.py | 31 ++++++- api/firewall_manager.py | 31 +++++-- api/peer_registry.py | 20 ++++ api/routes/cells.py | 5 +- api/routes/peers.py | 71 ++++++++++++++ api/wireguard_manager.py | 74 +++++++++++++++ tests/test_cell_link_manager.py | 122 +++++++++++++++++++++++- tests/test_cells_endpoints.py | 1 + tests/test_peer_registry.py | 30 ++++++ tests/test_peer_route_via.py | 160 ++++++++++++++++++++++++++++++++ 11 files changed, 547 insertions(+), 11 deletions(-) create mode 100644 tests/test_peer_route_via.py diff --git a/api/app.py b/api/app.py index 70261d8..197f1d7 100644 --- a/api/app.py +++ b/api/app.py @@ -275,6 +275,19 @@ def _apply_startup_enforcement(): firewall_manager.apply_all_dns_rules(peers, COREFILE_PATH, _configured_domain(), cell_links=cell_links) logger.info(f"Applied enforcement rules for {len(peers)} peers, {len(cell_links)} cells on startup") + # Phase 3: reapply policy routing rules for peers whose internet traffic is + # routed through an exit cell (ip rule entries don't survive container restart) + cell_links_map = {l['cell_name']: l for l in cell_links} + for peer in peers: + via_cell = peer.get('route_via') + if not via_cell: + continue + link = cell_links_map.get(via_cell) + if not link: + continue + peer_ip = peer.get('ip', '').split('/')[0] + if peer_ip: + wireguard_manager.apply_peer_route_via(peer_ip, via_wg_ip=link['dns_ip']) sync_summary = cell_link_manager.replay_pending_pushes() if sync_summary.get('attempted'): logger.info(f"Startup permission sync: {sync_summary}") diff --git a/api/cell_link_manager.py b/api/cell_link_manager.py index 5052189..02e3fe0 100644 --- a/api/cell_link_manager.py +++ b/api/cell_link_manager.py @@ -84,6 +84,13 @@ class CellLinkManager: if 'remote_exit_offered' not in link: link['remote_exit_offered'] = False changed = True + # Phase 3 migration: per-peer internet routing + if 'exit_relay_active' not in link: + link['exit_relay_active'] = False + changed = True + if 'remote_exit_relay_active' not in link: + link['remote_exit_relay_active'] = False + changed = True if changed: self._save(links) return links @@ -154,6 +161,7 @@ class CellLinkManager: 'inbound': dict(perms.get('outbound', {})), }, 'exit_offered': bool(link.get('exit_offered', False)), + 'use_as_exit_relay': bool(link.get('exit_relay_active', False)), 'sent_at': datetime.utcnow().isoformat() + 'Z', } payload = json.dumps(body) @@ -239,7 +247,8 @@ class CellLinkManager: def apply_remote_permissions(self, from_public_key: str, permissions: Dict[str, Any], - exit_offered: bool = False) -> Dict[str, Any]: + exit_offered: bool = False, + use_as_exit_relay: bool = False) -> Dict[str, Any]: """Store permissions pushed by a remote cell (identified by WG public key). Validates service names, persists, and re-applies local iptables rules. @@ -257,13 +266,15 @@ class CellLinkManager: link['permissions'] = {'inbound': clean_inbound, 'outbound': clean_outbound} link['remote_exit_offered'] = bool(exit_offered) + link['remote_exit_relay_active'] = bool(use_as_exit_relay) link['last_remote_update_at'] = datetime.utcnow().isoformat() self._save(links) inbound_list = [s for s, v in clean_inbound.items() if v] try: import firewall_manager as _fm - _fm.apply_cell_rules(link['cell_name'], link['vpn_subnet'], inbound_list) + _fm.apply_cell_rules(link['cell_name'], link['vpn_subnet'], inbound_list, + exit_relay=use_as_exit_relay) except Exception as e: logger.warning( f"apply_cell_rules after remote push for '{link['cell_name']}' " @@ -470,6 +481,22 @@ class CellLinkManager: self._try_push(cell_name, link) return link + def set_exit_relay_active(self, cell_name: str, active: bool) -> Dict[str, Any]: + """Record that THIS cell is routing a peer's internet traffic via cell_name. + + Persists the flag locally and pushes updated state to the remote cell so + it can enable/disable the FORWARD-to-eth0 rule on its side. + Returns the updated link record. + """ + links = self._load() + link = next((l for l in links if l['cell_name'] == cell_name), None) + if not link: + raise ValueError(f"Cell '{cell_name}' not found") + link['exit_relay_active'] = bool(active) + self._save(links) + self._try_push(cell_name, link) + return link + def get_connection_status(self, cell_name: str) -> Dict[str, Any]: """Return link record enriched with live WireGuard handshake status.""" links = self._load() diff --git a/api/firewall_manager.py b/api/firewall_manager.py index 2dd5f64..14efac9 100644 --- a/api/firewall_manager.py +++ b/api/firewall_manager.py @@ -262,7 +262,8 @@ def _get_cell_api_ip() -> Optional[str]: return r.stdout.strip() -def apply_cell_rules(cell_name: str, vpn_subnet: str, inbound_services: List[str]) -> bool: +def apply_cell_rules(cell_name: str, vpn_subnet: str, inbound_services: List[str], + exit_relay: bool = False) -> bool: """Apply FORWARD rules for a cell-to-cell peer. Traffic from vpn_subnet is allowed only to service VIPs listed in @@ -270,10 +271,15 @@ def apply_cell_rules(cell_name: str, vpn_subnet: str, inbound_services: List[str internet or peer access — only explicit service VIPs, plus the cell-api port (3000) for permission-sync pushes arriving via DNAT. - Rule insertion order (last inserted → top of chain): - 1. Catch-all DROP for the subnet (inserted first → bottom) - 2. Per-service ACCEPT/DROP (inserted in reversed() order) - 3. API-sync ACCEPT (inserted last → top, above catch-all) + When exit_relay=True, the remote cell's peers can route internet + traffic through this cell (Phase 3). A broad ACCEPT for traffic + going out eth0 is added below per-service rules but above catch-all. + + Rule insertion order (first inserted = bottom, last inserted = top): + 1. Catch-all DROP for the subnet (inserted first → bottom) + 2. Exit relay ACCEPT (-o eth0) (if exit_relay, above catch-all) + 3. Per-service ACCEPT/DROP (inserted in reversed() order) + 4. API-sync ACCEPT (inserted last → top) """ try: tag = _cell_tag(cell_name) @@ -283,6 +289,13 @@ def apply_cell_rules(cell_name: str, vpn_subnet: str, inbound_services: List[str _iptables(['-I', 'FORWARD', '-s', vpn_subnet, '-m', 'comment', '--comment', tag, '-j', 'DROP']) + # Exit relay ACCEPT — allow internet-bound traffic from this cell's peers. + # Inserted ABOVE catch-all but BELOW per-service rules so service-level + # DROP rules still take effect for specific service VIPs. + if exit_relay: + _iptables(['-I', 'FORWARD', '-s', vpn_subnet, '-o', 'eth0', + '-m', 'comment', '--comment', tag, '-j', 'ACCEPT']) + # Per-service rules — inserted in reverse dict order, highest-priority last for service, svc_ip in reversed(list(SERVICE_IPS.items())): target = 'ACCEPT' if service in inbound_services else 'DROP' @@ -298,7 +311,10 @@ def apply_cell_rules(cell_name: str, vpn_subnet: str, inbound_services: List[str '-p', 'tcp', '--dport', '3000', '-m', 'comment', '--comment', tag, '-j', 'ACCEPT']) - logger.info(f"Applied cell rules for {cell_name} ({vpn_subnet}): inbound={inbound_services}") + logger.info( + f"Applied cell rules for {cell_name} ({vpn_subnet}): " + f"inbound={inbound_services} exit_relay={exit_relay}" + ) return True except Exception as e: logger.error(f"apply_cell_rules({cell_name}): {e}") @@ -314,7 +330,8 @@ def apply_all_cell_rules(cell_links: List[Dict[str, Any]]) -> None: continue perms = link.get('permissions', {}) inbound = [s for s, v in perms.get('inbound', {}).items() if v] - apply_cell_rules(name, subnet, inbound) + exit_relay = bool(link.get('remote_exit_relay_active', False)) + apply_cell_rules(name, subnet, inbound, exit_relay=exit_relay) def ensure_cell_api_dnat() -> bool: diff --git a/api/peer_registry.py b/api/peer_registry.py index 86a14c2..bd5083c 100644 --- a/api/peer_registry.py +++ b/api/peer_registry.py @@ -193,6 +193,14 @@ class PeerRegistry(BaseServiceManager): except Exception as e: self.logger.error(f"Error loading peers: {e}") self.peers = [] + # Phase 3 migration: per-peer internet routing + changed = False + for peer in self.peers: + if 'route_via' not in peer: + peer['route_via'] = None + changed = True + if changed: + self._save_peers() else: self.peers = [] self.logger.info("No peers file found, starting with empty registry") @@ -326,6 +334,18 @@ class PeerRegistry(BaseServiceManager): self.logger.error(f"Error updating peer {name} IP: {e}") return False + def set_route_via(self, peer_name: str, via_cell: Optional[str]) -> Dict[str, Any]: + """Set or clear the route_via field on a peer. Returns the updated peer dict.""" + with self.lock: + for peer in self.peers: + if peer.get('peer') == peer_name: + peer['route_via'] = via_cell + peer['updated_at'] = datetime.utcnow().isoformat() + self._save_peers() + self.logger.info(f"Set route_via for {peer_name}: {via_cell!r}") + return dict(peer) + raise ValueError(f"Peer '{peer_name}' not found") + def get_peer_stats(self) -> Dict[str, Any]: """Get peer registry statistics""" try: diff --git a/api/routes/cells.py b/api/routes/cells.py index 27637f5..ef804ce 100644 --- a/api/routes/cells.py +++ b/api/routes/cells.py @@ -219,8 +219,11 @@ def peer_sync_permissions(): return jsonify({'ok': False, 'error': f'unknown service: {svc!r}'}), 400 exit_offered = bool(data.get('exit_offered', False)) + use_as_exit_relay = bool(data.get('use_as_exit_relay', False)) from app import cell_link_manager - cell_link_manager.apply_remote_permissions(sender_pubkey, perms, exit_offered=exit_offered) + cell_link_manager.apply_remote_permissions(sender_pubkey, perms, + exit_offered=exit_offered, + use_as_exit_relay=use_as_exit_relay) return jsonify({'ok': True, 'applied_at': datetime.utcnow().isoformat()}) except ValueError as e: return jsonify({'ok': False, 'error': str(e)}), 404 diff --git a/api/routes/peers.py b/api/routes/peers.py index e714a02..253c5b4 100644 --- a/api/routes/peers.py +++ b/api/routes/peers.py @@ -183,6 +183,77 @@ def update_peer(peer_name): return jsonify({"error": str(e)}), 500 +@bp.route('/api/peers//route-via', methods=['PUT']) +def set_peer_route_via(peer_name): + """Route a peer's internet traffic through a connected exit cell. + + Body: {"via_cell": "cellB"} to enable, {"via_cell": null} to disable. + On enable: updates WG AllowedIPs and adds policy routing rule inside + cell-wireguard so the peer's packets egress through the exit cell. + On disable: reverts AllowedIPs and removes the ip rule. + Also signals the exit cell to add/remove the FORWARD-to-eth0 firewall rule. + """ + try: + from app import peer_registry, wireguard_manager, cell_link_manager + data = request.get_json(silent=True) + if data is None or 'via_cell' not in data: + return jsonify({'error': 'via_cell field required (string or null)'}), 400 + + via_cell = data['via_cell'] + if via_cell is not None and not isinstance(via_cell, str): + return jsonify({'error': 'via_cell must be a string or null'}), 400 + + peer = peer_registry.get_peer(peer_name) + if not peer: + return jsonify({'error': 'Peer not found'}), 404 + + peer_ip = peer.get('ip', '').split('/')[0] + if not peer_ip: + return jsonify({'error': 'Peer has no IP assigned'}), 400 + + old_via = peer.get('route_via') + + # Remove old routing if switching away from a previous exit cell + if old_via and old_via != via_cell: + old_link = next( + (l for l in cell_link_manager.list_connections() + if l['cell_name'] == old_via), None + ) + if old_link: + wireguard_manager.update_cell_peer_allowed_ips( + old_link['public_key'], old_link['vpn_subnet'], + add_default_route=False) + wireguard_manager.remove_peer_route_via(peer_ip) + try: + cell_link_manager.set_exit_relay_active(old_via, False) + except Exception as e: + logger.warning(f"set_exit_relay_active(False) for {old_via!r} failed: {e}") + + # Apply new routing + if via_cell: + link = next( + (l for l in cell_link_manager.list_connections() + if l['cell_name'] == via_cell), None + ) + if not link: + return jsonify({'error': f"Cell {via_cell!r} not connected"}), 404 + wireguard_manager.update_cell_peer_allowed_ips( + link['public_key'], link['vpn_subnet'], add_default_route=True) + wireguard_manager.apply_peer_route_via(peer_ip, via_wg_ip=link['dns_ip']) + try: + cell_link_manager.set_exit_relay_active(via_cell, True) + except Exception as e: + logger.warning(f"set_exit_relay_active(True) for {via_cell!r} failed: {e}") + + updated_peer = peer_registry.set_route_via(peer_name, via_cell) + return jsonify({'message': f"Route-via for '{peer_name}' updated", 'peer': updated_peer}) + except ValueError as e: + return jsonify({'error': str(e)}), 404 + except Exception as e: + logger.error(f"Error setting route-via for {peer_name}: {e}") + return jsonify({'error': str(e)}), 500 + + @bp.route('/api/peers//clear-reinstall', methods=['POST']) def clear_peer_reinstall(peer_name): try: diff --git a/api/wireguard_manager.py b/api/wireguard_manager.py index 74704b1..cc374ef 100644 --- a/api/wireguard_manager.py +++ b/api/wireguard_manager.py @@ -573,6 +573,80 @@ class WireGuardManager(BaseServiceManager): except Exception as e: logger.warning(f'sync_cell_routes failed (non-fatal): {e}') + def update_cell_peer_allowed_ips(self, public_key: str, + vpn_subnet: str, + add_default_route: bool = False) -> bool: + """Update AllowedIPs for a cell peer to optionally include 0.0.0.0/0. + + When add_default_route=True, appends 0.0.0.0/0 to AllowedIPs so that + traffic destined for the internet is routed through this cell peer + (Phase 3 exit-relay routing). + + At most one WireGuard peer can have 0.0.0.0/0 in AllowedIPs per + interface (v1 constraint); callers must remove it from any other cell + peer before adding it here. + """ + try: + if add_default_route: + new_ips = f'{vpn_subnet}, 0.0.0.0/0' + else: + new_ips = vpn_subnet + return self.update_peer_ip(public_key, new_ips) + except Exception as e: + logger.error(f'update_cell_peer_allowed_ips failed: {e}') + return False + + def apply_peer_route_via(self, peer_ip: str, via_wg_ip: str, + table: int = 100) -> bool: + """Add policy routing inside cell-wireguard so peer_ip traffic uses via_wg_ip. + + Creates a dedicated routing table (table) with a default route via the + exit cell's WG IP, then adds an ip rule so traffic sourced from peer_ip + uses that table. + + This is the cell-A side of Phase 3 routing: after this, traffic from + the local peer (at peer_ip) is sent through the WG tunnel to cell B + (at via_wg_ip) for internet egress. + + Idempotent: adding a duplicate route/rule returns 0 or EEXIST, which is + treated as success. + """ + real_conf = self._config_file() + if '/tmp/' in real_conf or 'pytest' in real_conf: + return True + try: + def _wg(cmd): + return subprocess.run( + ['docker', 'exec', 'cell-wireguard'] + cmd, + capture_output=True, text=True, timeout=5 + ) + # Add default route in the policy table (idempotent via || true) + _wg(['ip', 'route', 'add', 'default', 'via', via_wg_ip, + 'dev', 'wg0', 'table', str(table)]) + # Add ip rule: traffic FROM peer_ip uses this table + _wg(['ip', 'rule', 'add', 'from', f'{peer_ip}/32', + 'pref', str(table), 'lookup', str(table)]) + logger.info(f'apply_peer_route_via: {peer_ip} → {via_wg_ip} table {table}') + return True + except Exception as e: + logger.error(f'apply_peer_route_via failed: {e}') + return False + + def remove_peer_route_via(self, peer_ip: str, table: int = 100) -> None: + """Remove the ip rule for peer_ip added by apply_peer_route_via. Non-fatal.""" + real_conf = self._config_file() + if '/tmp/' in real_conf or 'pytest' in real_conf: + return + try: + subprocess.run( + ['docker', 'exec', 'cell-wireguard', + 'ip', 'rule', 'del', 'from', f'{peer_ip}/32', + 'pref', str(table), 'lookup', str(table)], + capture_output=True, timeout=5 + ) + except Exception: + pass + def remove_peer(self, public_key: str) -> bool: """Remove the [Peer] block matching public_key from wg0.conf.""" try: diff --git a/tests/test_cell_link_manager.py b/tests/test_cell_link_manager.py index 5e5f5f7..0a028a4 100644 --- a/tests/test_cell_link_manager.py +++ b/tests/test_cell_link_manager.py @@ -641,7 +641,8 @@ class TestPermissionSync(unittest.TestCase): {'inbound': {'calendar': True, 'files': False, 'mail': False, 'webdav': False}, 'outbound': {}}, ) - mock_rules.assert_called_once_with('office', '10.1.0.0/24', ['calendar']) + mock_rules.assert_called_once_with('office', '10.1.0.0/24', ['calendar'], + exit_relay=False) # ── replay_pending_pushes ───────────────────────────────────────────────── @@ -848,5 +849,124 @@ class TestExitOffer(unittest.TestCase): self.assertFalse(links[0]['remote_exit_offered']) +class TestExitRelay(unittest.TestCase): + """Tests for Phase 3: per-peer internet routing via exit cell.""" + + INVITE = { + 'cell_name': 'office', + 'public_key': 'officepubkey=', + 'endpoint': '5.5.5.5:51820', + 'vpn_subnet': '10.1.0.0/24', + 'dns_ip': '10.1.0.1', + 'domain': 'office', + } + + def setUp(self): + self.test_dir = tempfile.mkdtemp() + self.data_dir = os.path.join(self.test_dir, 'data') + self.config_dir = os.path.join(self.test_dir, 'config') + os.makedirs(self.data_dir, exist_ok=True) + os.makedirs(self.config_dir, exist_ok=True) + self.wg = _make_wg_mock() + self.net = _make_nm_mock() + self.mgr = CellLinkManager(self.data_dir, self.config_dir, self.wg, self.net) + + def tearDown(self): + shutil.rmtree(self.test_dir) + + def _add_office(self): + with patch('cell_link_manager.CellLinkManager._local_identity', + return_value={'cell_name': 'home', 'public_key': 'homepubkey='}), \ + patch('cell_link_manager.CellLinkManager._push_permissions_to_remote', + return_value={'ok': True, 'error': None}): + self.mgr.add_connection(self.INVITE) + + def test_phase3_migration_adds_exit_relay_fields(self): + """Existing links without Phase 3 fields get them on load.""" + raw = [{'cell_name': 'office', 'public_key': 'officepubkey=', + 'vpn_subnet': '10.1.0.0/24', 'dns_ip': '10.1.0.1', + 'domain': 'office', 'endpoint': '5.5.5.5:51820', + 'permissions': {'inbound': {}, 'outbound': {}}, + 'remote_api_url': 'http://10.1.0.1:3000', + 'last_push_status': 'ok', 'last_push_at': None, + 'last_push_error': None, 'pending_push': False, + 'last_remote_update_at': None, + 'exit_offered': False, 'remote_exit_offered': False}] + with open(os.path.join(self.data_dir, 'cell_links.json'), 'w') as f: + json.dump(raw, f) + links = self.mgr.list_connections() + self.assertIn('exit_relay_active', links[0]) + self.assertIn('remote_exit_relay_active', links[0]) + self.assertFalse(links[0]['exit_relay_active']) + self.assertFalse(links[0]['remote_exit_relay_active']) + + def test_set_exit_relay_active_persists(self): + self._add_office() + with patch('cell_link_manager.CellLinkManager._try_push'): + link = self.mgr.set_exit_relay_active('office', True) + self.assertTrue(link['exit_relay_active']) + self.assertTrue(self.mgr.list_connections()[0]['exit_relay_active']) + + def test_set_exit_relay_active_false_persists(self): + self._add_office() + with patch('cell_link_manager.CellLinkManager._try_push'): + self.mgr.set_exit_relay_active('office', True) + link = self.mgr.set_exit_relay_active('office', False) + self.assertFalse(link['exit_relay_active']) + + def test_set_exit_relay_active_triggers_push(self): + self._add_office() + push_mock = MagicMock() + with patch('cell_link_manager.CellLinkManager._try_push', push_mock): + self.mgr.set_exit_relay_active('office', True) + push_mock.assert_called_once() + + def test_set_exit_relay_active_unknown_cell_raises(self): + with self.assertRaises(ValueError): + self.mgr.set_exit_relay_active('nobody', True) + + def test_push_includes_use_as_exit_relay_true(self): + self._add_office() + with patch('cell_link_manager.CellLinkManager._try_push'): + self.mgr.set_exit_relay_active('office', True) + captured = [] + + def fake_run(cmd, **kw): + captured.append(cmd) + r = MagicMock() + r.returncode = 0 + r.stdout = '200' + return r + + import json as _json + with patch('cell_link_manager.CellLinkManager._local_identity', + return_value={'cell_name': 'home', 'public_key': 'homepubkey='}), \ + patch('cell_link_manager.subprocess.run', side_effect=fake_run), \ + patch('cell_link_manager.CellLinkManager._local_wg_ip', return_value=None): + self.mgr._try_push('office', self.mgr.list_connections()[0]) + + flat = [arg for cmd in captured for arg in cmd] + payload_str = next(a for a in flat if a.startswith('{')) + body = _json.loads(payload_str) + self.assertIn('use_as_exit_relay', body) + self.assertTrue(body['use_as_exit_relay']) + + def test_apply_remote_permissions_stores_remote_exit_relay_active(self): + self._add_office() + with patch('firewall_manager.apply_cell_rules'): + self.mgr.apply_remote_permissions('officepubkey=', {}, + use_as_exit_relay=True) + link = self.mgr.list_connections()[0] + self.assertTrue(link['remote_exit_relay_active']) + + def test_apply_remote_permissions_calls_apply_cell_rules_with_exit_relay_true(self): + self._add_office() + with patch('firewall_manager.apply_cell_rules') as mock_rules: + self.mgr.apply_remote_permissions('officepubkey=', {}, + use_as_exit_relay=True) + mock_rules.assert_called_once_with('office', '10.1.0.0/24', [], + exit_relay=True) + + if __name__ == '__main__': unittest.main() diff --git a/tests/test_cells_endpoints.py b/tests/test_cells_endpoints.py index 6a539b9..896c3da 100644 --- a/tests/test_cells_endpoints.py +++ b/tests/test_cells_endpoints.py @@ -572,6 +572,7 @@ class TestPeerSyncPermissionsEndpoint(unittest.TestCase): mock_clm.apply_remote_permissions.assert_called_once_with( 'officepubkey=', self._VALID_BODY['permissions'], exit_offered=False, + use_as_exit_relay=False, ) @patch('app.cell_link_manager') diff --git a/tests/test_peer_registry.py b/tests/test_peer_registry.py index a07b7cd..fa9771a 100644 --- a/tests/test_peer_registry.py +++ b/tests/test_peer_registry.py @@ -77,5 +77,35 @@ class TestPeerRegistry(unittest.TestCase): registry = PeerRegistry(data_dir=self.test_dir, config_dir=self.test_dir) self.assertEqual(registry.list_peers(), []) + def test_route_via_migration_adds_field(self): + """Existing peers without route_via get it as None on load.""" + peers_file = os.path.join(self.test_dir, 'peers.json') + raw = [{'peer': 'alice', 'ip': '10.0.0.5', 'public_key': 'key=', + 'active': True, 'created_at': '2026-01-01T00:00:00'}] + with open(peers_file, 'w') as f: + json.dump(raw, f) + reg = PeerRegistry(data_dir=self.test_dir, config_dir=self.test_dir) + peer = reg.get_peer('alice') + self.assertIn('route_via', peer) + self.assertIsNone(peer['route_via']) + + def test_set_route_via_persists(self): + self.registry.add_peer({'peer': 'alice', 'ip': '10.0.0.5'}) + updated = self.registry.set_route_via('alice', 'exit-cell') + self.assertEqual(updated['route_via'], 'exit-cell') + # Verify it survives a reload + reloaded = PeerRegistry(data_dir=self.test_dir, config_dir=self.test_dir) + self.assertEqual(reloaded.get_peer('alice')['route_via'], 'exit-cell') + + def test_set_route_via_clear(self): + self.registry.add_peer({'peer': 'alice', 'ip': '10.0.0.5'}) + self.registry.set_route_via('alice', 'exit-cell') + updated = self.registry.set_route_via('alice', None) + self.assertIsNone(updated['route_via']) + + def test_set_route_via_unknown_peer_raises(self): + with self.assertRaises(ValueError): + self.registry.set_route_via('nobody', 'exit-cell') + if __name__ == '__main__': unittest.main() \ No newline at end of file diff --git a/tests/test_peer_route_via.py b/tests/test_peer_route_via.py new file mode 100644 index 0000000..b8884ce --- /dev/null +++ b/tests/test_peer_route_via.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +""" +Tests for PUT /api/peers//route-via. + +Covers: + - missing via_cell field → 400 + - peer not found → 404 + - connected cell not found → 404 + - enable route-via: calls wg methods and set_exit_relay_active(True) + - disable route-via: clears wg methods and set_exit_relay_active(False) + - switching from one exit cell to another removes old routing first +""" + +import sys +import json +import unittest +from pathlib import Path +from unittest.mock import patch, MagicMock + +api_dir = Path(__file__).parent.parent / 'api' +sys.path.insert(0, str(api_dir)) + +from app import app + +_PEER = { + 'peer': 'alice', + 'ip': '10.0.0.5', + 'public_key': 'alicepubkey=', + 'route_via': None, + 'internet_access': True, + 'service_access': [], +} + +_CELL_LINK = { + 'cell_name': 'exit-cell', + 'public_key': 'exitcellpubkey=', + 'vpn_subnet': '10.1.0.0/24', + 'dns_ip': '10.1.0.1', + 'remote_exit_relay_active': False, +} + + +class TestSetPeerRouteVia(unittest.TestCase): + + def setUp(self): + app.config['TESTING'] = True + self.client = app.test_client() + + def _put(self, peer_name, body): + return self.client.put( + f'/api/peers/{peer_name}/route-via', + data=json.dumps(body), + content_type='application/json', + ) + + @patch('app.cell_link_manager') + @patch('app.wireguard_manager') + @patch('app.peer_registry') + def test_missing_via_cell_returns_400(self, mock_reg, mock_wg, mock_clm): + r = self._put('alice', {}) + self.assertEqual(r.status_code, 400) + data = json.loads(r.data) + self.assertIn('via_cell', data['error']) + + @patch('app.cell_link_manager') + @patch('app.wireguard_manager') + @patch('app.peer_registry') + def test_peer_not_found_returns_404(self, mock_reg, mock_wg, mock_clm): + mock_reg.get_peer.return_value = None + r = self._put('nobody', {'via_cell': 'exit-cell'}) + self.assertEqual(r.status_code, 404) + + @patch('app.cell_link_manager') + @patch('app.wireguard_manager') + @patch('app.peer_registry') + def test_unknown_cell_returns_404(self, mock_reg, mock_wg, mock_clm): + mock_reg.get_peer.return_value = _PEER + mock_clm.list_connections.return_value = [] + r = self._put('alice', {'via_cell': 'no-such-cell'}) + self.assertEqual(r.status_code, 404) + + @patch('app.cell_link_manager') + @patch('app.wireguard_manager') + @patch('app.peer_registry') + def test_enable_route_via_calls_wg_methods(self, mock_reg, mock_wg, mock_clm): + mock_reg.get_peer.return_value = _PEER + mock_reg.set_route_via.return_value = dict(_PEER, route_via='exit-cell') + mock_clm.list_connections.return_value = [_CELL_LINK] + r = self._put('alice', {'via_cell': 'exit-cell'}) + self.assertEqual(r.status_code, 200) + mock_wg.update_cell_peer_allowed_ips.assert_called_once_with( + 'exitcellpubkey=', '10.1.0.0/24', add_default_route=True) + mock_wg.apply_peer_route_via.assert_called_once_with( + '10.0.0.5', via_wg_ip='10.1.0.1') + + @patch('app.cell_link_manager') + @patch('app.wireguard_manager') + @patch('app.peer_registry') + def test_enable_route_via_signals_exit_relay_active(self, mock_reg, mock_wg, mock_clm): + mock_reg.get_peer.return_value = _PEER + mock_reg.set_route_via.return_value = dict(_PEER, route_via='exit-cell') + mock_clm.list_connections.return_value = [_CELL_LINK] + self._put('alice', {'via_cell': 'exit-cell'}) + mock_clm.set_exit_relay_active.assert_called_once_with('exit-cell', True) + + @patch('app.cell_link_manager') + @patch('app.wireguard_manager') + @patch('app.peer_registry') + def test_disable_route_via_clears_old_routing(self, mock_reg, mock_wg, mock_clm): + peer_with_via = dict(_PEER, route_via='exit-cell') + mock_reg.get_peer.return_value = peer_with_via + mock_reg.set_route_via.return_value = dict(_PEER, route_via=None) + mock_clm.list_connections.return_value = [_CELL_LINK] + r = self._put('alice', {'via_cell': None}) + self.assertEqual(r.status_code, 200) + mock_wg.update_cell_peer_allowed_ips.assert_called_once_with( + 'exitcellpubkey=', '10.1.0.0/24', add_default_route=False) + mock_wg.remove_peer_route_via.assert_called_once_with('10.0.0.5') + mock_clm.set_exit_relay_active.assert_called_once_with('exit-cell', False) + + @patch('app.cell_link_manager') + @patch('app.wireguard_manager') + @patch('app.peer_registry') + def test_enable_route_via_persists_field(self, mock_reg, mock_wg, mock_clm): + mock_reg.get_peer.return_value = _PEER + updated = dict(_PEER, route_via='exit-cell') + mock_reg.set_route_via.return_value = updated + mock_clm.list_connections.return_value = [_CELL_LINK] + r = self._put('alice', {'via_cell': 'exit-cell'}) + self.assertEqual(r.status_code, 200) + mock_reg.set_route_via.assert_called_once_with('alice', 'exit-cell') + data = json.loads(r.data) + self.assertEqual(data['peer']['route_via'], 'exit-cell') + + @patch('app.cell_link_manager') + @patch('app.wireguard_manager') + @patch('app.peer_registry') + def test_switch_exit_cell_removes_old_before_adding_new(self, mock_reg, mock_wg, mock_clm): + old_link = dict(_CELL_LINK, cell_name='old-exit', public_key='oldpubkey=', + vpn_subnet='10.2.0.0/24', dns_ip='10.2.0.1') + new_link = dict(_CELL_LINK, cell_name='new-exit', public_key='newpubkey=') + peer_with_old = dict(_PEER, route_via='old-exit') + mock_reg.get_peer.return_value = peer_with_old + mock_reg.set_route_via.return_value = dict(_PEER, route_via='new-exit') + mock_clm.list_connections.return_value = [old_link, new_link] + self._put('alice', {'via_cell': 'new-exit'}) + calls = mock_wg.update_cell_peer_allowed_ips.call_args_list + # First call removes old (add_default_route=False), second adds new (True) + self.assertEqual(len(calls), 2) + self.assertFalse(calls[0].kwargs['add_default_route']) + self.assertTrue(calls[1].kwargs['add_default_route']) + # set_exit_relay_active called twice: False for old, True for new + relay_calls = mock_clm.set_exit_relay_active.call_args_list + self.assertEqual(len(relay_calls), 2) + self.assertEqual(relay_calls[0].args, ('old-exit', False)) + self.assertEqual(relay_calls[1].args, ('new-exit', True)) + + +if __name__ == '__main__': + unittest.main()