init

backend/app/auth.py

@@ -0,0 +1,95 @@
+from functools import lru_cache
+from typing import Any, Dict, Optional
+from urllib.parse import urljoin
+
+import httpx
+from fastapi import HTTPException, Security
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from jose import JWTError, jwk, jwt
+from pydantic import BaseModel, BaseSettings, Field, HttpUrl
+
+security = HTTPBearer(auto_error=False)
+
+class Settings(BaseSettings):
+    oidc_issuer: HttpUrl = Field(..., env="OIDC_ISSUER")
+    oidc_audience: str = Field(..., env="OIDC_AUDIENCE")
+    oidc_client_id: str = Field(..., env="OIDC_CLIENT_ID")
+    oidc_client_secret: Optional[str] = Field(None, env="OIDC_CLIENT_SECRET")
+    oidc_verify_iss: bool = Field(True, env="OIDC_VERIFY_ISS")
+
+    class Config:
+        env_file = ".env"
+
+settings = Settings()
+
+class TokenClaims(BaseModel):
+    iss: str
+    sub: str
+    aud: Any
+    exp: int
+    iat: Optional[int]
+    azp: Optional[str]
+    email: Optional[str]
+    preferred_username: Optional[str]
+
+@lru_cache(maxsize=1)
+def get_discovery() -> Dict[str, Any]:
+    issuer = str(settings.oidc_issuer).rstrip("/")
+    discovery_url = urljoin(issuer + "/", ".well-known/openid-configuration")
+    with httpx.Client(timeout=10.0) as client:
+        response = client.get(discovery_url)
+        if response.status_code != 200:
+            raise RuntimeError(f"Unable to load OIDC discovery document from {discovery_url}")
+        return response.json()
+
+@lru_cache(maxsize=1)
+def get_jwks() -> Dict[str, Any]:
+    discovery = get_discovery()
+    jwks_uri = discovery.get("jwks_uri")
+    if not jwks_uri:
+        raise RuntimeError("OIDC discovery document does not contain jwks_uri")
+    with httpx.Client(timeout=10.0) as client:
+        response = client.get(jwks_uri)
+        if response.status_code != 200:
+            raise RuntimeError(f"Unable to load JWKS from {jwks_uri}")
+        return response.json()
+
+def get_signing_key(token: str) -> Dict[str, Any]:
+    try:
+        headers = jwt.get_unverified_header(token)
+    except JWTError as exc:
+        raise HTTPException(status_code=401, detail="Invalid JWT header") from exc
+
+    kid = headers.get("kid")
+    if not kid:
+        raise HTTPException(status_code=401, detail="JWT missing kid header")
+
+    jwks = get_jwks().get("keys", [])
+    for key in jwks:
+        if key.get("kid") == kid:
+            return key
+
+    raise HTTPException(status_code=401, detail="Unable to find matching JWKS key")
+
+async def validate_token(credentials: HTTPAuthorizationCredentials = Security(security)) -> TokenClaims:
+    if not credentials or credentials.scheme.lower() != "bearer":
+        raise HTTPException(status_code=401, detail="Missing Bearer authorization header")
+
+    token = credentials.credentials
+    key = get_signing_key(token)
+    public_key = jwk.construct(key)
+    try:
+        verified = jwt.decode(
+            token,
+            public_key,
+            algorithms=[key.get("alg", "RS256")],
+            audience=settings.oidc_audience,
+            issuer=str(settings.oidc_issuer) if settings.oidc_verify_iss else None,
+        )
+    except JWTError as exc:
+        raise HTTPException(status_code=401, detail="Invalid or expired token") from exc
+
+    try:
+        return TokenClaims(**verified)
+    except Exception as exc:
+        raise HTTPException(status_code=401, detail="Unable to parse token claims") from exc

backend/app/main.py

