mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-19 02:13:47 +00:00
feat(api): serve long-lived streams on the event loop
Move GET /api/events, the remote-device command stream and artifact downloads from Flask to Starlette routes mounted ahead of the Flask catch-all. On Flask each held an a2wsgi threadpool slot for as long as its response stayed open, and because uvicorn drops writes after a client disconnects, a closed tab never released it. - asgi_auth: one JWT/OIDC gate for Starlette routes; the chat reconnect reader uses it too - ClosingStreamingResponse closes the body iterator and releases the SSE slot or file handle even when the client leaves before the first frame - AsyncTopic liveness probe replaces the sync client's socket_timeout guard against half-open pub/sub sockets - ASYNC_REDIS_MAX_CONNECTIONS sizes the async Redis pool; every open stream holds a connection and redis-py defaults to 100
This commit is contained in:
1 parent
38773471a8
commit
54b540ea42
33 files changed
+2822
-1476
No files matched your search
+66
-28
@@ -6,7 +6,7 @@ import base64
|
||||
import hashlib
|
||||
import logging
|
||||
import time
|
||||
from typing import Optional, Tuple
|
||||
from typing import Callable, Mapping, Optional, Tuple
|
||||
|
||||
from flask import jsonify, make_response, request
|
||||
|
||||
@@ -17,6 +17,9 @@ from docsgpt.storage.db.session import db_readonly, db_session
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: ``(error_code, http_status)`` for a rejected device request.
|
||||
DeviceAuthFailure = Tuple[str, int]
|
||||
|
||||
|
||||
def hash_session_token(token: str) -> str:
|
||||
"""SHA-256 over the opaque session token. Stored in ``devices.token_hash``."""
|
||||
@@ -45,32 +48,46 @@ def _canonical_payload(method: str, path: str, ts: str, body: bytes) -> str:
|
||||
return f"{method} {path} {ts} {body_hash}"
|
||||
|
||||
|
||||
def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Optional[tuple]]:
|
||||
"""Validate the device session token on the current request.
|
||||
def authenticate_device(
|
||||
headers: Mapping[str, str],
|
||||
method: str,
|
||||
path: str,
|
||||
get_body: Callable[[], bytes],
|
||||
*,
|
||||
touch: bool = True,
|
||||
) -> Tuple[Optional[dict], Optional[DeviceAuthFailure]]:
|
||||
"""Validate a device session token from the raw request parts.
|
||||
|
||||
Returns ``(device_row, None)`` on success or ``(None, response)`` on
|
||||
failure, where ``response`` is a ``(json, status)`` Flask tuple.
|
||||
Framework-neutral core shared by the Flask routes and the native-async
|
||||
SSE stream, which calls it from a worker thread since it reads Postgres.
|
||||
|
||||
Args:
|
||||
headers: Request headers (a case-insensitive mapping).
|
||||
method: HTTP method, as the CLI signed it.
|
||||
path: Request path without the query string, as the CLI signed it.
|
||||
get_body: Returns the raw body; only called when signatures are required.
|
||||
touch: When True, bump ``last_seen_at`` on the device row.
|
||||
|
||||
Returns:
|
||||
tuple: ``(device_row, None)`` on success or ``(None, (code, status))``.
|
||||
"""
|
||||
auth = request.headers.get("Authorization", "")
|
||||
auth = headers.get("Authorization", "")
|
||||
if not auth.lower().startswith("bearer "):
|
||||
return None, _error("missing_token", 401)
|
||||
return None, ("missing_token", 401)
|
||||
token = auth.split(" ", 1)[1].strip()
|
||||
if not token:
|
||||
return None, _error("missing_token", 401)
|
||||
return None, ("missing_token", 401)
|
||||
|
||||
token_hash = hash_session_token(token)
|
||||
with db_readonly() as conn:
|
||||
device = DevicesRepository(conn).find_by_token_hash(token_hash)
|
||||
if device is None:
|
||||
return None, _error("invalid_token", 401)
|
||||
return None, ("invalid_token", 401)
|
||||
|
||||
if settings.REMOTE_DEVICE_REQUIRE_SIGNATURE:
|
||||
sig_error = _verify_signature(device)
|
||||
if sig_error is not None:
|
||||
return None, sig_error
|
||||
failure = _verify_signature(device, headers, method, path, get_body)
|
||||
if failure is not None:
|
||||
return None, failure
|
||||
|
||||
if touch:
|
||||
try:
|
||||
@@ -82,23 +99,46 @@ def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Option
|
||||
return device, None
|
||||
|
||||
|
||||
def _verify_signature(device: dict) -> Optional[tuple]:
|
||||
def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Optional[tuple]]:
|
||||
"""Validate the device session token on the current Flask request.
|
||||
|
||||
Returns ``(device_row, None)`` on success or ``(None, response)`` on
|
||||
failure, where ``response`` is a ready Flask error response.
|
||||
|
||||
Args:
|
||||
touch: When True, bump ``last_seen_at`` on the device row.
|
||||
"""
|
||||
device, failure = authenticate_device(
|
||||
request.headers, request.method, request.path, request.get_data, touch=touch
|
||||
)
|
||||
if failure is not None:
|
||||
return None, _error(*failure)
|
||||
return device, None
|
||||
|
||||
|
||||
def _verify_signature(
|
||||
device: dict,
|
||||
headers: Mapping[str, str],
|
||||
method: str,
|
||||
path: str,
|
||||
get_body: Callable[[], bytes],
|
||||
) -> Optional[DeviceAuthFailure]:
|
||||
"""Verify ``X-Device-Signature`` against the stored machine pubkey."""
|
||||
sig_b64 = request.headers.get("X-Device-Signature")
|
||||
ts = request.headers.get("X-Device-Timestamp")
|
||||
fp = request.headers.get("X-Device-Machine-Key")
|
||||
sig_b64 = headers.get("X-Device-Signature")
|
||||
ts = headers.get("X-Device-Timestamp")
|
||||
fp = headers.get("X-Device-Machine-Key")
|
||||
if not sig_b64 or not ts or not fp:
|
||||
return _error("missing_signature", 401)
|
||||
return ("missing_signature", 401)
|
||||
|
||||
try:
|
||||
ts_int = int(ts)
|
||||
except (TypeError, ValueError):
|
||||
return _error("invalid_timestamp", 401)
|
||||
return ("invalid_timestamp", 401)
|
||||
if abs(time.time() - ts_int) > 300:
|
||||
return _error("timestamp_skew", 401)
|
||||
return ("timestamp_skew", 401)
|
||||
|
||||
if fp != device.get("machine_pubkey_fingerprint"):
|
||||
return _error("fingerprint_mismatch", 401)
|
||||
return ("fingerprint_mismatch", 401)
|
||||
|
||||
# Defer cryptography import to keep cold-start light when signatures
|
||||
# are disabled (the dev default).
|
||||
@@ -111,26 +151,24 @@ def _verify_signature(device: dict) -> Optional[tuple]:
|
||||
logger.error(
|
||||
"Signature verification requested but ``cryptography`` is not installed."
|
||||
)
|
||||
return _error("signature_unsupported", 500)
|
||||
return ("signature_unsupported", 500)
|
||||
|
||||
# The full public key isn't stored — only the fingerprint. For signature
|
||||
# verification we accept the pubkey in a header too; we trust it iff
|
||||
# its fingerprint matches the stored one. This avoids a separate pubkey
|
||||
# column for MVP.
|
||||
pubkey_b64 = request.headers.get("X-Device-Machine-Pubkey")
|
||||
pubkey_b64 = headers.get("X-Device-Machine-Pubkey")
|
||||
if not pubkey_b64:
|
||||
return _error("missing_pubkey", 401)
|
||||
return ("missing_pubkey", 401)
|
||||
if fingerprint_pubkey(pubkey_b64) != device["machine_pubkey_fingerprint"]:
|
||||
return _error("pubkey_fingerprint_mismatch", 401)
|
||||
return ("pubkey_fingerprint_mismatch", 401)
|
||||
|
||||
payload = _canonical_payload(
|
||||
request.method, request.path, ts, request.get_data()
|
||||
).encode("utf-8")
|
||||
payload = _canonical_payload(method, path, ts, get_body()).encode("utf-8")
|
||||
try:
|
||||
pubkey = Ed25519PublicKey.from_public_bytes(base64.b64decode(pubkey_b64))
|
||||
pubkey.verify(base64.b64decode(sig_b64), payload)
|
||||
except (InvalidSignature, ValueError):
|
||||
return _error("invalid_signature", 401)
|
||||
return ("invalid_signature", 401)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
Reference in new issue
Block a user