mirror of
https://github.com/pyMC-dev/pyMC_Repeater.git
synced 2026-07-31 05:53:30 +02:00
feat(companion): enhance contact capacity management and bridge settings
- Introduced `CompanionContactCapacityError` to handle cases where persisted contacts exceed configured limits. - Added utility functions for parsing companion bridge settings and validating contact capacity. - Updated `RepeaterDaemon` to check contact capacity during companion loading and initialization. - Enhanced API endpoints to validate companion settings and manage contact limits effectively. - Implemented logging for bridge limits and errors related to contact capacity.
This commit is contained in:
Adam Gessaman
@@ -1,7 +1,46 @@
|
||||
"""Shared utilities for Companion (e.g. validation for config sync)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from pymc_core.companion.constants import DEFAULT_MAX_CONTACTS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_INVALID_NODE_NAME_CHARS = "\n\r\x00"
|
||||
|
||||
# Optional per-companion RepeaterCompanionBridge constructor settings (power-user).
|
||||
COMPANION_BRIDGE_SETTING_KEYS = frozenset({"max_contacts", "offline_queue_size"})
|
||||
|
||||
# Settings that must not be applied from config (fixed at pymc_core defaults).
|
||||
_COMPANION_IGNORED_BRIDGE_KEYS = frozenset({"max_channels", "adv_type"})
|
||||
|
||||
# MeshCore device-info reports max_contacts // 2 capped at 255.
|
||||
_DEVICE_INFO_MAX_CONTACTS_CEILING = 510
|
||||
|
||||
|
||||
class CompanionContactCapacityError(Exception):
|
||||
"""Persisted companion contacts exceed configured max_contacts."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
companion_hash: str,
|
||||
stored_count: int,
|
||||
max_contacts: int,
|
||||
companion_name: Optional[str] = None,
|
||||
) -> None:
|
||||
self.companion_hash = companion_hash
|
||||
self.stored_count = stored_count
|
||||
self.max_contacts = max_contacts
|
||||
self.companion_name = companion_name
|
||||
label = f"'{companion_name}'" if companion_name else companion_hash
|
||||
super().__init__(
|
||||
f"Companion {label}: {stored_count} contacts in storage exceeds "
|
||||
f"max_contacts={max_contacts}. Increase max_contacts or remove contacts before starting."
|
||||
)
|
||||
|
||||
|
||||
def normalize_companion_identity_key(identity_key: str) -> str:
|
||||
"""Strip whitespace and remove optional 0x prefix so fromhex() is consistent across installs."""
|
||||
@@ -23,3 +62,144 @@ def validate_companion_node_name(value: str) -> str:
|
||||
if any(c in s for c in _INVALID_NODE_NAME_CHARS):
|
||||
raise ValueError("node_name contains invalid characters")
|
||||
return s
|
||||
|
||||
|
||||
def parse_positive_int(value: Any, field_name: str, *, minimum: int = 1) -> int:
|
||||
"""Parse a positive integer from config or API input."""
|
||||
try:
|
||||
n = int(value)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError(f"{field_name} must be a positive integer") from e
|
||||
if n < minimum:
|
||||
raise ValueError(f"{field_name} must be >= {minimum}")
|
||||
return n
|
||||
|
||||
|
||||
def parse_companion_bridge_kwargs(settings: dict) -> Dict[str, int]:
|
||||
"""Extract optional RepeaterCompanionBridge kwargs from companion settings.
|
||||
|
||||
Only ``max_contacts`` and ``offline_queue_size`` are honored. ``max_channels`` and
|
||||
``adv_type`` are ignored with a warning if present.
|
||||
"""
|
||||
if not settings:
|
||||
return {}
|
||||
for key in _COMPANION_IGNORED_BRIDGE_KEYS:
|
||||
if key in settings:
|
||||
logger.warning(
|
||||
"Companion setting %r is not supported and will be ignored (fixed default)",
|
||||
key,
|
||||
)
|
||||
kwargs: Dict[str, int] = {}
|
||||
if "max_contacts" in settings:
|
||||
max_contacts = parse_positive_int(settings["max_contacts"], "max_contacts")
|
||||
if max_contacts > _DEVICE_INFO_MAX_CONTACTS_CEILING:
|
||||
logger.warning(
|
||||
"max_contacts=%s exceeds %s; MeshCore device-info will still report at most %s",
|
||||
max_contacts,
|
||||
_DEVICE_INFO_MAX_CONTACTS_CEILING,
|
||||
_DEVICE_INFO_MAX_CONTACTS_CEILING,
|
||||
)
|
||||
kwargs["max_contacts"] = max_contacts
|
||||
if "offline_queue_size" in settings:
|
||||
kwargs["offline_queue_size"] = parse_positive_int(
|
||||
settings["offline_queue_size"], "offline_queue_size"
|
||||
)
|
||||
return kwargs
|
||||
|
||||
|
||||
def effective_max_contacts(bridge_kwargs: Dict[str, int]) -> int:
|
||||
"""Return max_contacts from parsed kwargs or pymc_core default."""
|
||||
return bridge_kwargs.get("max_contacts", DEFAULT_MAX_CONTACTS)
|
||||
|
||||
|
||||
def merge_companion_settings_update(current_settings: dict, patch: dict) -> Dict[str, Any]:
|
||||
"""Merge a companion settings PATCH into current settings.
|
||||
|
||||
Raises:
|
||||
ValueError: Unknown setting or invalid bridge setting value.
|
||||
"""
|
||||
merged = dict(current_settings or {})
|
||||
for key, value in patch.items():
|
||||
if key not in COMPANION_SETTINGS_ALLOWLIST:
|
||||
raise ValueError(f"Unknown companion setting: {key}")
|
||||
if key in COMPANION_BRIDGE_SETTING_KEYS:
|
||||
parsed = parse_companion_bridge_kwargs({key: value})
|
||||
merged[key] = parsed[key]
|
||||
else:
|
||||
merged[key] = value
|
||||
return merged
|
||||
|
||||
|
||||
def validate_companion_config_capacity(
|
||||
identity: dict,
|
||||
sqlite_handler: Any,
|
||||
*,
|
||||
companion_name: Optional[str] = None,
|
||||
settings: Optional[dict] = None,
|
||||
) -> None:
|
||||
"""Raise CompanionContactCapacityError if persisted contacts exceed configured max_contacts."""
|
||||
if sqlite_handler is None:
|
||||
return
|
||||
identity_key = identity.get("identity_key")
|
||||
if not identity_key:
|
||||
return
|
||||
merged_settings = settings if settings is not None else (identity.get("settings") or {})
|
||||
max_contacts = effective_max_contacts(parse_companion_bridge_kwargs(merged_settings))
|
||||
companion_hash = companion_hash_str_from_identity_key(identity_key)
|
||||
check_companion_contact_capacity(
|
||||
companion_hash,
|
||||
max_contacts,
|
||||
sqlite_handler,
|
||||
companion_name=companion_name,
|
||||
)
|
||||
|
||||
|
||||
def check_companion_contact_capacity(
|
||||
companion_hash: str,
|
||||
max_contacts: int,
|
||||
sqlite_handler: Any,
|
||||
*,
|
||||
companion_name: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Raise CompanionContactCapacityError if persisted contacts exceed max_contacts."""
|
||||
if sqlite_handler is None:
|
||||
return
|
||||
stored_count = sqlite_handler.companion_count_contacts(companion_hash)
|
||||
if stored_count > max_contacts:
|
||||
raise CompanionContactCapacityError(
|
||||
companion_hash, stored_count, max_contacts, companion_name=companion_name
|
||||
)
|
||||
|
||||
|
||||
def format_companion_bridge_limits(bridge_kwargs: Dict[str, int]) -> str:
|
||||
"""Format non-default bridge limits for log lines."""
|
||||
if not bridge_kwargs:
|
||||
return ""
|
||||
parts = [f"{k}={v}" for k, v in sorted(bridge_kwargs.items())]
|
||||
return ", " + ", ".join(parts)
|
||||
|
||||
|
||||
def companion_hash_str_from_identity_key(identity_key: Any) -> str:
|
||||
"""Derive companion_hash storage key (0xHH) from an identity_key config value."""
|
||||
from pymc_core import LocalIdentity
|
||||
|
||||
if isinstance(identity_key, str):
|
||||
key_bytes = bytes.fromhex(normalize_companion_identity_key(identity_key))
|
||||
elif isinstance(identity_key, bytes):
|
||||
key_bytes = identity_key
|
||||
else:
|
||||
raise ValueError("identity_key has unknown type")
|
||||
pubkey_byte = LocalIdentity(seed=key_bytes).get_public_key()[0]
|
||||
return f"0x{pubkey_byte:02x}"
|
||||
|
||||
|
||||
# All companion settings writable via identity API (tcp + bridge power-user keys).
|
||||
COMPANION_SETTINGS_ALLOWLIST = frozenset(
|
||||
{
|
||||
"node_name",
|
||||
"tcp_port",
|
||||
"bind_address",
|
||||
"tcp_timeout",
|
||||
*COMPANION_BRIDGE_SETTING_KEYS,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -2229,6 +2229,20 @@ class SQLiteHandler:
|
||||
return 0
|
||||
|
||||
# Companion persistence methods
|
||||
def companion_count_contacts(self, companion_hash: str) -> int:
|
||||
"""Return the number of persisted contacts for a companion."""
|
||||
try:
|
||||
with self._connect() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT COUNT(*) FROM companion_contacts WHERE companion_hash = ?",
|
||||
(companion_hash,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return int(row[0]) if row else 0
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to count companion contacts: {e}")
|
||||
return 0
|
||||
|
||||
def companion_load_contacts(self, companion_hash: str) -> List[Dict]:
|
||||
"""Load contacts for a companion from storage."""
|
||||
try:
|
||||
|
||||
+39
-3
@@ -7,7 +7,15 @@ import sys
|
||||
import socket
|
||||
import time
|
||||
|
||||
from repeater.companion.utils import validate_companion_node_name, normalize_companion_identity_key
|
||||
from repeater.companion.utils import (
|
||||
CompanionContactCapacityError,
|
||||
check_companion_contact_capacity,
|
||||
format_companion_bridge_limits,
|
||||
normalize_companion_identity_key,
|
||||
parse_companion_bridge_kwargs,
|
||||
effective_max_contacts,
|
||||
validate_companion_node_name,
|
||||
)
|
||||
from repeater.config import get_radio_for_board, load_config, save_config
|
||||
from repeater.config_manager import ConfigManager
|
||||
from repeater.data_acquisition.glass_handler import GlassHandler
|
||||
@@ -525,6 +533,16 @@ class RepeaterDaemon:
|
||||
break
|
||||
return _sync
|
||||
|
||||
bridge_kwargs = parse_companion_bridge_kwargs(settings)
|
||||
max_contacts = effective_max_contacts(bridge_kwargs)
|
||||
if sqlite_handler:
|
||||
check_companion_contact_capacity(
|
||||
companion_hash_str,
|
||||
max_contacts,
|
||||
sqlite_handler,
|
||||
companion_name=name,
|
||||
)
|
||||
|
||||
bridge = RepeaterCompanionBridge(
|
||||
identity=identity,
|
||||
packet_injector=self.router.inject_packet,
|
||||
@@ -533,6 +551,7 @@ class RepeaterDaemon:
|
||||
sqlite_handler=sqlite_handler,
|
||||
companion_hash=companion_hash_str,
|
||||
on_prefs_saved=_make_sync_node_name_to_config(name),
|
||||
**bridge_kwargs,
|
||||
)
|
||||
|
||||
# Load contacts from SQLite
|
||||
@@ -616,11 +635,15 @@ class RepeaterDaemon:
|
||||
identity_type="companion",
|
||||
)
|
||||
|
||||
limits = format_companion_bridge_limits(bridge_kwargs)
|
||||
logger.info(
|
||||
f"Loaded companion '{name}': hash=0x{companion_hash:02x}, "
|
||||
f"port={tcp_port}, bind={bind_address}, client_idle_timeout_sec={client_idle_timeout_sec}"
|
||||
f"port={tcp_port}, bind={bind_address}, "
|
||||
f"client_idle_timeout_sec={client_idle_timeout_sec}{limits}"
|
||||
)
|
||||
|
||||
except CompanionContactCapacityError as e:
|
||||
logger.error("%s", e)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load companion '{name}': {e}", exc_info=True)
|
||||
|
||||
@@ -686,6 +709,16 @@ class RepeaterDaemon:
|
||||
tcp_timeout_raw = settings.get("tcp_timeout", 120)
|
||||
client_idle_timeout_sec = None if tcp_timeout_raw == 0 else int(tcp_timeout_raw)
|
||||
|
||||
bridge_kwargs = parse_companion_bridge_kwargs(settings)
|
||||
max_contacts = effective_max_contacts(bridge_kwargs)
|
||||
if sqlite_handler:
|
||||
check_companion_contact_capacity(
|
||||
companion_hash_str,
|
||||
max_contacts,
|
||||
sqlite_handler,
|
||||
companion_name=name,
|
||||
)
|
||||
|
||||
bridge = RepeaterCompanionBridge(
|
||||
identity=identity,
|
||||
packet_injector=self.router.inject_packet,
|
||||
@@ -693,6 +726,7 @@ class RepeaterDaemon:
|
||||
radio_config=radio_config,
|
||||
sqlite_handler=sqlite_handler,
|
||||
companion_hash=companion_hash_str,
|
||||
**bridge_kwargs,
|
||||
)
|
||||
|
||||
if sqlite_handler:
|
||||
@@ -769,9 +803,11 @@ class RepeaterDaemon:
|
||||
identity_type="companion",
|
||||
)
|
||||
|
||||
limits = format_companion_bridge_limits(bridge_kwargs)
|
||||
logger.info(
|
||||
f"Hot-reload: Loaded companion '{name}': hash=0x{companion_hash:02x}, "
|
||||
f"port={tcp_port}, bind={bind_address}, client_idle_timeout_sec={client_idle_timeout_sec}"
|
||||
f"port={tcp_port}, bind={bind_address}, "
|
||||
f"client_idle_timeout_sec={client_idle_timeout_sec}{limits}"
|
||||
)
|
||||
|
||||
async def _on_raw_rx_for_companions(self, data: bytes, rssi: int, snr: float) -> None:
|
||||
|
||||
@@ -15,6 +15,12 @@ from repeater.companion.identity_resolve import (
|
||||
find_companion_index,
|
||||
heal_companion_empty_names,
|
||||
)
|
||||
from repeater.companion.utils import (
|
||||
CompanionContactCapacityError,
|
||||
merge_companion_settings_update,
|
||||
parse_companion_bridge_kwargs,
|
||||
validate_companion_config_capacity,
|
||||
)
|
||||
from repeater.config import resolve_storage_dir, update_unscoped_flood_policy
|
||||
from repeater.service_utils import get_buildroot_image_info
|
||||
|
||||
@@ -3021,6 +3027,11 @@ class APIEndpoints:
|
||||
if any(str(c.get("name") or "").strip() == name for c in companions):
|
||||
return self._error(f"Companion with name '{name}' already exists")
|
||||
|
||||
try:
|
||||
bridge_settings = parse_companion_bridge_kwargs(settings)
|
||||
except ValueError as e:
|
||||
return self._error(str(e))
|
||||
|
||||
comp_settings = {
|
||||
"node_name": settings.get("node_name") or name,
|
||||
"tcp_port": settings.get("tcp_port", 5000),
|
||||
@@ -3028,12 +3039,28 @@ class APIEndpoints:
|
||||
}
|
||||
if "tcp_timeout" in settings:
|
||||
comp_settings["tcp_timeout"] = settings["tcp_timeout"]
|
||||
comp_settings.update(bridge_settings)
|
||||
new_identity = {
|
||||
"name": name,
|
||||
"identity_key": identity_key,
|
||||
"type": identity_type,
|
||||
"settings": comp_settings,
|
||||
}
|
||||
sqlite_handler = None
|
||||
if self.repeater_handler and self.repeater_handler.storage:
|
||||
sqlite_handler = self.repeater_handler.storage.sqlite_handler
|
||||
if sqlite_handler and identity_key:
|
||||
try:
|
||||
validate_companion_config_capacity(
|
||||
new_identity,
|
||||
sqlite_handler,
|
||||
companion_name=name,
|
||||
settings=comp_settings,
|
||||
)
|
||||
except CompanionContactCapacityError as e:
|
||||
return self._error(str(e))
|
||||
except (ValueError, TypeError) as e:
|
||||
return self._error(str(e))
|
||||
companions.append(new_identity)
|
||||
self.config["identities"]["companions"] = companions
|
||||
else:
|
||||
@@ -3234,12 +3261,31 @@ class APIEndpoints:
|
||||
pass
|
||||
|
||||
if "settings" in data:
|
||||
if "settings" not in identity:
|
||||
identity["settings"] = {}
|
||||
# Only allow companion settings
|
||||
for k, v in data["settings"].items():
|
||||
if k in ("node_name", "tcp_port", "bind_address", "tcp_timeout"):
|
||||
identity["settings"][k] = v
|
||||
try:
|
||||
merged_settings = merge_companion_settings_update(
|
||||
identity.get("settings") or {},
|
||||
data["settings"],
|
||||
)
|
||||
except ValueError as e:
|
||||
return self._error(str(e))
|
||||
|
||||
sqlite_handler = None
|
||||
if self.repeater_handler and self.repeater_handler.storage:
|
||||
sqlite_handler = self.repeater_handler.storage.sqlite_handler
|
||||
if sqlite_handler and identity.get("identity_key"):
|
||||
try:
|
||||
validate_companion_config_capacity(
|
||||
identity,
|
||||
sqlite_handler,
|
||||
companion_name=resolved_name,
|
||||
settings=merged_settings,
|
||||
)
|
||||
except CompanionContactCapacityError as e:
|
||||
return self._error(str(e))
|
||||
except (ValueError, TypeError) as e:
|
||||
return self._error(str(e))
|
||||
|
||||
identity["settings"] = merged_settings
|
||||
|
||||
companions[identity_index] = identity
|
||||
self.config["identities"]["companions"] = companions
|
||||
|
||||
@@ -257,6 +257,10 @@ class CompanionAPIEndpoints:
|
||||
"is_running": b.is_running,
|
||||
"contacts_count": b.contacts.get_count(),
|
||||
"channels_count": b.channels.get_count(),
|
||||
"max_contacts": b.contacts.max_contacts,
|
||||
"offline_queue_size": getattr(b.message_queue, "max_size", None)
|
||||
or getattr(b.message_queue, "_max_size", None)
|
||||
or 512,
|
||||
}
|
||||
)
|
||||
return self._success(items)
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Tests for per-companion bridge settings parsing and startup guard."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from repeater.companion.utils import (
|
||||
CompanionContactCapacityError,
|
||||
check_companion_contact_capacity,
|
||||
effective_max_contacts,
|
||||
merge_companion_settings_update,
|
||||
parse_companion_bridge_kwargs,
|
||||
parse_positive_int,
|
||||
validate_companion_config_capacity,
|
||||
)
|
||||
# pymc_core defaults (CompanionBridge / ContactStore)
|
||||
_DEFAULT_MAX_CONTACTS = 1000
|
||||
|
||||
|
||||
class TestParsePositiveInt:
|
||||
def test_valid(self):
|
||||
assert parse_positive_int("100", "max_contacts") == 100
|
||||
|
||||
def test_invalid_type(self):
|
||||
with pytest.raises(ValueError, match="max_contacts"):
|
||||
parse_positive_int("abc", "max_contacts")
|
||||
|
||||
def test_below_minimum(self):
|
||||
with pytest.raises(ValueError, match="max_contacts"):
|
||||
parse_positive_int(0, "max_contacts")
|
||||
|
||||
|
||||
class TestParseCompanionBridgeKwargs:
|
||||
def test_empty_settings(self):
|
||||
assert parse_companion_bridge_kwargs({}) == {}
|
||||
|
||||
def test_max_contacts_and_offline_queue(self):
|
||||
assert parse_companion_bridge_kwargs(
|
||||
{"max_contacts": 2000, "offline_queue_size": 1024}
|
||||
) == {"max_contacts": 2000, "offline_queue_size": 1024}
|
||||
|
||||
def test_ignored_keys_warn(self, caplog):
|
||||
caplog.set_level(logging.WARNING)
|
||||
result = parse_companion_bridge_kwargs(
|
||||
{"max_contacts": 500, "max_channels": 64, "adv_type": 2}
|
||||
)
|
||||
assert result == {"max_contacts": 500}
|
||||
assert any("max_channels" in r.message for r in caplog.records)
|
||||
assert any("adv_type" in r.message for r in caplog.records)
|
||||
|
||||
def test_invalid_max_contacts(self):
|
||||
with pytest.raises(ValueError):
|
||||
parse_companion_bridge_kwargs({"max_contacts": -1})
|
||||
|
||||
|
||||
class TestEffectiveMaxContacts:
|
||||
def test_default(self):
|
||||
assert effective_max_contacts({}) == _DEFAULT_MAX_CONTACTS
|
||||
|
||||
def test_override(self):
|
||||
assert effective_max_contacts({"max_contacts": 500}) == 500
|
||||
|
||||
|
||||
class TestMergeCompanionSettingsUpdate:
|
||||
def test_merges_bridge_settings(self):
|
||||
merged = merge_companion_settings_update(
|
||||
{"node_name": "a"},
|
||||
{"max_contacts": 500},
|
||||
)
|
||||
assert merged == {"node_name": "a", "max_contacts": 500}
|
||||
|
||||
def test_unknown_key_raises(self):
|
||||
with pytest.raises(ValueError, match="Unknown companion setting"):
|
||||
merge_companion_settings_update({}, {"max_channels": 64})
|
||||
|
||||
|
||||
class TestValidateCompanionConfigCapacity:
|
||||
def test_uses_merged_settings_not_stale_identity(self):
|
||||
identity = {
|
||||
"identity_key": "aa" * 32,
|
||||
"settings": {"max_contacts": 1000},
|
||||
}
|
||||
sqlite = MagicMock()
|
||||
sqlite.companion_count_contacts.return_value = 600
|
||||
with pytest.raises(CompanionContactCapacityError):
|
||||
validate_companion_config_capacity(
|
||||
identity,
|
||||
sqlite,
|
||||
settings={"max_contacts": 500},
|
||||
)
|
||||
sqlite.companion_count_contacts.assert_called_once()
|
||||
|
||||
|
||||
class TestCheckCompanionContactCapacity:
|
||||
def test_skips_without_sqlite(self):
|
||||
check_companion_contact_capacity("0x01", 100, None)
|
||||
|
||||
def test_passes_when_under_limit(self):
|
||||
sqlite = MagicMock()
|
||||
sqlite.companion_count_contacts.return_value = 100
|
||||
check_companion_contact_capacity("0x01", 500, sqlite)
|
||||
|
||||
def test_raises_when_over_limit(self):
|
||||
sqlite = MagicMock()
|
||||
sqlite.companion_count_contacts.return_value = 812
|
||||
with pytest.raises(CompanionContactCapacityError) as exc:
|
||||
check_companion_contact_capacity(
|
||||
"0xab", 500, sqlite, companion_name="BotCompanion"
|
||||
)
|
||||
assert exc.value.stored_count == 812
|
||||
assert exc.value.max_contacts == 500
|
||||
assert "BotCompanion" in str(exc.value)
|
||||
Reference in New Issue
Block a user