@@ -0,0 +1,96 @@
+import os
+from typing import AsyncGenerator
+
+import httpx
+from fastapi import Depends, FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+
+from app.auth import get_discovery, settings, validate_token
+from app.schemas import AuthCallbackRequest, AuthConfigResponse, ChatRequest
+
+OLLAMA_API_KEY = os.getenv("OLLAMA_API_KEY")
+OLLAMA_API_URL = os.getenv("OLLAMA_API_URL", "https://api.ollama.com/v1/chat/completions")
+if not OLLAMA_API_KEY:
+    raise RuntimeError("Environment variable OLLAMA_API_KEY is required")
+
+app = FastAPI(
+    title="OIDC-Protected Ollama Proxy",
+    description="FastAPI backend validating OIDC JWTs and streaming Ollama Cloud responses over SSE.",
+    version="0.1.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+async def event_stream(request: ChatRequest) -> AsyncGenerator[str, None]:
+    payload = {
+        "model": request.model,
+        "messages": [message.dict() for message in request.messages],
+        "stream": True,
+    }
+    headers = {
+        "Authorization": f"Bearer {OLLAMA_API_KEY}",
+        "Content-Type": "application/json",
+        "Accept": "text/event-stream",
+    }
+
+    async with httpx.AsyncClient(timeout=None) as client:
+        async with client.stream("POST", OLLAMA_API_URL, json=payload, headers=headers) as response:
+            if response.status_code >= 400:
+                text = await response.aread()
+                raise HTTPException(status_code=502, detail=f"Ollama Cloud error: {text.decode('utf-8', 'ignore')}" )
+
+            async for chunk in response.aiter_text():
+                if not chunk:
+                    continue
+                for line in chunk.splitlines():
+                    if line.strip():
+                        yield f"data: {line}\n\n"
+
+@app.get("/health")
+async def health() -> dict:
+    return {"status": "ok"}
+
+@app.get("/auth/config", response_model=AuthConfigResponse)
+async def auth_config() -> AuthConfigResponse:
+    discovery = get_discovery()
+    return AuthConfigResponse(
+        issuer=discovery["issuer"],
+        authorization_endpoint=discovery["authorization_endpoint"],
+        token_endpoint=discovery["token_endpoint"],
+    )
+
+@app.post("/auth/callback")
+async def auth_callback(payload: AuthCallbackRequest):
+    discovery = get_discovery()
+    token_endpoint = discovery.get("token_endpoint")
+    if not token_endpoint:
+        raise HTTPException(status_code=502, detail="OIDC discovery document missing token_endpoint")
+
+    data = {
+        "grant_type": "authorization_code",
+        "client_id": settings.oidc_client_id,
+        "code": payload.code,
+        "code_verifier": payload.code_verifier,
+        "redirect_uri": str(payload.redirect_uri),
+    }
+    if settings.oidc_client_secret:
+        data["client_secret"] = settings.oidc_client_secret
+
+    async with httpx.AsyncClient(timeout=15.0) as client:
+        response = await client.post(token_endpoint, data=data, headers={"Accept": "application/json"})
+
+    if response.status_code != 200:
+        raise HTTPException(status_code=400, detail=f"OIDC token exchange failed: {response.text}")
+
+    return response.json()
+
+@app.post("/chat", response_class=StreamingResponse)
+async def chat(request: ChatRequest, token=Depends(validate_token)):
+    return StreamingResponse(event_stream(request), media_type="text/event-stream")

backend/app/schemas.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel, Field, HttpUrl
+from typing import Literal, List
+
+class ChatMessage(BaseModel):
+    role: Literal["user", "assistant", "system"]
+    content: str
+
+class ChatRequest(BaseModel):
+    model: str = Field(default="llama2")
+    messages: List[ChatMessage] = Field(default_factory=list)
+
+class AuthCallbackRequest(BaseModel):
+    code: str
+    code_verifier: str
+    redirect_uri: HttpUrl
+
+class AuthConfigResponse(BaseModel):
+    issuer: HttpUrl
+    authorization_endpoint: HttpUrl
+    token_endpoint: HttpUrl