146 lines
7.7 KiB
Python
146 lines
7.7 KiB
Python
"""Переиспользуемый Blueprint слоя 2: POST /upload_refs (pull с ВМ-буфера в сессию).
|
||
|
||
Поведение:
|
||
- safe_name (защита от path traversal)
|
||
- SSRF-валидация url.startswith(vm_prefix)
|
||
- лимит на один файл -> delete + skip
|
||
- pull с ретраями через httpx stream
|
||
- лимит сессии -> skip; отсутствие сессии -> 404
|
||
- delete url с ВМ (best-effort)
|
||
- поддержка пофайлового транзита и пачек
|
||
- опциональный callback для интеграции/эмуляции Слоя 3 (on_file_received)
|
||
"""
|
||
|
||
import httpx
|
||
import logging
|
||
from urllib.parse import urlsplit
|
||
from flask import Blueprint, request, jsonify, current_app
|
||
|
||
from ..session import (create_session, add_file, get_files, file_count,
|
||
MAX_FILE_BYTES, configure)
|
||
from .config import PULL_RETRIES, PULL_RETRY_DELAY, VM_UPLOAD_PREFIX
|
||
from .safe_name import safe_name
|
||
from .pull_file import pull_file
|
||
|
||
log = logging.getLogger("upload.upload_refs")
|
||
|
||
|
||
def create_upload_refs_blueprint(cfg: dict = None) -> Blueprint:
|
||
"""Создать Blueprint с эндпоинтом upload_refs.
|
||
|
||
cfg (все ключи опциональны, есть дефолты):
|
||
apiPrefix (str) — префикс Blueprint, по умолчанию "/api"
|
||
vmUploadPrefix (str) — доверенный префикс ВМ-буфера (SSRF-валидация)
|
||
maxFileBytes (int) — лимит на один файл
|
||
maxSessionBytes (int) — суммарный лимит сессии (применяется к сессиям)
|
||
ttlSeconds (int) — TTL сессии
|
||
pullRetries (int) — ретраи pull
|
||
pullRetryDelay (int) — пауза между ретраями (сек)
|
||
pullTimeout (int) — таймаут одного GET pull
|
||
onFileReceived (func) — опциональный callback (sid, name, content) для Слоя 3
|
||
"""
|
||
cfg = cfg or {}
|
||
prefix = cfg.get("apiPrefix", "/api")
|
||
vm_prefix = cfg.get("vmUploadPrefix", VM_UPLOAD_PREFIX)
|
||
max_file_bytes = cfg.get("maxFileBytes", MAX_FILE_BYTES)
|
||
pull_retries = cfg.get("pullRetries", PULL_RETRIES)
|
||
pull_delay = cfg.get("pullRetryDelay", PULL_RETRY_DELAY)
|
||
pull_timeout = cfg.get("pullTimeout", 120)
|
||
on_file_received = cfg.get("onFileReceived")
|
||
|
||
# Применить лимиты сессии/TTL из конфига (глобально для всех сессий)
|
||
configure(
|
||
max_file_bytes=cfg.get("maxFileBytes"),
|
||
max_session_bytes=cfg.get("maxSessionBytes"),
|
||
ttl_seconds=cfg.get("ttlSeconds"),
|
||
)
|
||
|
||
is_path_only_prefix = vm_prefix.startswith("/")
|
||
|
||
bp = Blueprint("upload_refs", __name__, url_prefix=prefix)
|
||
|
||
@bp.route("/upload_refs", methods=["POST"])
|
||
def upload_refs():
|
||
"""Принять ссылки на файлы (загружены на ВМ-буфер), забрать по egress.
|
||
|
||
Вход: JSON {"session": "...", "files": [{"name": str, "size": int, "url": str}]}.
|
||
Каждый файл тянется исходящим GET с ВМ (egress не ограничен шлюзом),
|
||
читается по частям (stream), кладётся в сессию. После успешного pull файл
|
||
удаляется с ВМ (DELETE).
|
||
"""
|
||
data = request.get_json(silent=True) or {}
|
||
sid = data.get("session") or create_session()
|
||
refs = data.get("files") or []
|
||
if not refs:
|
||
log.warning("upload_refs: no files, sid=%s", sid)
|
||
return jsonify({"ok": False, "error": "No files"}), 400
|
||
added = 0
|
||
transport = cfg.get("httpxTransport") or current_app.config.get("UPLOAD_HTTPX_TRANSPORT")
|
||
try:
|
||
with httpx.Client(transport=transport, timeout=pull_timeout, follow_redirects=True) as client:
|
||
for ref in refs:
|
||
name = safe_name(ref.get("name") or "")
|
||
url = ref.get("url") or ""
|
||
if not name or not url:
|
||
continue
|
||
# SSRF-защита: тянуть можно ТОЛЬКО с доверенного ВМ-буфера
|
||
if is_path_only_prefix:
|
||
if not urlsplit(url).path.startswith(vm_prefix):
|
||
log.warning("upload_refs: unsafe URL path, skip sid=%s url=%r", sid, url)
|
||
continue
|
||
if not url.startswith(("http://", "https://")):
|
||
url = request.host_url.rstrip("/") + ("/" if not url.startswith("/") else "") + url
|
||
else:
|
||
if not url.startswith(vm_prefix):
|
||
log.warning("upload_refs: unsafe URL, skip sid=%s url=%r", sid, url)
|
||
continue
|
||
# Лимит на один файл: сверх лимита — пропускаем
|
||
if (ref.get("size") or 0) > max_file_bytes:
|
||
log.warning("upload_refs: file exceeds %dMB, skip sid=%s file=%r size=%s",
|
||
max_file_bytes // (1024 * 1024), sid, name, ref.get("size"))
|
||
try:
|
||
client.delete(url)
|
||
except Exception:
|
||
pass
|
||
continue
|
||
# Pull с ретраями
|
||
content = pull_file(client, url, pull_retries, pull_delay, sid=sid, name=name)
|
||
log.info("upload_refs: pulled sid=%s file=%r size=%d", sid, name, len(content))
|
||
if len(content) > max_file_bytes:
|
||
log.warning("upload_refs: pulled file exceeds %dMB, skip sid=%s file=%r size=%d",
|
||
max_file_bytes // (1024 * 1024), sid, name, len(content))
|
||
try:
|
||
client.delete(url)
|
||
except Exception:
|
||
pass
|
||
continue
|
||
if not add_file(sid, name, content):
|
||
if get_files(sid) is None:
|
||
log.warning("upload_refs: session not found, sid=%s file=%r", sid, name)
|
||
return jsonify({"ok": False, "error": "Session not found"}), 404
|
||
log.warning("upload_refs: session limit exceeded, skip sid=%s file=%r", sid, name)
|
||
try:
|
||
client.delete(url)
|
||
except Exception:
|
||
pass
|
||
continue
|
||
# Успешно добавлен в сессию — удаляем с ВМ-буфера
|
||
try:
|
||
client.delete(url)
|
||
except Exception:
|
||
pass
|
||
# Если передан callback для Слоя 3 (эмуляция или реальный процессинг)
|
||
if callable(on_file_received):
|
||
try:
|
||
on_file_received(sid, name, content)
|
||
except Exception as cb_err:
|
||
log.warning("upload_refs: on_file_received callback error: %r", cb_err)
|
||
added += 1
|
||
except Exception as e:
|
||
log.error("upload_refs: pull error sid=%s: %r", sid, e)
|
||
return jsonify({"ok": False, "error": "Pull failed: %s" % e}), 502
|
||
log.info("upload_refs: done sid=%s added=%d total=%d", sid, added, file_count(sid))
|
||
return jsonify({"ok": True, "session": sid, "count": file_count(sid), "added": added})
|
||
|
||
return bp
|