polish(network): outermost PIN gate + non-buffering ASGI middleware + listener test (#160)

* fix(network-share): mount RemoteAuthGate at outermost provider

Move the <RemoteAuthGate> wrap from App.jsx's main-studio return up to
main-app.jsx, inside QueryClientProvider and wrapping the entire app tree
(both the dictation widget and <App />). Previously the gate only wrapped
the studio return, so a remote device opening a bare URL (no ?pin=) during
first-run states — the /setup/status check, SetupWizard, or BootstrapSplash
early returns — would 401 with no gate rendered to collect the PIN. The QR
path was fine (PIN captured pre-fetch in client.ts); only bare-URL was broken.

Remove the App.jsx wrap to avoid double-gating (two PIN dialogs). No behavior
change for loopback or QR users.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* perf(network-share): make NetworkAccessMiddleware non-buffering ASGI

Rewrite NetworkAccessMiddleware from a starlette BaseHTTPMiddleware into a
pure ASGI middleware (class with __init__(app) and __call__(scope, receive,
send)). BaseHTTPMiddleware buffers StreamingResponse/SSE bodies before
forwarding them, so PIN'd LAN clients on streaming endpoints (dictation SSE,
tts streaming, /system/logs/stream) got buffered/laggy responses. Loopback was
unaffected (bypasses early), but remote-share streaming was degraded.

The ASGI form forwards send untouched on every pass-through path, and only
wraps send to inject Set-Cookie on the http.response.start message for the
first valid-PIN request — the body keeps streaming chunk-by-chunk. request.app
resolves in ASGI scope (Starlette sets scope["app"]), so the inert/loopback/
shell/PIN logic is identical to before. Registered after CORS (unchanged) so
CORS stays outermost.

All 5 existing behavior tests pass unchanged. Adds three tests: a guard that
the middleware is not a BaseHTTPMiddleware subclass, a StreamingResponse
pass-through (401 without PIN, full chunked stream with PIN, no buffered
Content-Length), and a Set-Cookie-via-ASGI assertion.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(network-share): integration test for real listener lifecycle

Add tests/test_network_share_lifecycle.py exercising the real second uvicorn
listener: await network_share.enable(app) on a minimal FastAPI app, assert
get_state().enabled is True with a share_port set and a live TCP listener on
that port (real socket connect), then await disable(app) and assert the state
resets and the port stops accepting connections.

