roof/pic
1
0
Fork 0
Code Issues Pull Requests Actions Wiki Activity
Files
aba2b0d33f234c8cec5fa7296446b073004d97d5
pic/tests/test_egress_manager.py
T
roof 89aed4efe0
Unit Tests / test (push) Successful in 12m6s
Details
feat: connectivity redesign phase 2 — instance-aware routing + reference connections by id 
apply_routes now iterates over connection instances rather than types:
each instance gets its own fwmark, routing table, interface, and
redirect_port via _routing_connections / _resolve_peer_connection /
_apply_connection_for_src; kill-switch is enforced per iface-instance.
Old per-type MARKS/TABLES constants are kept only as migration scaffolding.

peer_registry: exit_via is now stored as a connection id (or 'default');
_migrate_exit_via_to_connection_id runs on _load_peers to upgrade legacy
type-string values; set_peer_exit_via validates against known connection
ids; VALID_EXIT_VIA removed; config_manager wired in from managers.py.

egress_manager: egress_overrides keyed by service_id → connection_id;
local MARKS/TABLES/EXIT_TYPES/_REDIRECT_PORTS/_add_tor_redirect removed;
(mark, table, redirect_port) resolved at apply-time via
connectivity_manager.get_connection; manifest egress.allowed still
enforced by connection type.

api/app.py + api.js: PUT peer/service exit endpoints accept {connection_id};
back-compat shim resolves a legacy type string to its single active instance.

Tests extended: two same-type instances produce distinct marks/tables/ports;
peer exit_via and egress override id migrations round-trip correctly;
single-instance behaviour is equivalent to the old type-keyed path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-10 17:35:28 -04:00

659 lines
25 KiB
Python
Raw Blame History

