Spaces:

devmandan
/

qr-scan-difficulty-check

Running

qr-scan-difficulty-check / server /main.py

Second-round review fixes

6c26597 verified 3 days ago

6.33 kB

"""QR Scan Difficulty Check: the measured bell, and a place on it for your own pictures.
GET  /                 the page (web/)
POST /api/measure      one picture (JPEG, PNG or WebP, at most 25 MB) as the request body; returns its reads by
                       the five scanners at the seven frame sizes and its place on the scale
GET  /api/health       what is running
Pictures are read from a temporary file that is deleted when the reading ends; nothing is stored, and no access log is
kept.
"""

import asyncio
import os
import tempfile
from contextlib import asynccontextmanager
from importlib.metadata import version
from pathlib import Path

from starlette.applications import Starlette
from starlette.concurrency import run_in_threadpool
from starlette.requests import ClientDisconnect, Request
from starlette.responses import JSONResponse
from starlette.routing import Mount, Route
from starlette.staticfiles import StaticFiles

from .place import place
from .reading import FRAMES, SCANNERS, PictureError, Reader

WEB = Path(__file__).resolve().parent.parent / "web"
MAX_BYTES = 25 * 1024 * 1024
SLOTS = int(os.environ.get("QSD_SLOTS", "2"))   # pictures read at once: one decoder worker each
MAX_WAITING = 16      # pictures waiting beyond those (arriving, or waiting for a slot); more are turned away
SLOT_WAIT_S = 10      # the longest a picture waits for a slot, so that a request ends within a minute
BODY_S = 60           # the longest a picture may take to arrive; a stalled upload doesn't keep its place
BUSY = "Many pictures are being read right now. Try again in a minute."
CSP = "; ".join([
    "default-src 'self'", "img-src 'self' blob: data:", "style-src 'self' 'unsafe-inline'", "script-src 'self'",
    "connect-src 'self'", "object-src 'none'", "base-uri 'none'", "form-action 'none'",
    "frame-ancestors 'self' https://huggingface.co",
])
STATE: dict = {}


def sniff(data: bytes) -> str | None:
    """The picture's format from its first bytes: only JPEG, PNG and WebP are opened here (the page converts others)."""
    if data[:3] == b"\xff\xd8\xff":
        return ".jpg"
    if data[:8] == b"\x89PNG\r\n\x1a\n":
        return ".png"
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return ".webp"
    return None


def problem(status: int, message: str) -> JSONResponse:
    return JSONResponse({"error": message}, status_code=status, headers={"Cache-Control": "no-store"})


async def body(request: Request) -> bytes | None:
    """The request body, or None past MAX_BYTES (read no further)."""
    n = request.headers.get("content-length", "")
    if n.isdigit() and int(n) > MAX_BYTES:
        return None
    buf = bytearray()   # one buffer, so a picture costs its own size in memory however small its chunks
    async for chunk in request.stream():
        buf += chunk
        if len(buf) > MAX_BYTES:
            return None
    return bytes(buf)


def read_picture(data: bytes, suffix: str) -> dict:
    with tempfile.NamedTemporaryFile(prefix="qsd-", suffix=suffix) as f:
        f.write(data)
        f.flush()
        result = STATE["reader"].read(f.name)
    result["order"] = FRAMES   # the frame size behind each position of "reads" and "frames"
    result["place"] = place(result["reads"], FRAMES)
    result["count1280"] = sum(result["reads"][k][FRAMES.index(1280)] for k in SCANNERS)
    return result


async def measure(request: Request) -> JSONResponse:
    if STATE["busy"] >= SLOTS + MAX_WAITING:
        return problem(503, BUSY)
    STATE["busy"] += 1   # counted from the first byte, so at most SLOTS + MAX_WAITING pictures are held in memory
    try:
        try:
            async with asyncio.timeout(BODY_S):
                data = await body(request)
        except TimeoutError:
            return problem(408, "This picture took too long to arrive. Try again, or send a smaller copy.")
        except ClientDisconnect:   # the visitor left mid-upload: nobody will read this answer
            return problem(400, "The picture did not arrive.")
        if data is None:
            return problem(413, "This picture is over 25 MB.")
        suffix = sniff(data)
        if suffix is None:
            return problem(415, "Send a JPEG, PNG or WebP picture.")
        try:
            await asyncio.wait_for(STATE["gate"].acquire(), SLOT_WAIT_S)
        except TimeoutError:
            return problem(503, BUSY)
        try:
            result = await run_in_threadpool(read_picture, data, suffix)
        finally:
            STATE["gate"].release()
    except PictureError as e:
        return problem(422, str(e))
    finally:
        STATE["busy"] -= 1
    return JSONResponse(result, headers={"Cache-Control": "no-store"})


