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:
co-authored by
Claude Opus 4.8
parent
fa1503c4eb
commit
f7cfd33994
+36
-13
@@ -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(
|
||||
|
||||
@@ -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
@@ -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>,
|
||||
);
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()))
|
||||
Reference in New Issue
Block a user