Claude commited on
Commit
f58bb83
·
unverified ·
1 Parent(s): 575f8ed

fix(api): anchor CORS regex and add CORP middleware (BUG-002, BUG-003)

Browse files

BUG-002 - CORS Security Fix:
- Change regex from permissive `.*stroke-viewer-frontend.*` to
anchored `vibecodermcswaggins-stroke-viewer-frontend\.hf\.space`
- Set allow_credentials=False (not needed)
- Restrict allow_methods to ["GET", "POST"]
- Restrict allow_headers to ["Content-Type"]

BUG-003 - CORP Middleware for WebGL:
- Add CORPMiddleware that sets Cross-Origin-Resource-Policy: cross-origin
- Required for frontend COEP header to enable SharedArrayBuffer
- Add GPU availability check at startup with warning log

src/stroke_deepisles_demo/api/main.py CHANGED
@@ -19,9 +19,10 @@ from contextlib import asynccontextmanager
19
  from pathlib import Path
20
  from typing import Any
21
 
22
- from fastapi import FastAPI
23
  from fastapi.middleware.cors import CORSMiddleware
24
  from fastapi.staticfiles import StaticFiles
 
25
 
26
  from stroke_deepisles_demo.api.job_store import init_job_store
27
  from stroke_deepisles_demo.api.routes import router
@@ -47,6 +48,18 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
47
  # Startup
48
  logger.info("Starting stroke segmentation API...")
49
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  # Create results directory
51
  RESULTS_DIR.mkdir(parents=True, exist_ok=True)
52
 
@@ -68,6 +81,23 @@ app = FastAPI(
68
  lifespan=lifespan,
69
  )
70
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
71
  # CORS configuration
72
  FRONTEND_ORIGIN = os.environ.get("FRONTEND_ORIGIN", "")
73
  CORS_ORIGINS = [
@@ -77,16 +107,18 @@ CORS_ORIGINS = [
77
  if FRONTEND_ORIGIN:
78
  CORS_ORIGINS.append(FRONTEND_ORIGIN)
79
 
 
 
 
 
80
  app.add_middleware(
81
  CORSMiddleware,
82
  allow_origins=CORS_ORIGINS,
83
- # Match HF Spaces URLs in both formats:
84
- # - Direct: https://username-spacename.hf.space
85
- # - Proxy: https://username--spacename--hash.hf.space
86
- allow_origin_regex=r"https://.*stroke-viewer-frontend.*\.hf\.space",
87
- allow_credentials=True,
88
- allow_methods=["*"],
89
- allow_headers=["*"],
90
  )
91
 
92
  # API routes
 
19
  from pathlib import Path
20
  from typing import Any
21
 
22
+ from fastapi import FastAPI, Request, Response
23
  from fastapi.middleware.cors import CORSMiddleware
24
  from fastapi.staticfiles import StaticFiles
25
+ from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
26
 
27
  from stroke_deepisles_demo.api.job_store import init_job_store
28
  from stroke_deepisles_demo.api.routes import router
 
48
  # Startup
49
  logger.info("Starting stroke segmentation API...")
50
 
51
+ # Check for GPU availability (DeepISLES requires GPU)
52
+ try:
53
+ import torch
54
+
55
+ if not torch.cuda.is_available():
56
+ logger.warning(
57
+ "GPU not available! DeepISLES requires GPU for inference. "
58
+ "This Space should be configured with t4-small or better hardware."
59
+ )
60
+ except ImportError:
61
+ pass # torch may not be available in all environments
62
+
63
  # Create results directory
64
  RESULTS_DIR.mkdir(parents=True, exist_ok=True)
65
 
 
81
  lifespan=lifespan,
82
  )
83
 
84
+ # Cross-Origin Resource Policy middleware (required for COEP)
85
+ # This must be added BEFORE CORSMiddleware for proper header ordering
86
+ class CORPMiddleware(BaseHTTPMiddleware):
87
+ """Add Cross-Origin-Resource-Policy header to all responses.
88
+
89
+ Required when frontend uses COEP (Cross-Origin-Embedder-Policy: require-corp)
90
+ to enable SharedArrayBuffer for WebGL performance optimizations.
91
+ """
92
+
93
+ async def dispatch(
94
+ self, request: Request, call_next: RequestResponseEndpoint
95
+ ) -> Response:
96
+ response = await call_next(request)
97
+ response.headers["Cross-Origin-Resource-Policy"] = "cross-origin"
98
+ return response
99
+
100
+
101
  # CORS configuration
102
  FRONTEND_ORIGIN = os.environ.get("FRONTEND_ORIGIN", "")
103
  CORS_ORIGINS = [
 
107
  if FRONTEND_ORIGIN:
108
  CORS_ORIGINS.append(FRONTEND_ORIGIN)
109
 
110
+ # Add CORP middleware first (for COEP compatibility)
111
+ app.add_middleware(CORPMiddleware)
112
+
113
+ # Add CORS middleware with strict security settings
114
  app.add_middleware(
115
  CORSMiddleware,
116
  allow_origins=CORS_ORIGINS,
117
+ # Anchored regex: only allow our specific HF Space (security fix for BUG-002)
118
+ allow_origin_regex=r"https://vibecodermcswaggins-stroke-viewer-frontend\.hf\.space",
119
+ allow_credentials=False, # Not needed - no cookies/auth
120
+ allow_methods=["GET", "POST"], # Only methods we use
121
+ allow_headers=["Content-Type"], # Only headers we need
 
 
122
  )
123
 
124
  # API routes