"""
Tests for EgressManager — per-service egress enforcement via host iptables.
Connectivity v2: a service routes through a connection *instance* (by id),
sharing the connection's fwmark / routing table / redirect port. The egress
override map is service_id → connection_id, and (mark, table, redirect_port)
are resolved from ConnectivityManager.get_connection(id). EgressManager no
longer owns its own per-type MARKS/TABLES.
All subprocess calls (iptables, iptables-save, iptables-restore, ip rule,
docker inspect) and manager state are mocked so these tests run without any
live infrastructure or root privileges.
"""
import os
import sys
import unittest
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'api'))
import egress_manager as em_module
from egress_manager import EgressManager, EGRESS_CHAIN
# ---------------------------------------------------------------------------
# Connection fixtures — mirror the v2 allocator output.
# ---------------------------------------------------------------------------
CONNECTIONS = {
'conn_wg': {
'id': 'conn_wg', 'type': 'wireguard_ext', 'name': 'Work VPN',
'enabled': True, 'mark': 0x1000, 'table': 1000,
'iface': 'wgext_aaaa', 'redirect_port': None,
},
'conn_ovpn': {
'id': 'conn_ovpn', 'type': 'openvpn', 'name': 'OVPN',
'enabled': True, 'mark': 0x1010, 'table': 1001,
'iface': 'ovpn_bbbb', 'redirect_port': None,
},
'conn_tor': {
'id': 'conn_tor', 'type': 'tor', 'name': 'Tor',
'enabled': True, 'mark': 0x1020, 'table': 1002,
'iface': None, 'redirect_port': 9100,
},
}
def _make_manager(installed=None, overrides=None, connections=None):
"""Build an EgressManager backed by mock config + connectivity managers."""
cm = MagicMock()
cm.get_installed_services.return_value = installed or {}
cm.configs = {'egress_overrides': overrides or {}}
cm._save_all_configs = MagicMock()
conns = connections if connections is not None else CONNECTIONS
conn_list = list(conns.values())
cm.list_connections.return_value = conn_list
cm.get_connection.side_effect = lambda cid: conns.get(cid)
connectivity = MagicMock()
connectivity.list_connections.return_value = conn_list
connectivity.get_connection.side_effect = lambda cid: conns.get(cid)
mgr = EgressManager(config_manager=cm, connectivity_manager=connectivity)
return mgr, cm
def _subprocess_ok(stdout=''):
return MagicMock(returncode=0, stdout=stdout, stderr='')
def _subprocess_fail(stderr='error', stdout=''):
return MagicMock(returncode=1, stdout=stdout, stderr=stderr)
def _make_manifest(has_egress=True, egress_default='conn_wg',
allowed=None, container_name='cell-myapp'):
"""Return a minimal manifest dict with optional egress configuration.
`allowed` is a list of connection *types* (manifests are type-scoped).
"""
m = {
'id': 'myapp',
'name': 'My App',
'container_name': container_name,
}
if has_egress:
m['has_egress'] = True
m['egress'] = {
'default': egress_default,
'allowed': allowed if allowed is not None
else ['wireguard_ext', 'openvpn', 'tor', 'sshuttle', 'proxy'],
}
else:
m['has_egress'] = False
return m
def _installed_with_manifest(manifest, service_id='myapp'):
return {service_id: {'id': service_id, 'manifest': manifest}}
# ---------------------------------------------------------------------------
# 1. default exit → no iptables rule-adding
# ---------------------------------------------------------------------------
class TestApplyServiceDefaultExit(unittest.TestCase):
def test_apply_service_default_exit_no_iptables_calls(self):
manifest = _make_manifest(egress_default='default')
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
with patch('subprocess.run') as mock_run:
mock_run.return_value = _subprocess_ok(stdout='172.20.0.50\n')
result = mgr.apply_service('myapp')
self.assertTrue(result['ok'])
self.assertEqual(result.get('exit_via'), 'default')
rule_add_calls = [
c for c in mock_run.call_args_list
if c.args and c.args[0][:1] == ['iptables']
and any(a in c.args[0] for a in ('-A', '-I', 'MARK', 'REDIRECT'))
]
self.assertEqual(rule_add_calls, [])
# ---------------------------------------------------------------------------
# 2. wireguard_ext connection → mark rule with the connection's own mark
# ---------------------------------------------------------------------------
class TestApplyServiceWireguardExt(unittest.TestCase):
def test_apply_service_wireguard_ext_adds_mark_rule(self):
manifest = _make_manifest(egress_default='conn_wg')
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
calls_made = []
def fake_run(cmd, **kwargs):
calls_made.append(cmd)
if 'docker' in cmd and 'inspect' in cmd:
return _subprocess_ok(stdout='172.20.0.50\n')
if 'iptables-save' in cmd:
return _subprocess_ok(stdout='')
if 'iptables-restore' in cmd:
return _subprocess_ok()
if cmd[:3] == ['ip', 'rule', 'del']:
return _subprocess_fail()
return _subprocess_ok()
with patch('subprocess.run', side_effect=fake_run):
result = mgr.apply_service('myapp')
self.assertTrue(result['ok'], result)
self.assertEqual(result['exit_via'], 'conn_wg')
mark_calls = [
c for c in calls_made
if 'iptables' in str(c) and 'MARK' in c and '--set-mark' in c
]
self.assertGreater(len(mark_calls), 0, 'No MARK rule was added')
mark_cmd = ' '.join(mark_calls[0])
self.assertIn('0x1000', mark_cmd) # the connection's mark
self.assertIn('pic-egr-myapp', mark_cmd)
self.assertIn('mangle', mark_cmd)
# The ip rule must point fwmark→the connection's table.
ip_rule_add = [
c for c in calls_made
if c[:3] == ['ip', 'rule', 'add']
]
self.assertGreater(len(ip_rule_add), 0)
self.assertIn('1000', ' '.join(ip_rule_add[0])) # table
# ---------------------------------------------------------------------------
# 3. openvpn connection → its own mark
# ---------------------------------------------------------------------------
class TestApplyServiceOpenVPN(unittest.TestCase):
def test_apply_service_openvpn_adds_mark_rule(self):
manifest = _make_manifest(egress_default='conn_ovpn')
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
calls_made = []
def fake_run(cmd, **kwargs):
calls_made.append(cmd)
if 'docker' in cmd and 'inspect' in cmd:
return _subprocess_ok(stdout='172.20.0.51\n')
if 'iptables-save' in cmd:
return _subprocess_ok(stdout='')
if 'iptables-restore' in cmd:
return _subprocess_ok()
if cmd[:3] == ['ip', 'rule', 'del']:
return _subprocess_fail()
return _subprocess_ok()
with patch('subprocess.run', side_effect=fake_run):
result = mgr.apply_service('myapp')
self.assertTrue(result['ok'], result)
self.assertEqual(result['exit_via'], 'conn_ovpn')
mark_calls = [
c for c in calls_made
if 'iptables' in str(c) and 'MARK' in c and '--set-mark' in c
]
self.assertGreater(len(mark_calls), 0)
self.assertIn('0x1010', ' '.join(mark_calls[0]))
# ---------------------------------------------------------------------------
# 4. tor (redirect-style) connection → mark + REDIRECT to its redirect_port
# ---------------------------------------------------------------------------
class TestApplyServiceTor(unittest.TestCase):
def test_apply_service_tor_adds_mark_and_redirect(self):
manifest = _make_manifest(egress_default='conn_tor')
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
calls_made = []
def fake_run(cmd, **kwargs):
calls_made.append(cmd)
if 'docker' in cmd and 'inspect' in cmd:
return _subprocess_ok(stdout='172.20.0.52\n')
if 'iptables-save' in cmd:
return _subprocess_ok(stdout='')
if 'iptables-restore' in cmd:
return _subprocess_ok()
if cmd[:3] == ['ip', 'rule', 'del']:
return _subprocess_fail()
return _subprocess_ok()
with patch('subprocess.run', side_effect=fake_run):
result = mgr.apply_service('myapp')
self.assertTrue(result['ok'], result)
self.assertEqual(result['exit_via'], 'conn_tor')
mark_calls = [
c for c in calls_made
if 'iptables' in str(c) and 'MARK' in c and '--set-mark' in c
]
self.assertGreater(len(mark_calls), 0, 'No MARK rule found')
self.assertIn('0x1020', ' '.join(mark_calls[0]))
redirect_calls = [
c for c in calls_made
if 'iptables' in str(c) and 'REDIRECT' in c
]
self.assertGreater(len(redirect_calls), 0, 'No REDIRECT rule found')
redirect_cmd = ' '.join(redirect_calls[0])
self.assertIn('9100', redirect_cmd) # the connection's redirect_port
self.assertIn('nat', redirect_cmd)
# ---------------------------------------------------------------------------
# 5. no container IP → error
# ---------------------------------------------------------------------------
class TestApplyServiceNoContainerIP(unittest.TestCase):
def test_apply_service_no_container_ip_returns_error(self):
manifest = _make_manifest(egress_default='conn_wg')
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
def fake_run(cmd, **kwargs):
if 'docker' in cmd and 'inspect' in cmd:
return _subprocess_ok(stdout='\n')
if 'iptables-save' in cmd:
return _subprocess_ok(stdout='')
if 'iptables-restore' in cmd:
return _subprocess_ok()
return _subprocess_ok()
with patch('subprocess.run', side_effect=fake_run):
result = mgr.apply_service('myapp')
self.assertFalse(result['ok'])
self.assertIn('container IP not discoverable', result.get('error', ''))
# ---------------------------------------------------------------------------
# 6. container IP retries
# ---------------------------------------------------------------------------
class TestApplyServiceRetries(unittest.TestCase):
def test_apply_service_container_ip_retries(self):
manifest = _make_manifest(egress_default='conn_wg')
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
inspect_count = [0]
def fake_run(cmd, **kwargs):
if 'docker' in cmd and 'inspect' in cmd:
inspect_count[0] += 1
if inspect_count[0] == 1:
return _subprocess_ok(stdout='\n')
return _subprocess_ok(stdout='172.20.0.50\n')
if 'iptables-save' in cmd:
return _subprocess_ok(stdout='')
if 'iptables-restore' in cmd:
return _subprocess_ok()
if cmd[:3] == ['ip', 'rule', 'del']:
return _subprocess_fail()
return _subprocess_ok()
with patch('subprocess.run', side_effect=fake_run):
with patch('time.sleep'):
result = mgr.apply_service('myapp')
self.assertTrue(result['ok'], result)
self.assertGreaterEqual(inspect_count[0], 2)
# ---------------------------------------------------------------------------
# 7. has_egress False → skipped
# ---------------------------------------------------------------------------
class TestHasEgressFalse(unittest.TestCase):
def test_has_egress_false_skips_rules(self):
manifest = _make_manifest(has_egress=False)
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
with patch('subprocess.run') as mock_run:
mock_run.return_value = _subprocess_ok(stdout='')
result = mgr.apply_service('myapp')
self.assertTrue(result['ok'])
self.assertTrue(result.get('skipped'))
rule_add_calls = [
c for c in mock_run.call_args_list
if c.args and c.args[0][:1] == ['iptables']
and any(a in c.args[0] for a in ('-A', '-I', 'MARK', 'REDIRECT'))
]
self.assertEqual(rule_add_calls, [])
# ---------------------------------------------------------------------------
# 8. has_egress True but no egress block → skipped
# ---------------------------------------------------------------------------
class TestHasEgressMissingBlock(unittest.TestCase):
def test_has_egress_missing_egress_block_skips(self):
manifest = {
'id': 'myapp',
'container_name': 'cell-myapp',
'has_egress': True,
}
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
with patch('subprocess.run') as mock_run:
mock_run.return_value = _subprocess_ok(stdout='')
result = mgr.apply_service('myapp')
self.assertTrue(result['ok'])
self.assertTrue(result.get('skipped'))
# ---------------------------------------------------------------------------
# 9. clear_service removes only the tagged rules
# ---------------------------------------------------------------------------
class TestClearService(unittest.TestCase):
def test_clear_service_removes_tagged_rules(self):
mgr, _ = _make_manager()
mangle_rules = (
'-A PIC_EGRESS -s 172.20.0.50 -j MARK --set-mark 0x1000 '
'-m comment --comment "pic-egr-myapp"\n'
'-A PIC_EGRESS -s 172.20.0.99 -j MARK --set-mark 0x1000 '
'-m comment --comment "pic-egr-otherapp"\n'
)
nat_rules = ''
restore_inputs = {}
def fake_run(cmd, input=None, **kwargs):
if cmd == ['iptables-save', '-t', 'mangle']:
return _subprocess_ok(stdout=mangle_rules)
if cmd == ['iptables-save', '-t', 'nat']:
return _subprocess_ok(stdout=nat_rules)
if cmd == ['iptables-restore', '-T', 'mangle']:
restore_inputs['mangle'] = input
return _subprocess_ok()
if cmd == ['iptables-restore', '-T', 'nat']:
restore_inputs['nat'] = input
return _subprocess_ok()
return _subprocess_ok()
with patch('subprocess.run', side_effect=fake_run):
result = mgr.clear_service('myapp')
self.assertTrue(result['ok'])
restored = restore_inputs.get('mangle', '')
self.assertNotIn('pic-egr-myapp', restored)
self.assertIn('pic-egr-otherapp', restored)
# ---------------------------------------------------------------------------
# 10. set_service_exit rejects a connection whose type is not in allowed
# ---------------------------------------------------------------------------
class TestSetServiceExitRejectNotAllowed(unittest.TestCase):
def test_set_service_exit_rejects_type_not_in_allowed(self):
manifest = _make_manifest(
egress_default='default',
allowed=['tor'], # wireguard_ext not allowed
)
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
result = mgr.set_service_exit('myapp', 'conn_wg')
self.assertFalse(result['ok'])
self.assertIn('error', result)
self.assertIn('allowed', result['error'])
# ---------------------------------------------------------------------------
# 11. set_service_exit persists the connection id and applies
# ---------------------------------------------------------------------------
class TestSetServiceExitPersistsAndApplies(unittest.TestCase):
def test_set_service_exit_persists_and_applies(self):
manifest = _make_manifest(egress_default='default')
mgr, cm = _make_manager(installed=_installed_with_manifest(manifest))
apply_calls = []
mgr.apply_service = lambda sid: apply_calls.append(sid) or {
'ok': True, 'exit_via': 'conn_tor'}
result = mgr.set_service_exit('myapp', 'conn_tor')
self.assertTrue(result['ok'], result)
self.assertIn('myapp', apply_calls)
cm._save_all_configs.assert_called()
self.assertEqual(cm.configs['egress_overrides'].get('myapp'), 'conn_tor')
def test_set_service_exit_default_clears_override(self):
manifest = _make_manifest(egress_default='conn_wg')
mgr, cm = _make_manager(
installed=_installed_with_manifest(manifest),
overrides={'myapp': 'conn_tor'},
)
mgr.apply_service = lambda sid: {'ok': True, 'exit_via': 'default'}
result = mgr.set_service_exit('myapp', 'default')
self.assertTrue(result['ok'])
self.assertEqual(cm.configs['egress_overrides'].get('myapp'), 'default')
def test_set_service_exit_unknown_connection_rejected(self):
manifest = _make_manifest(egress_default='default')
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
result = mgr.set_service_exit('myapp', 'conn_ghost')
self.assertFalse(result['ok'])
self.assertIn('unknown connection', result['error'])
def test_set_service_exit_legacy_type_resolves_to_single_instance(self):
"""Back-compat shim: a legacy type resolves to the one instance of it."""
manifest = _make_manifest(egress_default='default', allowed=['tor'])
mgr, cm = _make_manager(installed=_installed_with_manifest(manifest))
mgr.apply_service = lambda sid: {'ok': True}
result = mgr.set_service_exit('myapp', 'tor')
self.assertTrue(result['ok'], result)
self.assertEqual(cm.configs['egress_overrides'].get('myapp'), 'conn_tor')
# ---------------------------------------------------------------------------
# 12. apply_all iterates installed services
# ---------------------------------------------------------------------------
class TestApplyAll(unittest.TestCase):
def test_apply_all_iterates_installed_services(self):
manifests = {
'svc1': _make_manifest(egress_default='conn_wg'),
'svc2': _make_manifest(egress_default='conn_ovpn'),
'svc3': _make_manifest(egress_default='conn_tor'),
}
installed = {
sid: {'id': sid, 'manifest': m}
for sid, m in manifests.items()
}
mgr, _ = _make_manager(installed=installed)
applied = []
mgr.apply_service = lambda sid: applied.append(sid) or {'ok': True}
result = mgr.apply_all()
self.assertTrue(result['ok'])
self.assertEqual(sorted(applied), ['svc1', 'svc2', 'svc3'])
# ---------------------------------------------------------------------------
# 13. service + peer on the same connection share the same mark
# ---------------------------------------------------------------------------
class TestSharedMarkWithConnection(unittest.TestCase):
def test_service_uses_connection_mark_not_a_private_table(self):
"""A service routed via a connection uses that connection's mark/table,
i.e. the SAME resources a peer on that connection would use."""
manifest = _make_manifest(egress_default='conn_wg')
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
captured = {}
def fake_run(cmd, **kwargs):
if 'docker' in cmd and 'inspect' in cmd:
return _subprocess_ok(stdout='172.20.0.50\n')
if 'iptables-save' in cmd:
return _subprocess_ok(stdout='')
if 'iptables-restore' in cmd:
return _subprocess_ok()
if cmd[:3] == ['ip', 'rule', 'del']:
return _subprocess_fail()
if cmd[:3] == ['ip', 'rule', 'add']:
captured['ip_rule'] = ' '.join(cmd)
if 'MARK' in cmd:
captured['mark'] = ' '.join(cmd)
return _subprocess_ok()
with patch('subprocess.run', side_effect=fake_run):
mgr.apply_service('myapp')
# conn_wg has mark 0x1000 / table 1000 — both must appear, proving the
# service inherits the connection's shared routing resources.
self.assertIn('0x1000', captured.get('mark', ''))
self.assertIn('1000', captured.get('ip_rule', ''))
# ---------------------------------------------------------------------------
# 14. apply_service with an unknown connection id → error
# ---------------------------------------------------------------------------
class TestApplyServiceUnknownExit(unittest.TestCase):
def test_apply_service_unknown_connection_rejected(self):
manifest = {
'id': 'myapp',
'container_name': 'cell-myapp',
'has_egress': True,
'egress': {
'default': 'conn_ghost', # no such connection
'allowed': ['wireguard_ext'],
},
}
mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))
def fake_run(cmd, **kwargs):
if 'docker' in cmd and 'inspect' in cmd:
return _subprocess_ok(stdout='172.20.0.50\n')
if 'iptables-save' in cmd:
return _subprocess_ok(stdout='')
if 'iptables-restore' in cmd:
return _subprocess_ok()
return _subprocess_ok()
with patch('subprocess.run', side_effect=fake_run):
result = mgr.apply_service('myapp')
self.assertFalse(result['ok'])
self.assertIn('error', result)
# ---------------------------------------------------------------------------
# _has_egress edge cases
# ---------------------------------------------------------------------------
class TestHasEgressLogic(unittest.TestCase):
def setUp(self):
self.mgr, _ = _make_manager()
def test_has_egress_both_required(self):
m = {'has_egress': True, 'egress': {'default': 'conn_tor', 'allowed': ['tor']}}
self.assertTrue(self.mgr._has_egress(m))
def test_has_egress_false_field(self):
m = {'has_egress': False, 'egress': {'default': 'conn_tor', 'allowed': ['tor']}}
self.assertFalse(self.mgr._has_egress(m))
def test_has_egress_missing_has_egress_key(self):
m = {'egress': {'default': 'conn_tor', 'allowed': ['tor']}}
self.assertFalse(self.mgr._has_egress(m))
def test_has_egress_empty_egress_dict(self):
m = {'has_egress': True, 'egress': {}}
self.assertFalse(self.mgr._has_egress(m))
# ---------------------------------------------------------------------------
# _resolve_exit (now returns connection ids)
# ---------------------------------------------------------------------------
class TestResolveExit(unittest.TestCase):
def test_override_takes_precedence(self):
mgr, _ = _make_manager(overrides={'myapp': 'conn_ovpn'})
manifest = _make_manifest(egress_default='conn_wg')
self.assertEqual(mgr._resolve_exit('myapp', manifest), 'conn_ovpn')
def test_manifest_default_used_when_no_override(self):
mgr, _ = _make_manager(overrides={})
manifest = _make_manifest(egress_default='conn_tor')
self.assertEqual(mgr._resolve_exit('myapp', manifest), 'conn_tor')
def test_fallback_to_default_when_no_egress_block(self):
mgr, _ = _make_manager(overrides={})
manifest = {'id': 'myapp'}
self.assertEqual(mgr._resolve_exit('myapp', manifest), 'default')
def test_legacy_type_override_migrates_to_connection_id(self):
"""An old override holding a type string resolves to the migrated id."""
mgr, _ = _make_manager(overrides={'myapp': 'wireguard_ext'})
manifest = _make_manifest(egress_default='default')
self.assertEqual(mgr._resolve_exit('myapp', manifest), 'conn_wg')
def test_legacy_type_default_with_no_instance_falls_back(self):
"""A legacy type with no matching instance falls back to 'default'."""
mgr, _ = _make_manager(connections={})
manifest = _make_manifest(egress_default='tor')
self.assertEqual(mgr._resolve_exit('myapp', manifest), 'default')
# ---------------------------------------------------------------------------
# apply_service with missing manifest
# ---------------------------------------------------------------------------
class TestApplyServiceMissingManifest(unittest.TestCase):
def test_apply_service_missing_manifest_returns_error(self):
mgr, _ = _make_manager(installed={})
result = mgr.apply_service('ghost')
self.assertFalse(result['ok'])
self.assertIn('error', result)
if __name__ == '__main__':
unittest.main()
Reference in New Issue View Git Blame Copy Permalink