"""Policy, quota, and lifecycle ledger for browser-direct S3 uploads.""" from __future__ import annotations import logging import re import unicodedata import uuid from base64 import b64decode, b64encode from collections.abc import Awaitable from datetime import UTC, datetime, timedelta from typing import Final, cast from sqlalchemy import delete as sa_delete from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.core.config import settings from app.core.exceptions import AppException from app.core.redis_client import get_redis from app.core.storage import resolve_storage_config, stat_object from app.integrations.models import StorageUploadGrant logger = logging.getLogger(__name__) PURPOSE_MAX_BYTES: Final[dict[str, int]] = { "inventory_evidence": 10 * 1024 * 1024, "asset_disposal_evidence": 10 * 1024 * 1024, "feed_post_image": 10 * 1024 * 1024, "feed_announcement": 10 * 1024 * 1024, # 安全宣贯素材 —— 漫画页 / 文章配图 / 封面. A single comic page # scanned at print resolution routinely lands near 8 MB, so this # shares the platform's 10 MB image ceiling rather than getting a # tighter one that would push the security team into re-compressing # artwork by hand. "awareness_material": 10 * 1024 * 1024, # AI Labs inputs. Images are base64-inlined into the vendor request, so # the wire cost is ~1.33x the file — 10 MB is the point past which a # chat turn starts to feel broken. Documents are read server-side and # only their extracted text travels, so a larger file is cheap. Audio # is inlined like images; 15 MB is roughly 15 minutes of speech. "ai_labs_image": 10 * 1024 * 1024, "ai_labs_document": 20 * 1024 * 1024, "ai_labs_audio": 15 * 1024 * 1024, } ALLOWED_PURPOSES: Final[frozenset[str]] = frozenset(PURPOSE_MAX_BYTES) ALLOWED_CONTENT_TYPES: Final[frozenset[str]] = frozenset({ "image/jpeg", "image/png", "image/webp", "image/heic", "image/heif", }) #: Per-purpose content-type allow-list, consulted **before** #: :data:`ALLOWED_CONTENT_TYPES`. #: #: The global set is image-only, which was right while every upload surface #: was a photo. AI Labs needs documents and audio, and widening the global #: set to admit them would also let a PDF be posted as a feed photo. A #: per-purpose map keeps each surface as narrow as it actually is: a purpose #: absent from this map keeps its previous behaviour exactly. PURPOSE_CONTENT_TYPES: Final[dict[str, frozenset[str]]] = { "ai_labs_image": frozenset({ "image/jpeg", "image/png", "image/webp", }), "ai_labs_document": frozenset({ "application/pdf", ( "application/vnd.openxmlformats-officedocument" ".wordprocessingml.document" ), "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "text/plain", "text/markdown", "text/csv", }), "ai_labs_audio": frozenset({ "audio/mpeg", "audio/mp4", "audio/wav", "audio/x-wav", "audio/webm", "audio/aac", "audio/ogg", }), } def content_types_for_purpose(purpose: str) -> frozenset[str]: """The types this purpose accepts, defaulting to the global image set.""" return PURPOSE_CONTENT_TYPES.get(purpose, ALLOWED_CONTENT_TYPES) _MAX_FILENAME_LEN = 120 _FILENAME_SAFE_RE = re.compile(r"[^A-Za-z0-9._\-]") _UPLOAD_QUOTA_SCRIPT = """ local calls = tonumber(redis.call('GET', KEYS[1]) or '0') local bytes = tonumber(redis.call('GET', KEYS[2]) or '0') local add_bytes = tonumber(ARGV[1]) local max_calls = tonumber(ARGV[2]) local max_bytes = tonumber(ARGV[3]) if calls + 1 > max_calls then return -1 end if bytes + add_bytes > max_bytes then return -2 end redis.call('INCR', KEYS[1]) redis.call('INCRBY', KEYS[2], add_bytes) if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], 3600) end if redis.call('TTL', KEYS[2]) < 0 then redis.call('EXPIRE', KEYS[2], 172800) end return 1 """ #: Characters never allowed in a name that will be shown to a person. #: #: C0/C1 controls (they break a header and a log line) and the path #: separators, plus the Unicode **bidi and format** controls. Those last #: ones are why the class is not just "controls": ``x\u202egnp.exe`` #: renders as ``xexe.png``, so a name can display as a different file type #: than it is. Only the uploader ever sees an AI Labs name, which is why #: this is a small hazard rather than a large one — but a sanitiser whose #: docstring says it strips controls should strip the ones that matter. #: #: This class covers the C0/C1 controls and the path separators. The #: **format** controls are removed separately, by Unicode category, in #: :func:`safe_display_filename`: an enumerated range missed U+00AD SOFT #: HYPHEN, U+2060 WORD JOINER, U+180E and the U+FFF9-FFFB annotation #: controls, all of which are exactly what the docstring claimed to #: remove, and a hand-listed range goes stale the way that one did. _DISPLAY_UNSAFE_RE: Final = re.compile(r"[\x00-\x1f\x7f-\x9f/\\]+") #: Cap on a display name. Long enough for a real document title, short #: enough that a 4 KB filename cannot pad a row or a prompt. _MAX_DISPLAY_NAME_LEN: Final[int] = 200 def safe_display_filename(raw: str) -> str: """Sanitise a filename that will be **shown**, not used as a key. Deliberately NOT :func:`safe_upload_filename`. That one produces an S3-key segment, so it replaces every character outside ``[A-Za-z0-9._-]`` — which for a CJK or Japanese name means the whole stem becomes underscores and then ``.strip("._")`` eats them **and the dot**: ``年度报告.pdf`` came out as ``pdf``, ``離職面談.m4a`` as ``m4a``, ``面談メモ.docx`` as ``docx``. Using it on a display name destroyed exactly the value AI Labs now sends from the client to keep (the object key no longer carries it), and in a company whose primary language is Chinese that is the common case, not the edge case. Keeps every printable character including CJK; strips controls and path separators; collapses whitespace; bounds the length. Falls back to ``file`` so a caller always has something to render. """ cleaned = _DISPLAY_UNSAFE_RE.sub("", raw or "") # Then every format control, by Unicode category. A bidi override # (U+202E) makes ``xgnp.exe`` render as ``xexe.png``, so a name # can display as a different file type than it is; the rest are # invisible joiners and soft hyphens that let two different names look # identical. ``Cf`` is the category that covers all of them, and it # cannot go stale the way a hand-listed range did. cleaned = "".join(c for c in cleaned if unicodedata.category(c) != "Cf") cleaned = " ".join(cleaned.split()).strip() if len(cleaned) > _MAX_DISPLAY_NAME_LEN: cleaned = cleaned[:_MAX_DISPLAY_NAME_LEN].rstrip() return cleaned or "file" def safe_upload_filename(raw: str) -> str: no_path = raw.replace("\\", "/").rsplit("/", 1)[-1] cleaned = _FILENAME_SAFE_RE.sub("_", no_path).strip("._") or "file" return cleaned[-_MAX_FILENAME_LEN:] #: Upper bound on how long any single stored object key may be — matches #: ``storage_router.PresignGetRequest.object_key``'s own ``max_length``. MAX_OBJECT_KEY_LEN: Final[int] = 512 def sanitize_object_key_list( raw: list[str] | None, *, field: str, max_items: int, ) -> list[str] | None: """Normalise + bound-check a caller-supplied list of S3 object keys. Shared by every domain that persists browser-uploaded keys on its own row (inventory evidence photos, asset disposal evidence, …). Returns ``None`` for ``None`` / empty input, otherwise the stripped list. Raises ``ValueError`` on bad entries so a Pydantic ``field_validator`` can adopt it verbatim; service layers wrap it in their own domain exception. Blocks obvious path-traversal / scheme-prefix attempts. Real keys are produced by ``/uploads/presign-put`` and are never user-typed, but this list arrives on the wire and defence-in-depth costs nothing — the same key later reaches ``presign_get`` verbatim. """ if not raw: return None if len(raw) > max_items: raise ValueError(f"{field} too long: {len(raw)} > {max_items}") cleaned: list[str] = [] seen: set[str] = set() for key in raw: if not isinstance(key, str): raise ValueError(f"{field} entries must be strings") stripped = key.strip() if not stripped or len(stripped) > MAX_OBJECT_KEY_LEN: raise ValueError(f"{field} key invalid length") if ( ".." in stripped or stripped.startswith("/") or "://" in stripped ): raise ValueError(f"{field} key contains invalid path characters") if stripped in seen: raise ValueError(f"{field} contains duplicate object keys") seen.add(stripped) cleaned.append(stripped) return cleaned #: How long a grant stays claimable, per purpose. #: #: The default (10 minutes) fits every other consumer, where the gap #: between the presigned upload and the row that references it is one form #: submit. **AI Labs is different in kind**: the browser uploads at the #: moment the employee picks the file, and the claim happens when they #: press send — so the gap is "how long they spend writing their #: question", or composing a prompt in the generation studio. At the #: default, attaching a PDF and then thinking for twelve minutes produced #: a 409 ``UPLOAD_GRANT_EXPIRED`` that cost the employee their typed text #: and their attachment, with a raw platform string as the explanation. #: #: Four hours, which is the shape of a working session rather than of a #: form submit. The grant is still single-purpose, still owner-scoped, and #: still verified against the landed object's size, type and checksum at #: claim time. #: #: It did weaken one thing, and saying so is the point: the orphan sweep #: selected on ``created_at`` alone and never read ``expires_at``, which #: was safe only because every TTL used to be shorter than the sweep's #: one-hour floor. A 4h TTL inverted that ordering, so #: ``storage_handlers._eligible_claim`` now requires the grant to be #: expired as well — the ordering is an invariant rather than a #: coincidence, whatever TTL anyone sets here later. GRANT_TTL_BY_PURPOSE: Final[dict[str, int]] = { "ai_labs_image": 4 * 60 * 60, "ai_labs_document": 4 * 60 * 60, "ai_labs_audio": 4 * 60 * 60, } #: Default grant lifetime for every other purpose. DEFAULT_GRANT_TTL_SECONDS: Final[int] = 600 def grant_ttl_for_purpose(purpose: str) -> int: """Seconds a freshly minted grant for ``purpose`` stays claimable.""" return GRANT_TTL_BY_PURPOSE.get(purpose, DEFAULT_GRANT_TTL_SECONDS) #: Purposes whose stored key must NOT carry the uploader's filename. #: #: Every other purpose keeps it, and that is genuinely useful: an admin #: resolving an inventory dispute wants to see ``receipt-2026-03.pdf`` in #: the audit trail. AI Labs is the one place where the filename is the #: employee's private content — ``离职谈话录音.m4a``, ``分手信.txt`` — and #: the product tells them nobody at Helix reads what they put in there. #: #: The key is the wrong place to fix that by filtering readers, because it #: is copied into three stores that outlive the employee's own delete: the #: ``storage_upload_grants`` ledger row, the ``storage.upload.authorize`` #: audit row (company-readable through ``/audit/logs`` and its CSV export), #: and the S3 listing itself. Not writing the name is the only version of #: this that stays true. The display name still exists — it is stored on #: the employee's own message row, which the hard delete does remove. OPAQUE_NAME_PURPOSES: Final[frozenset[str]] = frozenset({ "ai_labs_image", "ai_labs_document", "ai_labs_audio", }) #: Extension characters we are willing to echo into a key. Deliberately #: narrow: the extension is kept only so an object is recognisable to a #: human debugging a bucket, and anything unusual is dropped rather than #: sanitised into something that might still carry a word. _EXTENSION_RE: Final = re.compile(r"^[a-z0-9]{1,8}$") def _opaque_object_name(filename: str) -> str: """A random name, keeping only a plain lowercase extension. ``rpartition`` is checked for an actual separator, not just a non-empty tail: on a dotless name it returns ``('', '', )``, so the **entire filename** became the "extension" whenever it matched the narrow pattern. An upload called ``resign``, ``payslip`` or ``nda`` produced the key ``.resign`` — which is then copied into the grant ledger and the company-readable ``storage.upload.authorize`` audit row and survives the employee's own hard delete, i.e. exactly the chain this function exists to break. """ stem = uuid.uuid4().hex # The extension comes from the ORIGINAL name, not from # ``safe_upload_filename``'s output. Running the key sanitiser first # was a bug where two fixes met: it collapses a CJK stem to nothing, # so ``离职谈话录音.m4a`` became the dotless ``m4a`` and this function # then found no extension at all — and ``attachments.load`` derives an # attachment's CONTENT TYPE from the key's extension, so a # Chinese-named upload lost its type as well as its name. head, separator, raw_ext = (filename or "").rpartition(".") if not separator or not head.strip(): # Dotless (``payslip``), or a leading-dot name (``.bashrc``): there # is no extension, and the whole name must not be echoed. Checking # the separator matters — ``rpartition`` on a dotless string # returns the entire string as the tail, so the earlier version # produced the key ``.payslip``, which then travelled into # the grant ledger and the company-readable audit row. return stem ext = raw_ext.strip().lower() return f"{stem}.{ext}" if _EXTENSION_RE.match(ext) else stem def build_relative_upload_key( *, purpose: str, person_id: uuid.UUID, filename: str, ) -> str: """The stored key for one upload. Layout is fixed at ``{purpose}/{yyyy}/{mm}/{dd}/{person_id}/{name}`` — ``storage_router.parse_object_key`` requires exactly those five segments after the purpose to recover the uploader for its read ACL, so the opaque branch below changes only the last segment. """ today = datetime.now(UTC).strftime("%Y/%m/%d") if purpose in OPAQUE_NAME_PURPOSES: return f"{purpose}/{today}/{person_id}/{_opaque_object_name(filename)}" rand = uuid.uuid4().hex[:12] return f"{purpose}/{today}/{person_id}/{rand}-{safe_upload_filename(filename)}" def _as_utc(value: datetime) -> datetime: return value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC) def normalize_sha256_checksum(value: str) -> str: """Return one canonical base64-encoded 256-bit SHA-256 digest.""" try: raw = b64decode(value, validate=True) except (TypeError, ValueError) as exc: raise AppException( "UPLOAD_CHECKSUM_INVALID", "checksum_sha256 must be a base64-encoded SHA-256 digest", 422, ) from exc if len(raw) != 32: raise AppException( "UPLOAD_CHECKSUM_INVALID", "checksum_sha256 must be a base64-encoded SHA-256 digest", 422, ) return b64encode(raw).decode("ascii") async def _enforce_upload_quota( person_id: uuid.UUID, size_bytes: int, *, namespace: str, max_calls: int, max_bytes: int, ) -> None: """Atomically enforce a per-person count and byte budget.""" now = datetime.now(UTC) calls_key = f"{namespace}:calls:{person_id}:{now:%Y%m%d%H}" bytes_key = f"{namespace}:bytes:{person_id}:{now:%Y%m%d}" try: redis = get_redis() raw_result = await cast( Awaitable[str | bytes | int], redis.eval( _UPLOAD_QUOTA_SCRIPT, 2, calls_key, bytes_key, size_bytes, max_calls, max_bytes, ), ) result = int(raw_result) except Exception as exc: logger.exception("upload presign quota unavailable") raise AppException( code="UPLOAD_QUOTA_UNAVAILABLE", message="upload quota service is temporarily unavailable", status_code=503, ) from exc if result == -1: raise AppException( code="UPLOAD_RATE_LIMITED", message="too many upload authorizations; try again later", status_code=429, headers={"Retry-After": "3600"}, ) if result == -2: raise AppException( code="UPLOAD_DAILY_QUOTA_EXCEEDED", message="daily upload byte quota exceeded", status_code=429, headers={"Retry-After": "86400"}, ) if result != 1: raise AppException( code="UPLOAD_QUOTA_UNAVAILABLE", message="upload quota service returned an invalid result", status_code=503, ) async def enforce_upload_presign_quota(person_id: uuid.UUID, size_bytes: int) -> None: """Atomically enforce browser-direct presign budgets.""" await _enforce_upload_quota( person_id, size_bytes, namespace="upload_presign", max_calls=settings.UPLOAD_PRESIGN_MAX_PER_HOUR, max_bytes=settings.UPLOAD_PRESIGN_MAX_BYTES_PER_DAY, ) async def enforce_approval_upload_quota(person_id: uuid.UUID, size_bytes: int) -> None: """Bound multipart approval uploads before any external storage call.""" await _enforce_upload_quota( person_id, size_bytes, namespace="approval_upload", max_calls=settings.APPROVAL_UPLOAD_MAX_PER_HOUR, max_bytes=settings.APPROVAL_UPLOAD_MAX_BYTES_PER_DAY, ) async def create_upload_grant( db: AsyncSession, *, person_id: uuid.UUID, purpose: str, filename: str, content_type: str, size_bytes: int, checksum_sha256: str, expires_in: int | None = None, ) -> StorageUploadGrant: if purpose not in ALLOWED_PURPOSES: raise AppException("UPLOAD_PURPOSE_NOT_ALLOWED", "upload purpose not allowed", 400) allowed_types = content_types_for_purpose(purpose) if content_type not in allowed_types: raise AppException("UPLOAD_TYPE_NOT_ALLOWED", "upload content type not allowed", 400) return await create_scoped_upload_grant( db, person_id=person_id, purpose=purpose, relative_key=build_relative_upload_key( purpose=purpose, person_id=person_id, filename=filename, ), content_type=content_type, size_bytes=size_bytes, checksum_sha256=checksum_sha256, allowed_content_types=allowed_types, max_bytes=PURPOSE_MAX_BYTES[purpose], expires_in=expires_in, ) async def create_scoped_upload_grant( db: AsyncSession, *, person_id: uuid.UUID, purpose: str, relative_key: str, content_type: str, size_bytes: int, checksum_sha256: str, allowed_content_types: frozenset[str], max_bytes: int, expires_in: int | None = None, ) -> StorageUploadGrant: """Issue a durable, quota-controlled grant for a domain-owned key. ``expires_in`` defaults to the purpose's own TTL — see :func:`grant_ttl_for_purpose`, which exists because AI Labs uploads at file-pick time and claims at send, a gap measured in minutes of typing rather than in one form submit. """ if expires_in is None: expires_in = grant_ttl_for_purpose(purpose) if content_type not in allowed_content_types: raise AppException("UPLOAD_TYPE_NOT_ALLOWED", "upload content type not allowed", 400) if size_bytes <= 0 or size_bytes > max_bytes: raise AppException( code="UPLOAD_SIZE_NOT_ALLOWED", message=f"upload size must be between 1 and {max_bytes} bytes", status_code=413, ) checksum = normalize_sha256_checksum(checksum_sha256) await enforce_upload_presign_quota(person_id, size_bytes) cfg = await resolve_storage_config(db) grant = StorageUploadGrant( person_id=person_id, purpose=purpose, object_key=cfg["key_prefix"] + relative_key.lstrip("/"), content_type=content_type, declared_size_bytes=size_bytes, checksum_sha256=checksum, status="issued", expires_at=datetime.now(UTC) + timedelta(seconds=expires_in), ) db.add(grant) await db.flush() return grant async def forget_upload_grants( db: AsyncSession, *, object_keys: list[str], person_id: uuid.UUID | None, ) -> int: """Hard-delete this person's grant ledger rows for these keys. The counterpart to :func:`mark_uploads_referenced` for a domain whose delete is a **hard** delete of the content itself, which today means AI Labs. Everywhere else a grant row outliving its resource is a harmless audit breadcrumb; there it is not. A ``referenced`` grant row keeps ``person_id``, ``purpose``, ``checksum_sha256`` — a full content fingerprint — ``declared_size_bytes``, ``content_type`` and ``created_at``, and nothing ever removes it. Someone holding a copy of a document could hash it, find the row, and learn that a named employee put exactly that file into AI Labs at a given minute, months after the employee deleted the conversation. That is the same promise the opaque object keys exist to keep, so the same delete has to reach this row. Scoped to ``person_id`` so a request-path caller cannot release someone else's grant by passing a key it happens to know. ``None`` means "no owner scoping", which is only for a system sweep acting on rows it has itself just deleted — pass an id on every path that starts with a request. Returns how many rows went. """ keys = [key for key in dict.fromkeys(object_keys) if key] if not keys: return 0 stmt = sa_delete(StorageUploadGrant).where( StorageUploadGrant.object_key.in_(keys), ) if person_id is not None: stmt = stmt.where(StorageUploadGrant.person_id == person_id) result = await db.execute(stmt) # ``rowcount`` is on ``CursorResult`` — see the note in # ``ai_labs.executor``; the declared ``Result`` union lacks it. return int(getattr(result, "rowcount", 0) or 0) async def mark_uploads_referenced( db: AsyncSession, *, object_keys: list[str], person_id: uuid.UUID, allowed_purposes: frozenset[str], resource_type: str, resource_id: str, ) -> dict[str, StorageUploadGrant]: """Verify landed objects and atomically claim their durable references.""" keys = list(dict.fromkeys(object_keys)) if not keys: return {} rows = (await db.execute( select(StorageUploadGrant).where(StorageUploadGrant.object_key.in_(keys)) )).scalars().all() by_key = {row.object_key: row for row in rows} if set(by_key) != set(keys): raise AppException("UPLOAD_GRANT_NOT_FOUND", "upload authorization not found", 400) now = datetime.now(UTC) for key in keys: row = by_key[key] if row.person_id != person_id or row.purpose not in allowed_purposes: raise AppException("UPLOAD_GRANT_FORBIDDEN", "upload authorization mismatch", 403) if row.status == "referenced": if ( row.referenced_resource_type == resource_type and row.referenced_resource_id == resource_id ): continue raise AppException("UPLOAD_ALREADY_REFERENCED", "upload is already referenced", 409) if row.status != "issued" or _as_utc(row.expires_at) < now: raise AppException("UPLOAD_GRANT_EXPIRED", "upload authorization expired", 409) actual = await stat_object(db, object_key=key) if actual is None: raise AppException("UPLOAD_NOT_FOUND", "uploaded object not found", 400) if ( actual["size_bytes"] != row.declared_size_bytes or actual["content_type"] != row.content_type or actual["checksum_sha256"] != row.checksum_sha256 ): raise AppException("UPLOAD_METADATA_MISMATCH", "uploaded object metadata mismatch", 400) locked = (await db.execute( select(StorageUploadGrant) .where(StorageUploadGrant.object_key.in_(keys)) .with_for_update() )).scalars().all() locked_by_key = {row.object_key: row for row in locked} for key in keys: row = locked_by_key[key] if row.status == "referenced": if ( row.referenced_resource_type == resource_type and row.referenced_resource_id == resource_id ): continue raise AppException("UPLOAD_ALREADY_REFERENCED", "upload is already referenced", 409) if row.status != "issued" or _as_utc(row.expires_at) < datetime.now(UTC): raise AppException("UPLOAD_GRANT_EXPIRED", "upload authorization expired", 409) row.status = "referenced" row.referenced_resource_type = resource_type row.referenced_resource_id = resource_id row.referenced_at = datetime.now(UTC) return locked_by_key