refactor
This commit is contained in:
tao.chen
+221
-66
@@ -1,84 +1,239 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, Request, status
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
|
||||
from sqlalchemy import select
|
||||
|
||||
from backend.dependencies import (
|
||||
RequestContext,
|
||||
request_context,
|
||||
)
|
||||
from backend.dependencies import database_session
|
||||
from backend.runtime_client import RuntimeClientError
|
||||
from backend.schemas import (
|
||||
CreateJupyterAccessTicketRequest,
|
||||
)
|
||||
from common.db.models import Scripts, Users, WorkspaceMembers, Workspaces
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
|
||||
router = APIRouter(tags=["jupyter"])
|
||||
security = HTTPBearer(auto_error=False)
|
||||
|
||||
|
||||
def cookie_secure() -> bool:
|
||||
return os.getenv("COOKIE_SECURE", "false").strip().lower() in {
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
}
|
||||
JWT_SECRET = os.environ.get("JWT_SECRET", "dev-only-not-for-production")
|
||||
JWT_ALGORITHM = "HS256"
|
||||
|
||||
|
||||
@router.post(
|
||||
"/api/v1/jupyter/access-tickets",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def create_jupyter_access_ticket(
|
||||
payload: CreateJupyterAccessTicketRequest,
|
||||
request: Request,
|
||||
context: RequestContext = Depends(request_context),
|
||||
) -> JSONResponse:
|
||||
def _b64decode(value: str) -> bytes:
|
||||
padding = "=" * (-len(value) % 4)
|
||||
return base64.urlsafe_b64decode(value + padding)
|
||||
|
||||
|
||||
def verify_jwt_token(token: str) -> dict:
|
||||
"""Parse a signed JWT from the access_token cookie.
|
||||
|
||||
Returns a payload dict containing at least ``sub`` (user_id) and
|
||||
``exp``; raises 401 on missing, malformed, or expired tokens. The
|
||||
current implementation is intentionally minimal — replace with a
|
||||
real verification once a public-key/HSM source is wired in.
|
||||
"""
|
||||
if not token:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Missing Authentication Token",
|
||||
)
|
||||
|
||||
try:
|
||||
ticket_data = (
|
||||
await request.app.state.runtime_client
|
||||
.create_jupyter_access_ticket(
|
||||
payload.edit_session_id,
|
||||
{
|
||||
"workspace_id": context.workspace.workspace_id,
|
||||
"user_id": context.user.user_id,
|
||||
"lock_token": payload.lock_token,
|
||||
},
|
||||
)
|
||||
)
|
||||
except RuntimeClientError as exc:
|
||||
error = exc.detail
|
||||
if not isinstance(error, dict) or "code" not in error:
|
||||
error = {
|
||||
"code": "JUPYTER_TICKET_FAILED",
|
||||
"message": str(error),
|
||||
"retryable": exc.status_code >= 500,
|
||||
"details": {},
|
||||
}
|
||||
return JSONResponse(
|
||||
status_code=exc.status_code,
|
||||
content={"request_id": context.request_id, "error": error},
|
||||
header_b64, payload_b64, signature_b64 = token.split(".", 2)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Invalid Authentication Token",
|
||||
) from exc
|
||||
|
||||
signing_input = f"{header_b64}.{payload_b64}".encode()
|
||||
expected = hmac.new(
|
||||
JWT_SECRET.encode(),
|
||||
signing_input,
|
||||
hashlib.sha256,
|
||||
).digest()
|
||||
try:
|
||||
signature = _b64decode(signature_b64)
|
||||
except Exception as exc: # pragma: no cover - malformed b64
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Invalid Authentication Token",
|
||||
) from exc
|
||||
if not hmac.compare_digest(expected, signature):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Invalid Authentication Token",
|
||||
)
|
||||
|
||||
raw_ticket = ticket_data.pop("ticket")
|
||||
max_age = int(ticket_data.pop("expires_in_seconds"))
|
||||
response = JSONResponse(
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
content={
|
||||
"request_id": context.request_id,
|
||||
"data": ticket_data,
|
||||
"meta": {},
|
||||
},
|
||||
try:
|
||||
payload = json.loads(_b64decode(payload_b64))
|
||||
except Exception as exc:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Invalid Authentication Token",
|
||||
) from exc
|
||||
|
||||
exp = payload.get("exp")
|
||||
if not isinstance(exp, (int, float)) or exp < time.time():
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Token expired",
|
||||
)
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def extract_notebook_path(
|
||||
uri: str,
|
||||
workspace_id: str,
|
||||
) -> Optional[str]:
|
||||
"""Pull the relative notebook path out of the original request URI.
|
||||
|
||||
Only ``/jupyter/{workspace_id}/notebooks/*.ipynb`` requests are
|
||||
subject to file-level lock checks; everything else (tree views,
|
||||
``/api/contents`` and WebSocket upgrades) bypasses the lock.
|
||||
"""
|
||||
pattern = rf"^/jupyter/{re.escape(workspace_id)}/notebooks/(.+\.ipynb)"
|
||||
match = re.match(pattern, uri)
|
||||
if match:
|
||||
return match.group(1)
|
||||
return None
|
||||
|
||||
|
||||
async def check_notebook_is_locked(
|
||||
session: AsyncSession,
|
||||
workspace_id: str,
|
||||
notebook_path: str,
|
||||
user_id: str,
|
||||
) -> bool:
|
||||
"""Decide whether the request is allowed to read the notebook.
|
||||
|
||||
Returns ``True`` only when the notebook exists, is owned by a
|
||||
different user, and is currently locked. The owner is always let
|
||||
through; a missing row is treated as "not owned yet" and allowed.
|
||||
"""
|
||||
statement = select(Scripts.owner_user_id, Scripts.is_locked).where(
|
||||
Scripts.workspace_id == workspace_id,
|
||||
Scripts.script_name == notebook_path,
|
||||
Scripts.script_type == "notebook",
|
||||
Scripts.is_deleted == 0,
|
||||
)
|
||||
response.set_cookie(
|
||||
key="jupyter_access",
|
||||
value=raw_ticket,
|
||||
max_age=max_age,
|
||||
path="/jupyter/",
|
||||
secure=cookie_secure(),
|
||||
httponly=True,
|
||||
samesite="lax",
|
||||
row = (await session.execute(statement)).one_or_none()
|
||||
if row is None:
|
||||
return False
|
||||
owner_user_id, is_locked = row
|
||||
if owner_user_id == user_id:
|
||||
return False
|
||||
return bool(is_locked)
|
||||
|
||||
|
||||
async def load_active_membership(
|
||||
session: AsyncSession,
|
||||
user_id: str,
|
||||
workspace_id: str,
|
||||
) -> tuple[Users, Workspaces]:
|
||||
"""Resolve the user's active membership in the workspace."""
|
||||
statement = (
|
||||
select(Users, Workspaces)
|
||||
.join(
|
||||
WorkspaceMembers,
|
||||
WorkspaceMembers.user_id == Users.user_id,
|
||||
)
|
||||
.join(
|
||||
Workspaces,
|
||||
Workspaces.workspace_id == WorkspaceMembers.workspace_id,
|
||||
)
|
||||
.where(
|
||||
Users.user_id == user_id,
|
||||
Users.status == "active",
|
||||
WorkspaceMembers.workspace_id == workspace_id,
|
||||
WorkspaceMembers.member_status == "active",
|
||||
Workspaces.status == "active",
|
||||
)
|
||||
)
|
||||
return response
|
||||
row = (await session.execute(statement)).one_or_none()
|
||||
if row is None:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="active workspace membership is required",
|
||||
)
|
||||
return row
|
||||
|
||||
|
||||
@router.get("/api/v1/auth/jupyter")
|
||||
async def verify_jupyter_access(
|
||||
request: Request,
|
||||
response: Response,
|
||||
auth: Optional[HTTPAuthorizationCredentials] = Depends(security),
|
||||
session: AsyncSession = Depends(database_session),
|
||||
) -> dict:
|
||||
"""Nginx auth_request subrequest handler.
|
||||
|
||||
Returns 200 with two response headers so Nginx can proxy the
|
||||
request to the user's Jupyter instance:
|
||||
|
||||
* ``x-upstream-addr`` = ``{base_url}:{port}``
|
||||
* ``x-jupyter-internal-token`` = the runtime-issued token
|
||||
"""
|
||||
workspace_id = request.headers.get("X-Original-Workspace-Id")
|
||||
original_uri = request.headers.get("X-Original-URI", "")
|
||||
|
||||
if not workspace_id:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Missing Workspace ID",
|
||||
)
|
||||
|
||||
cookie_token = request.cookies.get("access_token")
|
||||
bearer_token = auth.credentials if auth else None
|
||||
token = cookie_token or bearer_token
|
||||
payload = verify_jwt_token(token)
|
||||
user_id = payload.get("sub")
|
||||
if not user_id:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Invalid Authentication Token",
|
||||
)
|
||||
|
||||
await load_active_membership(session, user_id, workspace_id)
|
||||
|
||||
notebook_path = extract_notebook_path(original_uri, workspace_id)
|
||||
if notebook_path and await check_notebook_is_locked(
|
||||
session, workspace_id, notebook_path, user_id,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"Notebook '{notebook_path}' is currently locked",
|
||||
)
|
||||
|
||||
runtime_client = request.app.state.runtime_client
|
||||
ws_info = await runtime_client.get_workspace(workspace_id)
|
||||
if not ws_info or ws_info.get("status") != "running":
|
||||
try:
|
||||
ws_info = await runtime_client.start_workspace(workspace_id)
|
||||
except RuntimeClientError as exc:
|
||||
raise HTTPException(
|
||||
status_code=exc.status_code,
|
||||
detail=exc.detail,
|
||||
) from exc
|
||||
|
||||
target_port = ws_info.get("port")
|
||||
jupyter_token = ws_info.get("token")
|
||||
jupyter_base_url = ws_info.get("base_url")
|
||||
if not target_port:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail="Jupyter instance returned no port",
|
||||
)
|
||||
|
||||
response.headers["x-upstream-addr"] = f"{jupyter_base_url}:{target_port}"
|
||||
response.headers["x-jupyter-internal-token"] = jupyter_token or ""
|
||||
return {"status": "ok"}
|
||||
|
||||
@@ -1,147 +0,0 @@
|
||||
# coding=utf-8
|
||||
"""
|
||||
@Time :2026/7/27
|
||||
@Author :tao.chen
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
import httpx
|
||||
from fastapi import FastAPI, Request, Response, HTTPException, Depends, status
|
||||
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
|
||||
from typing import Optional
|
||||
from loguru import logger
|
||||
|
||||
app = FastAPI(title="Jupyter Auth & Router Backend")
|
||||
|
||||
RUNTIME_BASE_URL = os.getenv("RUNTIME_BASE_URL", "http://127.0.0.1:8001")
|
||||
security = HTTPBearer(auto_error=False)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 1. Runtime 交互 Client
|
||||
# ------------------------------------------------------------------
|
||||
class RuntimeClient:
|
||||
"""与 Runtime 进程管理器服务交互"""
|
||||
|
||||
@staticmethod
|
||||
async def get_workspace(workspace_id: str) -> Optional[dict]:
|
||||
"""按需查询单个 workspace 进程"""
|
||||
async with httpx.AsyncClient(base_url=RUNTIME_BASE_URL) as client:
|
||||
try:
|
||||
resp = await client.post(
|
||||
"/api/v1/jupyter",
|
||||
json={"action": "get", "workspace_id": workspace_id},
|
||||
timeout=3.0,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
return resp.json()
|
||||
return None
|
||||
except httpx.RequestError:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
async def start_workspace(workspace_id: str) -> dict:
|
||||
"""进程未运行时主动触发启动"""
|
||||
async with httpx.AsyncClient(base_url=RUNTIME_BASE_URL) as client:
|
||||
resp = await client.post(
|
||||
"/api/v1/jupyter",
|
||||
json={"action": "start", "workspace_id": workspace_id},
|
||||
timeout=10.0,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
return resp.json()
|
||||
raise HTTPException(
|
||||
status_code=500, detail="Failed to start Jupyter instance"
|
||||
)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 2. 数据库与权限模拟 (请根据实际 MySQL ORM 修改)
|
||||
# ------------------------------------------------------------------
|
||||
async def check_notebook_is_locked(workspace_id: str, notebook_path: str) -> bool:
|
||||
"""
|
||||
查数据库：判断特定 Notebook 文件是否被锁定
|
||||
:param workspace_id: 工作区 ID
|
||||
:param notebook_path: 相对路径，如 "test.ipynb" 或 "folder/demo.ipynb"
|
||||
"""
|
||||
# 模拟锁定数据库：假定 test_locked.ipynb 被锁定
|
||||
locked_notebooks = {
|
||||
("test1234", "test_locked.ipynb"): True,
|
||||
}
|
||||
return locked_notebooks.get((workspace_id, notebook_path), False)
|
||||
|
||||
|
||||
def verify_jwt_token(token: str) -> str:
|
||||
"""校验 JWT 令牌"""
|
||||
if token == "invalid-token":
|
||||
raise HTTPException(status_code=401, detail="Invalid Authentication Token")
|
||||
return "user_001"
|
||||
|
||||
|
||||
def extract_notebook_path(uri: str, workspace_id: str) -> Optional[str]:
|
||||
"""
|
||||
从原始请求 URI 中提取请求的 .ipynb 文件相对路径
|
||||
例如: /jupyter/test1234/notebooks/folder/test.ipynb -> folder/test.ipynb
|
||||
"""
|
||||
pattern = rf"^/jupyter/{re.escape(workspace_id)}/notebooks/(.+\.ipynb)"
|
||||
match = re.match(pattern, uri)
|
||||
if match:
|
||||
return match.group(1)
|
||||
return None
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 3. 核心 Auth 接口 (针对 Nginx auth_request)
|
||||
# ------------------------------------------------------------------
|
||||
@app.get("/api/v1/auth/jupyter")
|
||||
async def verify_jupyter_access(
|
||||
request: Request,
|
||||
response: Response,
|
||||
auth: Optional[HTTPAuthorizationCredentials] = Depends(security),
|
||||
):
|
||||
# 获取 Nginx 传入的元数据
|
||||
workspace_id = request.headers.get("X-Original-Workspace-Id")
|
||||
original_uri = request.headers.get("X-Original-URI", "")
|
||||
|
||||
cookie_token = request.cookies.get("access_token")
|
||||
bearer_token = auth.credentials if auth else None
|
||||
token = bearer_token or cookie_token
|
||||
|
||||
# if not token:
|
||||
# raise HTTPException(status_code=401, detail="Missing Authentication Token")
|
||||
|
||||
if not workspace_id:
|
||||
raise HTTPException(status_code=400, detail="Missing Workspace ID")
|
||||
|
||||
# 基础身份认证
|
||||
# current_user_id = verify_jwt_token(token)
|
||||
|
||||
# 精准锁校验：只有在访问 .ipynb 文件时才检查 is_locked
|
||||
notebook_path = extract_notebook_path(original_uri, workspace_id)
|
||||
if notebook_path:
|
||||
is_locked = await check_notebook_is_locked(workspace_id, notebook_path)
|
||||
if is_locked:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"Notebook '{notebook_path}' is currently locked",
|
||||
)
|
||||
|
||||
# 获取或启动 Jupyter 子进程
|
||||
ws_info = await RuntimeClient.get_workspace(workspace_id)
|
||||
|
||||
if not ws_info or ws_info.get("status") != "running":
|
||||
ws_info = await RuntimeClient.start_workspace(workspace_id)
|
||||
|
||||
target_port = ws_info.get("port")
|
||||
jupyter_token = ws_info.get("token")
|
||||
jupyter_base_url = ws_info.get("base_url")
|
||||
|
||||
if not target_port:
|
||||
raise HTTPException(
|
||||
status_code=500, detail="Jupyter instance returned no port"
|
||||
)
|
||||
|
||||
# 通过 Response Header 返回 Upstream 地址与 Token 给 Nginx
|
||||
response.headers["x-upstream-addr"] = f"{jupyter_base_url}:{target_port}"
|
||||
response.headers["x-jupyter-internal-token"] = jupyter_token or ""
|
||||
return {"status": "ok"}
|
||||
@@ -90,15 +90,34 @@ class RuntimeClient:
|
||||
)
|
||||
)["data"]
|
||||
|
||||
async def create_jupyter_access_ticket(
|
||||
async def get_workspace(
|
||||
self,
|
||||
edit_session_id: str,
|
||||
payload: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
return (
|
||||
await self._request(
|
||||
workspace_id: str,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Ask Runtime for a live workspace.
|
||||
|
||||
Returns the workspace descriptor when the Jupyter process is
|
||||
``running``; returns ``None`` when the workspace is not yet
|
||||
started or has been torn down so callers can fall through to
|
||||
:meth:`start_workspace`.
|
||||
"""
|
||||
try:
|
||||
return await self._request(
|
||||
"POST",
|
||||
f"/internal/v1/jupyter/access-tickets/{edit_session_id}",
|
||||
payload,
|
||||
"/api/v1/jupyter",
|
||||
{"action": "get", "workspace_id": workspace_id},
|
||||
)
|
||||
)["data"]
|
||||
except RuntimeClientError as exc:
|
||||
if exc.status_code == 404:
|
||||
return None
|
||||
raise
|
||||
|
||||
async def start_workspace(
|
||||
self,
|
||||
workspace_id: str,
|
||||
) -> dict[str, Any]:
|
||||
return await self._request(
|
||||
"POST",
|
||||
"/api/v1/jupyter",
|
||||
{"action": "start", "workspace_id": workspace_id},
|
||||
)
|
||||
|
||||
@@ -65,8 +65,3 @@ class DownloadUrlRequest(StrictModel):
|
||||
|
||||
class FileLockTokenRequest(StrictModel):
|
||||
lock_token: str = Field(min_length=32, max_length=256)
|
||||
|
||||
|
||||
class CreateJupyterAccessTicketRequest(StrictModel):
|
||||
edit_session_id: str = Field(min_length=26, max_length=26)
|
||||
lock_token: str = Field(min_length=32, max_length=256)
|
||||
|
||||
Reference in New Issue
Block a user