VoiceStudio/backend/services/network_share.py

"""Same-process LAN share listener + access PIN.

Enabling starts a SECOND uvicorn.Server bound to 0.0.0.0 on a dedicated port,
serving the SAME FastAPI app object — so the loaded model and in-flight jobs
are untouched (no restart). Disabling stops it, closing the 0.0.0.0 socket.
Loopback-only by default: nothing binds 0.0.0.0 until enable() is called.
"""
import asyncio
import secrets
import socket
from dataclasses import dataclass, field
from typing import Optional

import psutil
import uvicorn

BACKEND_PORT = 3900  # must match backend/main.py uvicorn.run(port=...)


@dataclass
class ShareState:
    enabled: bool = False
    share_port: Optional[int] = None
    pin: Optional[str] = None
    lan_addresses: list = field(default_factory=list)


_state = ShareState()
_server: Optional["uvicorn.Server"] = None
_task: Optional["asyncio.Task"] = None


def lan_ipv4_addresses() -> list:
    out, seen = [], set()
    for _name, addrs in psutil.net_if_addrs().items():
        for a in addrs:
            if a.family == socket.AF_INET:
                ip = a.address
                if ip.startswith("127.") or ip.startswith("169.254."):
                    continue
                if ip not in seen:
                    seen.add(ip)
                    out.append(ip)
    return out


def _gen_pin() -> str:
    return f"{secrets.randbelow(900000) + 100000}"  # 100000-999999


def _find_free_port(base: int, tries: int = 20) -> int:
    for p in range(base, base + tries):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind(("0.0.0.0", p))
                return p
            except OSError:
                continue
    raise RuntimeError("no free share port available")


def get_state() -> ShareState:
    return _state


async def enable(app) -> ShareState:
    global _server, _task, _state
    if _state.enabled:
        return _state
    port = _find_free_port(BACKEND_PORT + 1)
    pin = _gen_pin()
    config = uvicorn.Config(app, host="0.0.0.0", port=port, log_level="warning")
    server = uvicorn.Server(config)
    server.install_signal_handlers = lambda: None  # never hijack signals in-process
    _task = asyncio.create_task(server.serve())
    for _ in range(100):  # ~5s for the socket to bind
        if getattr(server, "started", False):
            break
        await asyncio.sleep(0.05)
    if not getattr(server, "started", False):
        # Bind failed (e.g. the port was taken in the race after the
        # free-port probe). Tear down and stay Local — never report enabled
        # with a listener that isn't actually up (spec §7).
        server.should_exit = True
        try:
            await asyncio.wait_for(_task, timeout=2)
        except Exception:
            pass
        _task = None
        raise RuntimeError("share listener failed to start")
    _server = server
    _state = ShareState(True, port, pin, lan_ipv4_addresses())
    app.state.network_share = _state
    return _state


async def disable(app) -> ShareState:
    global _server, _task, _state
    if _server is not None:
        _server.should_exit = True
        if _task is not None:
            try:
                await asyncio.wait_for(_task, timeout=5)
            except Exception:
                pass
    _server = _task = None
    _state = ShareState()
    app.state.network_share = _state
    return _state