async def health(request: Request) -> JSONResponse:
    return JSONResponse({
        "ok": True, "scanners": SCANNERS, "frames": FRAMES, "slots": SLOTS,
        "versions": {"opencv-contrib": version("opencv-contrib-python-headless"), "zxing-cpp": version("zxing-cpp"),
                     "zxing": "3.5.2"},
    }, headers={"Cache-Control": "no-store"})


@asynccontextmanager
async def lifespan(app):
    STATE["reader"] = await run_in_threadpool(Reader, SLOTS)
    STATE["gate"] = asyncio.Semaphore(SLOTS)
    STATE["busy"] = 0
    yield
    STATE["reader"].close()


class SecurityHeaders:
    """Plain ASGI middleware: the same headers on every response."""

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            return await self.app(scope, receive, send)

        async def send_with_headers(message):
            if message["type"] == "http.response.start":
                message.setdefault("headers", [])
                message["headers"] += [(b"content-security-policy", CSP.encode()),
                                       (b"x-content-type-options", b"nosniff"),
                                       (b"referrer-policy", b"strict-origin-when-cross-origin")]
            await send(message)

        return await self.app(scope, receive, send_with_headers)


app = SecurityHeaders(Starlette(lifespan=lifespan, routes=[
    Route("/api/measure", measure, methods=["POST"]),
    Route("/api/health", health, methods=["GET"]),
    Mount("/", StaticFiles(directory=WEB, html=True), name="web"),
]))

	"""QR Scan Difficulty Check: the measured bell, and a place on it for your own pictures.

	GET / the page (web/)
	POST /api/measure one picture (JPEG, PNG or WebP, at most 25 MB) as the request body; returns its reads by
	the five scanners at the seven frame sizes and its place on the scale
	GET /api/health what is running
	Pictures are read from a temporary file that is deleted when the reading ends; nothing is stored, and no access log is
	kept.
	"""

	import asyncio
	import os
	import tempfile
	from contextlib import asynccontextmanager
	from importlib.metadata import version
	from pathlib import Path

	from starlette.applications import Starlette
	from starlette.concurrency import run_in_threadpool
	from starlette.requests import ClientDisconnect, Request
	from starlette.responses import JSONResponse
	from starlette.routing import Mount, Route
	from starlette.staticfiles import StaticFiles

	from .place import place
	from .reading import FRAMES, SCANNERS, PictureError, Reader

	WEB = Path(__file__).resolve().parent.parent / "web"
	MAX_BYTES = 25 * 1024 * 1024
	SLOTS = int(os.environ.get("QSD_SLOTS", "2")) # pictures read at once: one decoder worker each
	MAX_WAITING = 16 # pictures waiting beyond those (arriving, or waiting for a slot); more are turned away
	SLOT_WAIT_S = 10 # the longest a picture waits for a slot, so that a request ends within a minute
	BODY_S = 60 # the longest a picture may take to arrive; a stalled upload doesn't keep its place
	BUSY = "Many pictures are being read right now. Try again in a minute."
	CSP = "; ".join([
	"default-src 'self'", "img-src 'self' blob: data:", "style-src 'self' 'unsafe-inline'", "script-src 'self'",
	"connect-src 'self'", "object-src 'none'", "base-uri 'none'", "form-action 'none'",
	"frame-ancestors 'self' https://huggingface.co",
	])
	STATE: dict = {}


	def sniff(data: bytes) -> str \| None:
	"""The picture's format from its first bytes: only JPEG, PNG and WebP are opened here (the page converts others)."""
	if data[:3] == b"\xff\xd8\xff":
	return ".jpg"
	if data[:8] == b"\x89PNG\r\n\x1a\n":
	return ".png"
	if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
	return ".webp"
	return None


	def problem(status: int, message: str) -> JSONResponse:
	return JSONResponse({"error": message}, status_code=status, headers={"Cache-Control": "no-store"})


	async def body(request: Request) -> bytes \| None:
	"""The request body, or None past MAX_BYTES (read no further)."""
	n = request.headers.get("content-length", "")
	if n.isdigit() and int(n) > MAX_BYTES:
	return None
	buf = bytearray() # one buffer, so a picture costs its own size in memory however small its chunks
	async for chunk in request.stream():
	buf += chunk
	if len(buf) > MAX_BYTES:
	return None
	return bytes(buf)


	def read_picture(data: bytes, suffix: str) -> dict:
	with tempfile.NamedTemporaryFile(prefix="qsd-", suffix=suffix) as f:
	f.write(data)
	f.flush()
	result = STATE["reader"].read(f.name)
	result["order"] = FRAMES # the frame size behind each position of "reads" and "frames"
	result["place"] = place(result["reads"], FRAMES)
	result["count1280"] = sum(result["reads"][k][FRAMES.index(1280)] for k in SCANNERS)
	return result


	async def measure(request: Request) -> JSONResponse:
	if STATE["busy"] >= SLOTS + MAX_WAITING:
	return problem(503, BUSY)
	STATE["busy"] += 1 # counted from the first byte, so at most SLOTS + MAX_WAITING pictures are held in memory
	try:
	try:
	async with asyncio.timeout(BODY_S):
	data = await body(request)
	except TimeoutError:
	return problem(408, "This picture took too long to arrive. Try again, or send a smaller copy.")
	except ClientDisconnect: # the visitor left mid-upload: nobody will read this answer
	return problem(400, "The picture did not arrive.")
	if data is None:
	return problem(413, "This picture is over 25 MB.")
	suffix = sniff(data)
	if suffix is None:
	return problem(415, "Send a JPEG, PNG or WebP picture.")
	try:
	await asyncio.wait_for(STATE["gate"].acquire(), SLOT_WAIT_S)
	except TimeoutError:
	return problem(503, BUSY)
	try:
	result = await run_in_threadpool(read_picture, data, suffix)
	finally:
	STATE["gate"].release()
	except PictureError as e:
	return problem(422, str(e))
	finally:
	STATE["busy"] -= 1
	return JSONResponse(result, headers={"Cache-Control": "no-store"})


	async def health(request: Request) -> JSONResponse:
	return JSONResponse({
	"ok": True, "scanners": SCANNERS, "frames": FRAMES, "slots": SLOTS,
	"versions": {"opencv-contrib": version("opencv-contrib-python-headless"), "zxing-cpp": version("zxing-cpp"),
	"zxing": "3.5.2"},
	}, headers={"Cache-Control": "no-store"})


	@asynccontextmanager
	async def lifespan(app):
	STATE["reader"] = await run_in_threadpool(Reader, SLOTS)
	STATE["gate"] = asyncio.Semaphore(SLOTS)
	STATE["busy"] = 0
	yield
	STATE["reader"].close()


	class SecurityHeaders:
	"""Plain ASGI middleware: the same headers on every response."""

	def __init__(self, app):
	self.app = app

	async def __call__(self, scope, receive, send):
	if scope["type"] != "http":
	return await self.app(scope, receive, send)

	async def send_with_headers(message):
	if message["type"] == "http.response.start":
	message.setdefault("headers", [])
	message["headers"] += [(b"content-security-policy", CSP.encode()),
	(b"x-content-type-options", b"nosniff"),
	(b"referrer-policy", b"strict-origin-when-cross-origin")]
	await send(message)

	return await self.app(scope, receive, send_with_headers)


	app = SecurityHeaders(Starlette(lifespan=lifespan, routes=[
	Route("/api/measure", measure, methods=["POST"]),
	Route("/api/health", health, methods=["GET"]),
	Mount("/", StaticFiles(directory=WEB, html=True), name="web"),
	]))