fix(security): harden uploads and wecom binding
This commit is contained in:
+41
-20
@@ -1,12 +1,15 @@
|
||||
import uuid
|
||||
import io
|
||||
import warnings
|
||||
from datetime import date
|
||||
from pathlib import PurePosixPath
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.database import get_db
|
||||
from app.middleware.auth import get_current_user, require_any_role
|
||||
from app.services.minio_client import generate_presigned_upload_url, generate_presigned_download_url
|
||||
from app.services.minio_client import generate_presigned_download_url, put_validated_image
|
||||
from app.models.visit import Visit
|
||||
|
||||
router = APIRouter(prefix="/upload", tags=["Upload"])
|
||||
@@ -16,16 +19,17 @@ _IMAGE_TYPES = {
|
||||
"image/png": "png",
|
||||
"image/webp": "webp",
|
||||
}
|
||||
MAX_IMAGE_BYTES = 10 * 1024 * 1024
|
||||
MAX_IMAGE_PIXELS = 25_000_000
|
||||
|
||||
|
||||
def validate_image_upload(filename: str, content_type: str) -> str:
|
||||
"""Validate client-declared image metadata before issuing a short-lived PUT URL."""
|
||||
extension = _IMAGE_TYPES.get(content_type.lower())
|
||||
if not extension:
|
||||
"""Validate image metadata before accepting bytes for server-side verification."""
|
||||
if content_type.lower() not in _IMAGE_TYPES:
|
||||
raise HTTPException(status_code=400, detail="Only JPEG, PNG, and WebP images are supported")
|
||||
if PurePosixPath(filename).suffix.lower() not in {".jpg", ".jpeg", ".png", ".webp"}:
|
||||
raise HTTPException(status_code=400, detail="Invalid image filename extension")
|
||||
return extension
|
||||
return _IMAGE_TYPES[content_type.lower()]
|
||||
|
||||
|
||||
def build_object_key(user_id: str, extension: str) -> str:
|
||||
@@ -44,22 +48,39 @@ def is_owned_upload_key(object_key: str, user_id: str) -> bool:
|
||||
)
|
||||
|
||||
|
||||
@router.post("/presigned-url")
|
||||
async def get_presigned_upload_url(
|
||||
filename: str,
|
||||
content_type: str = "image/jpeg",
|
||||
def normalize_image(image_bytes: bytes) -> bytes:
|
||||
"""Verify image bytes, enforce a pixel limit, and strip active/metadata content by re-encoding."""
|
||||
if not image_bytes or len(image_bytes) > MAX_IMAGE_BYTES:
|
||||
raise HTTPException(status_code=413, detail="Image must be between 1 byte and 10 MB")
|
||||
try:
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("error", Image.DecompressionBombWarning)
|
||||
with Image.open(io.BytesIO(image_bytes)) as image:
|
||||
image.verify()
|
||||
with Image.open(io.BytesIO(image_bytes)) as image:
|
||||
if image.width * image.height > MAX_IMAGE_PIXELS:
|
||||
raise HTTPException(status_code=413, detail="Image dimensions are too large")
|
||||
normalized = image.convert("RGB")
|
||||
output = io.BytesIO()
|
||||
normalized.save(output, format="JPEG", quality=85, optimize=True)
|
||||
return output.getvalue()
|
||||
except HTTPException:
|
||||
raise
|
||||
except (Image.DecompressionBombError, Image.DecompressionBombWarning, OSError, UnidentifiedImageError, ValueError) as exc:
|
||||
raise HTTPException(status_code=400, detail="Invalid image content") from exc
|
||||
|
||||
|
||||
@router.post("/image")
|
||||
async def upload_image(
|
||||
file: UploadFile = File(...),
|
||||
current_user: dict = Depends(require_any_role),
|
||||
):
|
||||
"""Get a presigned PUT URL for direct MinIO upload."""
|
||||
extension = validate_image_upload(filename, content_type)
|
||||
object_key = build_object_key(current_user["user_id"], extension)
|
||||
|
||||
url = generate_presigned_upload_url(object_key)
|
||||
|
||||
return {
|
||||
"upload_url": url,
|
||||
"object_key": object_key,
|
||||
}
|
||||
"""Validate and normalize an image before persisting it to private object storage."""
|
||||
validate_image_upload(file.filename or "upload.jpg", file.content_type or "")
|
||||
normalized = normalize_image(await file.read(MAX_IMAGE_BYTES + 1))
|
||||
object_key = build_object_key(current_user["user_id"], "jpg")
|
||||
put_validated_image(object_key, normalized)
|
||||
return {"object_key": object_key}
|
||||
|
||||
|
||||
@router.get("/download-url")
|
||||
|
||||
+13
-18
@@ -9,8 +9,9 @@ from fastapi.responses import PlainTextResponse
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from app.config import settings
|
||||
from app.database import get_db
|
||||
from app.database import async_session, get_db
|
||||
from app.middleware.auth import get_current_user, require_director
|
||||
from app.models.user import User
|
||||
from app.services.wecom import wecom_client, store_bind_token, consume_bind_token
|
||||
@@ -116,21 +117,20 @@ async def wecom_callback_event(request: Request):
|
||||
# Triggers for binding
|
||||
if (msg_type == "text" and content == "绑定") or (msg_type == "event" and event == "click" and event_key == "BIND_ACCOUNT"):
|
||||
_log.info("wecom bind triggered")
|
||||
_handle_bind_request(from_user)
|
||||
await _handle_bind_request(from_user)
|
||||
return PlainTextResponse(content="success")
|
||||
|
||||
_log.info("wecom message ignored (no matching action)")
|
||||
return PlainTextResponse(content="success")
|
||||
|
||||
|
||||
def _handle_bind_request(wecom_userid: str):
|
||||
async def _handle_bind_request(wecom_userid: str):
|
||||
"""Generate bind token, store it, and push a bind link to the user."""
|
||||
if not wecom_userid:
|
||||
return
|
||||
|
||||
import asyncio
|
||||
|
||||
bind_token = store_bind_token(wecom_userid)
|
||||
async with async_session() as db:
|
||||
bind_token = await store_bind_token(db, wecom_userid)
|
||||
|
||||
content = (
|
||||
f"【账号绑定】\n\n"
|
||||
@@ -139,16 +139,7 @@ def _handle_bind_request(wecom_userid: str):
|
||||
f"绑定后可使用企微一键登录,并接收填报提醒通知。"
|
||||
)
|
||||
|
||||
try:
|
||||
# Must run async in sync context
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
loop.run_until_complete(
|
||||
wecom_client.send_text_message([wecom_userid], content)
|
||||
)
|
||||
await wecom_client.send_text_message([wecom_userid], content)
|
||||
|
||||
|
||||
# ── Bind confirmation (JWT-protected) ──
|
||||
@@ -163,7 +154,7 @@ async def bind_confirm(
|
||||
if not data.token:
|
||||
raise HTTPException(status_code=400, detail="token 不能为空")
|
||||
|
||||
wecom_userid = consume_bind_token(data.token)
|
||||
wecom_userid = await consume_bind_token(db, data.token)
|
||||
if not wecom_userid:
|
||||
raise HTTPException(status_code=400, detail="绑定链接已过期或无效,请重新在企微发送「绑定」")
|
||||
|
||||
@@ -181,7 +172,11 @@ async def bind_confirm(
|
||||
)
|
||||
user = result.scalar_one()
|
||||
user.wecom_userid = wecom_userid
|
||||
await db.commit()
|
||||
try:
|
||||
await db.commit()
|
||||
except IntegrityError as exc:
|
||||
await db.rollback()
|
||||
raise HTTPException(status_code=409, detail="该企业微信已绑定其他账号") from exc
|
||||
|
||||
return {"code": 200, "msg": "绑定成功", "wecom_userid": wecom_userid}
|
||||
|
||||
|
||||
@@ -51,6 +51,10 @@ async def lifespan(app: FastAPI):
|
||||
await conn.run_sync(lambda c: c.exec_driver_sql(
|
||||
"ALTER TABLE users ADD COLUMN IF NOT EXISTS color VARCHAR(7)"
|
||||
))
|
||||
await conn.run_sync(lambda c: c.exec_driver_sql(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS uq_users_wecom_userid "
|
||||
"ON users (wecom_userid) WHERE wecom_userid IS NOT NULL"
|
||||
))
|
||||
# system_config table for v0.5
|
||||
await conn.run_sync(lambda c: c.exec_driver_sql(
|
||||
"CREATE TABLE IF NOT EXISTS system_config (key VARCHAR(64) PRIMARY KEY, value TEXT DEFAULT '')"
|
||||
|
||||
@@ -10,6 +10,7 @@ from app.models.daily_note import DailyNote
|
||||
from app.models.ai_summary import AISummary
|
||||
from app.models.system_config import SystemConfig
|
||||
from app.models.leave import Leave
|
||||
from app.models.wecom_bind_token import WecomBindToken
|
||||
|
||||
__all__ = [
|
||||
"User",
|
||||
@@ -24,4 +25,5 @@ __all__ = [
|
||||
"AISummary",
|
||||
"SystemConfig",
|
||||
"Leave",
|
||||
"WecomBindToken",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import DateTime, String, func
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class WecomBindToken(Base):
|
||||
"""A single-use, expiring binding token issued after a verified WeCom callback."""
|
||||
|
||||
__tablename__ = "wecom_bind_tokens"
|
||||
|
||||
token: Mapped[str] = mapped_column(String(64), primary_key=True)
|
||||
wecom_userid: Mapped[str] = mapped_column(String(100), index=True)
|
||||
expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True)
|
||||
consumed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
|
||||
@@ -1,4 +1,5 @@
|
||||
from datetime import timedelta
|
||||
from io import BytesIO
|
||||
from minio import Minio
|
||||
from app.config import settings
|
||||
|
||||
@@ -20,10 +21,16 @@ def get_minio_client() -> Minio:
|
||||
return _client
|
||||
|
||||
|
||||
def generate_presigned_upload_url(object_key: str, expires: int = 600) -> str:
|
||||
"""Generate a presigned PUT URL for direct upload to MinIO."""
|
||||
def put_validated_image(object_key: str, image_bytes: bytes) -> None:
|
||||
"""Store a normalized image only after the API has validated its content."""
|
||||
client = get_minio_client()
|
||||
return client.presigned_put_object(settings.MINIO_BUCKET, object_key, expires=timedelta(seconds=expires))
|
||||
client.put_object(
|
||||
settings.MINIO_BUCKET,
|
||||
object_key,
|
||||
BytesIO(image_bytes),
|
||||
length=len(image_bytes),
|
||||
content_type="image/jpeg",
|
||||
)
|
||||
|
||||
|
||||
def generate_presigned_download_url(object_key: str, expires: int = 3600) -> str:
|
||||
|
||||
@@ -1,30 +1,41 @@
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime, timedelta, timezone
|
||||
import httpx
|
||||
from sqlalchemy import delete, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.config import settings
|
||||
from app.models.wecom_bind_token import WecomBindToken
|
||||
|
||||
# In-memory bind token store (TTL 600s). Replace with Redis if scaling to multiple workers.
|
||||
_bind_tokens: dict[str, tuple[str, float]] = {} # token → (wecom_userid, expires_at)
|
||||
|
||||
|
||||
def store_bind_token(wecom_userid: str, ttl: int = 600) -> str:
|
||||
"""Store a bind token → wecom_userid mapping. Returns the token."""
|
||||
async def store_bind_token(db: AsyncSession, wecom_userid: str, ttl: int = 600) -> str:
|
||||
"""Persist a one-time bind token so it works across application workers."""
|
||||
token = uuid.uuid4().hex
|
||||
_bind_tokens[token] = (wecom_userid, time.time() + ttl)
|
||||
# Cleanup expired tokens
|
||||
now = time.time()
|
||||
for k in list(_bind_tokens):
|
||||
if _bind_tokens[k][1] < now:
|
||||
del _bind_tokens[k]
|
||||
now = datetime.now(timezone.utc)
|
||||
await db.execute(delete(WecomBindToken).where(WecomBindToken.expires_at < now))
|
||||
db.add(WecomBindToken(
|
||||
token=token,
|
||||
wecom_userid=wecom_userid,
|
||||
expires_at=now + timedelta(seconds=ttl),
|
||||
))
|
||||
await db.commit()
|
||||
return token
|
||||
|
||||
|
||||
def consume_bind_token(token: str) -> str | None:
|
||||
"""Lookup and consume a bind token. Returns wecom_userid or None."""
|
||||
entry = _bind_tokens.pop(token, None)
|
||||
if entry and entry[1] > time.time():
|
||||
return entry[0]
|
||||
return None
|
||||
async def consume_bind_token(db: AsyncSession, token: str) -> str | None:
|
||||
"""Atomically consume an unexpired binding token and return its WeCom user ID."""
|
||||
now = datetime.now(timezone.utc)
|
||||
result = await db.execute(
|
||||
update(WecomBindToken)
|
||||
.where(
|
||||
WecomBindToken.token == token,
|
||||
WecomBindToken.consumed_at.is_(None),
|
||||
WecomBindToken.expires_at > now,
|
||||
)
|
||||
.values(consumed_at=now)
|
||||
.returning(WecomBindToken.wecom_userid)
|
||||
)
|
||||
await db.commit()
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
class WecomClient:
|
||||
|
||||
@@ -13,3 +13,4 @@ openpyxl==3.1.5
|
||||
apscheduler==3.11.0
|
||||
python-dotenv==1.0.1
|
||||
pycryptodome
|
||||
Pillow==11.1.0
|
||||
|
||||
@@ -1,7 +1,13 @@
|
||||
from fastapi import HTTPException
|
||||
import asyncio
|
||||
from io import BytesIO
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from app.api.upload import build_object_key, is_owned_upload_key, validate_image_upload
|
||||
from fastapi import HTTPException
|
||||
from PIL import Image
|
||||
|
||||
from app.api.upload import build_object_key, is_owned_upload_key, normalize_image, validate_image_upload
|
||||
from app.config import settings, validate_security_settings
|
||||
from app.services.wecom import consume_bind_token, store_bind_token
|
||||
|
||||
|
||||
USER_ID = "a2d4447f-7b67-4304-ad9a-953007e75ef2"
|
||||
@@ -32,6 +38,46 @@ def test_upload_rejects_an_invalid_filename_extension():
|
||||
raise AssertionError("invalid extensions must be rejected")
|
||||
|
||||
|
||||
def test_image_content_is_normalized_to_jpeg():
|
||||
source = BytesIO()
|
||||
Image.new("RGBA", (8, 8), color=(255, 0, 0, 128)).save(source, format="PNG")
|
||||
|
||||
normalized = normalize_image(source.getvalue())
|
||||
|
||||
with Image.open(BytesIO(normalized)) as image:
|
||||
assert image.format == "JPEG"
|
||||
assert image.mode == "RGB"
|
||||
|
||||
|
||||
def test_image_content_rejects_non_images():
|
||||
try:
|
||||
normalize_image(b"not an image")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 400
|
||||
else:
|
||||
raise AssertionError("non-image bytes must be rejected")
|
||||
|
||||
|
||||
def test_wecom_binding_tokens_are_persisted_and_consumed_atomically():
|
||||
session = MagicMock()
|
||||
session.execute = AsyncMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
token = asyncio.run(store_bind_token(session, "wecom-user"))
|
||||
|
||||
assert len(token) == 32
|
||||
assert session.add.call_args.args[0].wecom_userid == "wecom-user"
|
||||
session.commit.assert_awaited_once()
|
||||
|
||||
result = MagicMock()
|
||||
result.scalar_one_or_none.return_value = "wecom-user"
|
||||
session.execute = AsyncMock(return_value=result)
|
||||
session.commit = AsyncMock()
|
||||
|
||||
assert asyncio.run(consume_bind_token(session, token)) == "wecom-user"
|
||||
session.commit.assert_awaited_once()
|
||||
|
||||
|
||||
def test_production_rejects_placeholder_database_configuration(monkeypatch):
|
||||
monkeypatch.setattr(settings, "ENVIRONMENT", "production")
|
||||
monkeypatch.setattr(settings, "SECRET_KEY", "a" * 32)
|
||||
|
||||
Reference in New Issue
Block a user