118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
"""
|
||
Copyright 2024-2026 ReplyPier
|
||
|
||
Licensed under the Apache License, Version 2.0 (the "License");
|
||
you may not use this file except in compliance with the License.
|
||
You may obtain a copy of the License at
|
||
|
||
http://www.apache.org/licenses/LICENSE-2.0
|
||
|
||
Unless required by applicable law or agreed to in writing, software
|
||
distributed under the License is distributed on an "AS IS" BASIS,
|
||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
See the License for the specific language governing permissions and
|
||
limitations under the License.
|
||
|
||
Hardening for user-uploaded images (help-center logo + article images). Every
|
||
uploaded raster is:
|
||
- size-capped (bytes),
|
||
- checked for a real, decodable image (not a polyglot / renamed file),
|
||
- guarded against decompression bombs (pixel-count cap read from the header
|
||
BEFORE the full image is decoded),
|
||
- re-encoded from decoded pixels — which strips EXIF/metadata, trailing bytes
|
||
and any appended payload, and normalises the stored format.
|
||
|
||
SVG is intentionally NOT accepted here: it is active markup and would be stored
|
||
XSS on the public help center. Callers that previously took SVG rasterise on the
|
||
client (the cropper) or drop it.
|
||
"""
|
||
|
||
import io
|
||
|
||
from fastapi import HTTPException
|
||
from PIL import Image
|
||
|
||
# Refuse absurd pixel counts before decoding — a few-KB file can claim to be
|
||
# tens of thousands of pixels per side and exhaust memory on decode.
|
||
_MAX_PIXELS = 40_000_000 # 40 MP
|
||
Image.MAX_IMAGE_PIXELS = _MAX_PIXELS
|
||
|
||
# PIL format -> (content type, extension) for the formats we re-encode to.
|
||
_PNG = ("image/png", ".png")
|
||
_JPEG = ("image/jpeg", ".jpg")
|
||
|
||
_FORMAT_CONTENT_TYPE = {
|
||
"PNG": "image/png",
|
||
"JPEG": "image/jpeg",
|
||
"WEBP": "image/webp",
|
||
"GIF": "image/gif",
|
||
}
|
||
|
||
|
||
def _open_validated(content: bytes, *, max_bytes: int, max_dim: int) -> Image.Image:
|
||
"""Open bytes as an image, enforcing byte size, decodability and a pixel /
|
||
dimension cap (bomb guard) — all before the pixels are fully decoded."""
|
||
if not content:
|
||
raise HTTPException(status_code=400, detail="Empty image file.")
|
||
if len(content) > max_bytes:
|
||
raise HTTPException(status_code=400, detail=f"Image must be {max_bytes // (1024 * 1024)} MB or smaller.")
|
||
try:
|
||
img = Image.open(io.BytesIO(content))
|
||
width, height = img.size # header only; no full decode yet
|
||
except HTTPException:
|
||
raise
|
||
except Exception:
|
||
raise HTTPException(status_code=400, detail="Invalid or corrupt image file.")
|
||
if width <= 0 or height <= 0:
|
||
raise HTTPException(status_code=400, detail="Invalid image dimensions.")
|
||
if width > max_dim or height > max_dim or width * height > _MAX_PIXELS:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Image is too large — keep it within {max_dim}×{max_dim} pixels.",
|
||
)
|
||
try:
|
||
img.load() # dimensions are now known-safe, so decoding is bounded
|
||
except Exception:
|
||
raise HTTPException(status_code=400, detail="Invalid or corrupt image file.")
|
||
return img
|
||
|
||
|
||
def sanitize_image(
|
||
content: bytes,
|
||
*,
|
||
allowed_content_types: set[str],
|
||
max_bytes: int,
|
||
max_dim: int,
|
||
fit: int | None = None,
|
||
) -> tuple[bytes, str, str]:
|
||
"""Validate and re-encode an uploaded raster image.
|
||
|
||
Returns (safe_bytes, content_type, extension). Raises HTTPException(400) on
|
||
anything that is not a clean, in-policy raster image. `fit`, when given,
|
||
downscales the longest side to at most `fit` px (aspect preserved).
|
||
|
||
JPEG re-encodes to JPEG (keeps photos small); everything else re-encodes to
|
||
PNG (keeps transparency). Either path strips metadata and any non-pixel
|
||
payload because the output is built from decoded pixels only.
|
||
"""
|
||
img = _open_validated(content, max_bytes=max_bytes, max_dim=max_dim)
|
||
|
||
detected = _FORMAT_CONTENT_TYPE.get((img.format or "").upper())
|
||
if detected is None or detected not in allowed_content_types:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="Unsupported image type. Use PNG, JPEG or WebP.",
|
||
)
|
||
|
||
if fit and max(img.size) > fit:
|
||
img.thumbnail((fit, fit), Image.LANCZOS)
|
||
|
||
out = io.BytesIO()
|
||
if detected == "image/jpeg":
|
||
img.convert("RGB").save(out, format="JPEG", quality=85, optimize=True)
|
||
content_type, ext = _JPEG
|
||
else:
|
||
img.convert("RGBA").save(out, format="PNG", optimize=True)
|
||
content_type, ext = _PNG
|
||
return out.getvalue(), content_type, ext
|