Uses the returned share_port (never a hardcoded port) and tolerates teardown
timing by polling for socket close. Wrapped in asyncio.run inside a sync test
so it does not depend on a pytest-asyncio event-loop mode; skips gracefully if
binding 0.0.0.0 is not permitted in the sandbox. Defensive cleanup resets the
module-level state on any failure path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Palash Debnath
2026-05-30 11:45:54 +05:30
committed by GitHub
co-authored by Claude Opus 4.8
parent fa1503c4eb
commit f7cfd33994
5 changed files with 243 additions and 46 deletions
+36 -13
View File
@@ -256,7 +256,7 @@ from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, RedirectResponse, Response
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.datastructures import MutableHeaders
from scalar_fastapi import get_scalar_api_reference
import traceback
@@ -484,23 +484,38 @@ _LOOPBACK_CLIENTS = {"127.0.0.1", "::1"}
_SHELL_PATHS = {"/", "/index.html", "/favicon.ico", "/health"}
class NetworkAccessMiddleware(BaseHTTPMiddleware):
class NetworkAccessMiddleware:
"""When a share PIN is set, require it for non-loopback clients on API
routes. Inert when no PIN (default + docker deploys). Loopback (incl.
Tailscale-proxied) always bypasses; the SPA shell is always served so the
PIN gate UI can load."""
PIN gate UI can load.
async def dispatch(self, request, call_next):
Pure ASGI (not BaseHTTPMiddleware) so it never buffers the response body.
BaseHTTPMiddleware collects StreamingResponse/SSE bodies before forwarding,
which makes PIN'd LAN clients on streaming endpoints (dictation SSE, tts
streaming, /system/logs/stream) laggy. As a plain ASGI app we forward
`send` untouched on the pass-through paths and only wrap it to inject the
Set-Cookie header — the body still streams chunk-by-chunk."""
def __init__(self, app):
self.app = app
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
return await self.app(scope, receive, send)
from starlette.requests import Request
request = Request(scope, receive=receive)
ns = getattr(request.app.state, "network_share", None)
pin = getattr(ns, "pin", None) if ns else None
if not pin:
return await call_next(request)
client = request.client.host if request.client else None
return await self.app(scope, receive, send)
client = scope["client"][0] if scope.get("client") else None
if client in _LOOPBACK_CLIENTS:
return await call_next(request)
path = request.url.path
return await self.app(scope, receive, send)
path = scope["path"]
if path in _SHELL_PATHS or path.startswith("/assets/") or path.startswith("/favicon"):
return await call_next(request)
return await self.app(scope, receive, send)
supplied = (
request.headers.get("x-omnivoice-pin")
or request.query_params.get("pin")
@@ -508,11 +523,19 @@ class NetworkAccessMiddleware(BaseHTTPMiddleware):
or ""
)
if not secrets.compare_digest(supplied, pin):
return JSONResponse({"detail": "PIN required"}, status_code=401)
response = await call_next(request)
resp = JSONResponse({"detail": "PIN required"}, status_code=401)
return await resp(scope, receive, send)
# Valid PIN. Set the cookie by wrapping send to inject Set-Cookie on the
# http.response.start message — without ever materialising the body.
if request.cookies.get("ov_pin") != pin:
response.set_cookie("ov_pin", pin, samesite="lax")
return response
async def send_with_cookie(message):
if message["type"] == "http.response.start":
headers = MutableHeaders(scope=message)
headers.append("set-cookie", f"ov_pin={pin}; Path=/; SameSite=Lax")
await send(message)
return await self.app(scope, receive, send_with_cookie)
return await self.app(scope, receive, send)
_allowed = os.environ.get(
+4 -3
View File
@@ -30,7 +30,10 @@ import Header from './components/Header';
import NavRail from './components/NavRail';
import ErrorBoundary from './components/ErrorBoundary';
import FloatingPill from './components/FloatingPill';
import RemoteAuthGate from './components/RemoteAuthGate';
// RemoteAuthGate is mounted at the true outermost provider in main-app.jsx so
// it covers all app states (setup check / wizard / bootstrap), not just the
// main studio return below. Do not re-wrap here — double-gating renders two
// PIN dialogs.
import useRealtimeEvents from './hooks/useRealtimeEvents';
import { BootstrapSplash, useBootstrapStage } from './components/BootstrapSplash';
@@ -822,7 +825,6 @@ function App() {
}
return (
<RemoteAuthGate>
<div
className={[
'app-container',
@@ -1144,7 +1146,6 @@ function App() {
</Suspense>
</div>
</RemoteAuthGate>
);
}
+38 -30
View File
@@ -13,6 +13,7 @@ import './i18n'; // ← initialise i18next before any component renders
import './ui';
import './index.css';
import App from './App.jsx';
import RemoteAuthGate from './components/RemoteAuthGate';
import { installConsoleCapture } from './utils/consoleBuffer.js';
installConsoleCapture();
@@ -52,36 +53,43 @@ export async function bootstrapApp() {
createRoot(document.getElementById('root')).render(
<StrictMode>
<QueryClientProvider client={queryClient}>
{isWidget ? (
<Suspense
fallback={
<div
style={{
position: 'fixed',
inset: 0,
display: 'flex',
alignItems: 'center',
justifyContent: 'center',
background: 'rgba(18, 18, 22, 0.88)',
backdropFilter: 'blur(24px) saturate(180%)',
WebkitBackdropFilter: 'blur(24px) saturate(180%)',
border: '1px solid rgba(255, 255, 255, 0.08)',
borderRadius: '100px',
color: 'rgba(255, 255, 255, 0.9)',
fontFamily: '"Inter Variable", "Inter", -apple-system, sans-serif',
fontSize: 13,
userSelect: 'none',
}}
>
Loading dictation
</div>
}
>
<CaptureWidget />
</Suspense>
) : (
<App />
)}
{/* RemoteAuthGate is the TRUE outermost wrap so a remote device that
loads a bare URL (no ?pin=) during first-run setup states —
setup-status check, SetupWizard, BootstrapSplash — still gets the
PIN dialog instead of a silent 401. Loopback / QR users are
unaffected (the gate only shows on an ov:pin-required event). */}
<RemoteAuthGate>
{isWidget ? (
<Suspense
fallback={
<div
style={{
position: 'fixed',
inset: 0,
display: 'flex',
alignItems: 'center',
justifyContent: 'center',
background: 'rgba(18, 18, 22, 0.88)',
backdropFilter: 'blur(24px) saturate(180%)',
WebkitBackdropFilter: 'blur(24px) saturate(180%)',
border: '1px solid rgba(255, 255, 255, 0.08)',
borderRadius: '100px',
color: 'rgba(255, 255, 255, 0.9)',
fontFamily: '"Inter Variable", "Inter", -apple-system, sans-serif',
fontSize: 13,
userSelect: 'none',
}}
>
Loading dictation
</div>
}
>
<CaptureWidget />
</Suspense>
) : (
<App />
)}
</RemoteAuthGate>
</QueryClientProvider>
</StrictMode>,
);
+71
View File
@@ -41,3 +41,74 @@ def test_non_loopback_with_valid_pin_passes():
def test_spa_shell_served_without_pin():
c = TestClient(_app_with_pin(), client=("10.0.0.5", 1))
assert c.get("/health").status_code == 200
def test_middleware_is_plain_asgi_not_buffering():
# A pure ASGI middleware (class with __call__(scope, receive, send)) does
# NOT subclass starlette's BaseHTTPMiddleware, which buffers streaming
# responses. Guard against a regression back to the buffering base class.
from starlette.middleware.base import BaseHTTPMiddleware
from main import NetworkAccessMiddleware
assert not issubclass(NetworkAccessMiddleware, BaseHTTPMiddleware)
assert callable(getattr(NetworkAccessMiddleware, "__call__", None))
def test_streaming_response_passes_through_with_valid_pin():
# A PIN'd, non-loopback request to a StreamingResponse route must stream
# chunk-by-chunk, not be collected into one buffered body. We mount a tiny
# streaming route on a fresh app wrapped with the real middleware and
# confirm the response arrives chunked (multiple yields concatenated).
from fastapi import FastAPI
from starlette.responses import StreamingResponse
from main import NetworkAccessMiddleware
app = FastAPI()
app.add_middleware(NetworkAccessMiddleware)
app.state.network_share = ns.ShareState(
enabled=True, share_port=3901, pin="777888", lan_addresses=["10.0.0.9"]
)
@app.get("/stream")
def stream():
def gen():
for i in range(5):
yield f"chunk-{i}\n"
return StreamingResponse(gen(), media_type="text/plain")
c = TestClient(app, client=("10.0.0.5", 1))
# Without the PIN, the stream route is gated.
assert c.get("/stream").status_code == 401
# With the PIN, it streams the full body through the ASGI middleware.
r = c.get("/stream", headers={"X-OmniVoice-Pin": "777888"})
assert r.status_code == 200
body = r.text
for i in range(5):
assert f"chunk-{i}" in body
# Streaming responses carry no precomputed Content-Length — a buffering
# middleware would re-materialise the body and set one.
assert "content-length" not in {k.lower() for k in r.headers}
def test_valid_pin_sets_cookie_via_asgi():
from fastapi import FastAPI
from main import NetworkAccessMiddleware
app = FastAPI()
app.add_middleware(NetworkAccessMiddleware)
app.state.network_share = ns.ShareState(
enabled=True, share_port=3901, pin="424242", lan_addresses=["10.0.0.9"]
)
@app.get("/api/ping")
def ping():
return {"ok": True}
c = TestClient(app, client=("10.0.0.5", 1))
r = c.get("/api/ping", headers={"X-OmniVoice-Pin": "424242"})
assert r.status_code == 200
# The ASGI send-wrapper injects Set-Cookie on the first valid-PIN request
# (when the cookie isn't already present).
set_cookie = r.headers.get("set-cookie", "")
assert "ov_pin=424242" in set_cookie
+94
View File
@@ -0,0 +1,94 @@
"""Integration test for the real network-share listener lifecycle.
enable() starts a SECOND in-process uvicorn.Server bound to 0.0.0.0 on a
dedicated port serving the same app; disable() stops it. This test exercises
the real socket: it confirms a TCP connection succeeds on the reported
share_port while enabled, and is refused after disable().
Wrapped in asyncio.run inside a sync test so it does not depend on a
pytest-asyncio event-loop mode being configured.
"""
import asyncio
import socket
import pytest
from services import network_share as ns
def _can_connect(port: int, host: str = "127.0.0.1", timeout: float = 0.5) -> bool:
"""True if a TCP connection to host:port is accepted."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.settimeout(timeout)
try:
s.connect((host, port))
return True
except OSError:
return False
def _wait_closed(port: int, host: str = "127.0.0.1", tries: int = 40) -> bool:
"""Poll until the port stops accepting connections (tolerates teardown lag)."""
for _ in range(tries):
if not _can_connect(port, host):
return True
# Synchronous sleep is fine here — runs outside the event loop, between
# connect probes, after the server has been asked to exit.
socket_wait = 0.05
import time
time.sleep(socket_wait)
return False
async def _exercise_lifecycle():
# A minimal FastAPI app is enough — enable() only needs an ASGI app object
# and a place to stash app.state.network_share.
from fastapi import FastAPI
app = FastAPI()
# Sanity: starts Local (nothing bound to 0.0.0.0).
assert ns.get_state().enabled is False
state = await ns.enable(app)
try:
assert state.enabled is True
assert ns.get_state().enabled is True
port = state.share_port
assert isinstance(port, int) and port > 0
# app.state is updated to the enabled state.
assert app.state.network_share.enabled is True
assert app.state.network_share.share_port == port
# The listener is really up: a TCP connect to the reported port succeeds.
# The server binds 0.0.0.0; connect via loopback, which 0.0.0.0 covers.
assert _can_connect(port), f"expected a live listener on port {port}"
finally:
await ns.disable(app)
# After disable(): state reset and the socket is closed.
assert ns.get_state().enabled is False
assert app.state.network_share.enabled is False
assert _wait_closed(port), f"expected port {port} closed after disable()"
return port
def test_share_listener_lifecycle():
# Ensure a clean starting state regardless of test ordering.
if ns.get_state().enabled:
asyncio.run(ns.disable(__import__("fastapi").FastAPI()))
try:
# Probe whether binding 0.0.0.0 is permitted in this sandbox.
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
try:
probe.bind(("0.0.0.0", 0))
except OSError as e:
pytest.skip(f"binding 0.0.0.0 not permitted in this sandbox: {e}")
asyncio.run(_exercise_lifecycle())
finally:
# Defensive cleanup so a failure mid-test never leaves a stray listener
# or a dirty module-level _state for the next test.
if ns.get_state().enabled:
asyncio.run(ns.disable(__import__("fastapi").FastAPI()))