app/backend/fluksio/api/routes/artifacts.py

"""Artifacts over HTTP: the one way bytes get in and out of the store.

A node on this host could reach the directory itself, but a node on a remote
worker cannot — and having one path rather than two is what keeps a flow's
code the same wherever it runs.
"""

import re
import tempfile
from pathlib import Path
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
from fastapi.responses import FileResponse
from jwt.exceptions import InvalidTokenError
from pydantic import BaseModel
from sqlmodel import Session
from starlette.concurrency import run_in_threadpool

from fluksio.api.deps import (
    FlowControllerDep,
    get_current_active_superuser,
    user_from_token,
    worker_claims,
)
from fluksio.core.config import settings
from fluksio.core.db import engine
from fluksio.flow.artifacts import ArtifactStore
from fluksio.flow.runs import referenced_artifacts

#: What may be echoed back as a response content type. The caller holds the
#: reference and passes its media type, so this guards a header rather than
#: trusting one — anything else is served as bytes.
_MEDIA_TYPE = re.compile(r"^[\w.+-]+/[\w.+-]+$")

#: What a browser shows without running anything, so it is served as it is.
#: Every other type is served sandboxed: HTML, SVG and XML run script when
#: opened as a page, and so may a type nobody listed here.
_INERT = re.compile(r"^(audio|video)/|^image/(?!svg)|^application/(pdf|json)$")


def artifact_caller(request: Request) -> str:
    """Who may move artifacts: a signed-in person, or an attached worker.

    A worker's node stores its checkpoints through this endpoint, so its own
    credential has to open it — and only it. The token is no use anywhere else
    in the API, which is why this check is here rather than in the shared
    dependency every other route uses.
    """
    header = request.headers.get("Authorization", "")
    token = header[7:] if header.lower().startswith("bearer ") else ""
    if not token:
        raise HTTPException(status_code=401, detail="Not authenticated")
    try:
        claims = worker_claims(token)
    except InvalidTokenError:
        pass
    else:
        return f"worker:{claims.get('sub')}"
    with Session(engine) as session:
        # With the request, so a credential that is scoped by route — a wall
        # panel's — is judged against this one rather than waved through.
        user = user_from_token(session, token, request)
    if user is None:
        raise HTTPException(status_code=401, detail="Not authenticated")
    return user.email


router = APIRouter(
    prefix="/artifacts", tags=["artifacts"], dependencies=[Depends(artifact_caller)]
)


class ArtifactRef(BaseModel):
    digest: str
    size: int
    media_type: str
    name: str = ""


def _store(request: Request) -> ArtifactStore:
    store: ArtifactStore | None = getattr(request.app.state, "artifact_store", None)
    if store is None:
        raise HTTPException(status_code=503, detail="The artifact store is not ready")
    return store


@router.put("", response_model=ArtifactRef)
async def put_artifact(
    request: Request,
    name: str = Query(default=""),
    media_type: str = Query(default=""),
    volatile: bool = Query(default=False),
) -> Any:
    """Store the request body and answer with the reference to it.

    Spooled to disk as it arrives rather than buffered: a video segment is as
    legitimate a body here as a checkpoint, and neither should have to fit in
    memory twice. Capped, because nothing else here was: any account, and any
    worker credential, could otherwise fill the data volume.

    ``volatile`` puts it in the ring instead of the store — a frame a screen is
    watching now, which the oldest of falls out of memory rather than being
    kept. A worker on another host or a device publishing a camera sends this.
    It is read into memory, where the ring holds it, and never touches the
    disk.
    """
    store = _store(request)
    cap = settings.MAX_ARTIFACT_BYTES
    declared = request.headers.get("content-length")
    if cap and declared and declared.isdigit() and int(declared) > cap:
        raise HTTPException(
            status_code=413, detail=f"An artifact may be at most {cap} bytes"
        )
    ring = store.volatile if volatile else None
    if ring is not None:
        body = bytearray()
        async for chunk in request.stream():
            body += chunk
            if cap and len(body) > cap:
                raise HTTPException(
                    status_code=413,
                    detail=f"An artifact may be at most {cap} bytes",
                )
        return await run_in_threadpool(ring.put, [bytes(body)], name, media_type)
    handle = tempfile.NamedTemporaryFile(dir=store.root, delete=False)
    written = 0
    try:
        with handle:
            async for chunk in request.stream():
                written += len(chunk)
                # A chunked body declares no length, so the stream is what
                # actually holds the limit.
                if cap and written > cap:
                    raise HTTPException(
                        status_code=413,
                        detail=f"An artifact may be at most {cap} bytes",
                    )
                # Off the event loop: this is a write syscall per chunk, for
                # as long as the upload lasts.
                await run_in_threadpool(handle.write, chunk)
        return await run_in_threadpool(
            store.put_file, Path(handle.name), media_type, name, volatile
        )
    finally:
        Path(handle.name).unlink(missing_ok=True)


@router.get("/{digest}")
def get_artifact(
    digest: str, request: Request, media_type: str = Query(default="")
) -> Any:
    """Serve one artifact back.

    The caller passes the media type off the reference it holds, which is what
    lets a browser play a clip rather than download it; the store itself keeps
    only bytes. Ranged requests are answered because an audio or video element
    scrubbing through a file asks for them.
    """
    store = _store(request)
    path = store.path(digest)
    # A live frame is held in memory, not on disk.
    frame = store.volatile.get(digest) if path is None and store.volatile else None
    if path is None and frame is None:
        raise HTTPException(status_code=404, detail="No such artifact")
    if not _MEDIA_TYPE.match(media_type):
        media_type = "application/octet-stream"
    # The digest is the content, so it is also the perfect validator.
    headers = {"ETag": f'"{digest}"'}
    if not _INERT.match(media_type.lower()):
        # The caller names the type and any member can upload, so an artifact
        # asked for as text/html would otherwise run its script on this origin,
        # beside the session, for whoever follows a link to it.
        headers["Content-Security-Policy"] = "sandbox"
        headers["X-Content-Type-Options"] = "nosniff"
    if path is None:
        return Response(frame, media_type=media_type, headers=headers)
    return FileResponse(path, media_type=media_type, headers=headers)


@router.delete(
    "/{digest}",
    status_code=204,
    dependencies=[Depends(get_current_active_superuser)],
)
def delete_artifact(
    digest: str, request: Request, controller: FlowControllerDep
) -> Response:
    """Remove one artifact's bytes from the store.

    A superuser's call, because the store records no uploader and identical
    content from two callers is one file, so nobody else can be said to own
    it. Refused while anything still names the digest, the same references the
    hourly sweep keeps, so a reference that exists never answers 404. The
    volatile ring is left alone: its frames fall out on their own.
    """
    store = _store(request)
    if digest in referenced_artifacts(controller.state, controller.store):
        raise HTTPException(
            status_code=409,
            detail=(
                "The artifact is still referenced by a run, a current message or "
                "a flow input's initial value"
            ),
        )
    if not store.delete(digest):
        raise HTTPException(status_code=404, detail="No such artifact")
    return Response(status_code=204)