mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 08:16:38 +00:00
feat: image refactor
This commit is contained in:
1 parent
1c2b0f506a
commit
4b1bc17c77
28 files changed
+1192
-266
No files matched your search
@@ -244,7 +244,11 @@ def _format_agent_output(
|
||||
"slug": agent.get("slug", "") or "",
|
||||
"description": agent.get("description", "") or "",
|
||||
"image": (
|
||||
generate_image_url(agent["image"]) if agent.get("image") else ""
|
||||
generate_image_url(
|
||||
agent["image"], agent["id"], agent.get("user_id")
|
||||
)
|
||||
if agent.get("image")
|
||||
else ""
|
||||
),
|
||||
"source": source_value,
|
||||
"sources": sources_list,
|
||||
@@ -308,7 +312,7 @@ def _build_create_kwargs(data: dict, *, image_url: str, agent_type: str) -> dict
|
||||
allowed_fields = set(schema["fields"])
|
||||
|
||||
for key in (
|
||||
"description", "agent_type", "key", "image", "retriever",
|
||||
"description", "agent_type", "key", "retriever",
|
||||
"default_model_id",
|
||||
):
|
||||
if key in allowed_fields and data.get(key) not in (None, ""):
|
||||
@@ -787,8 +791,8 @@ class UpdateAgent(Resource):
|
||||
"description": fields.String(
|
||||
required=True, description="New description of the agent"
|
||||
),
|
||||
"image": fields.String(
|
||||
required=False, description="New image URL or identifier"
|
||||
"image": fields.Raw(
|
||||
required=False, description="Image file upload", type="file"
|
||||
),
|
||||
"source": fields.String(
|
||||
required=False, description="Source ID (legacy single source)"
|
||||
@@ -913,8 +917,12 @@ class UpdateAgent(Resource):
|
||||
404,
|
||||
)
|
||||
pg_agent_id = str(existing_agent["id"])
|
||||
existing_image = existing_agent.get("image", "") or ""
|
||||
image_url, image_error = handle_image_upload(
|
||||
request, existing_agent.get("image", "") or "", user, storage,
|
||||
request,
|
||||
existing_image,
|
||||
existing_agent.get("user_id") or user,
|
||||
storage,
|
||||
)
|
||||
if image_error:
|
||||
return image_error
|
||||
@@ -923,7 +931,6 @@ class UpdateAgent(Resource):
|
||||
allowed_fields = [
|
||||
"name",
|
||||
"description",
|
||||
"image",
|
||||
"source",
|
||||
"sources",
|
||||
"chunks",
|
||||
@@ -1133,7 +1140,7 @@ class UpdateAgent(Resource):
|
||||
f"Field '{field}' cannot be empty", user, field
|
||||
)
|
||||
update_fields[field] = value
|
||||
if image_url:
|
||||
if image_url and image_url != existing_image:
|
||||
update_fields["image"] = image_url
|
||||
if not update_fields:
|
||||
return _reject("No valid update data provided", user)
|
||||
@@ -1543,7 +1550,11 @@ class PinnedAgents(Resource):
|
||||
"name": agent.get("name", ""),
|
||||
"description": agent.get("description", ""),
|
||||
"image": (
|
||||
generate_image_url(agent["image"]) if agent.get("image") else ""
|
||||
generate_image_url(
|
||||
agent["image"], agent["id"], agent.get("user_id")
|
||||
)
|
||||
if agent.get("image")
|
||||
else ""
|
||||
),
|
||||
"source": str(source_id) if source_id else "",
|
||||
"chunks": str(agent["chunks"]) if agent.get("chunks") is not None else "",
|
||||
@@ -1591,7 +1602,13 @@ class GetTemplateAgents(Resource):
|
||||
"id": str(agent["id"]),
|
||||
"name": agent.get("name"),
|
||||
"description": agent.get("description") or "",
|
||||
"image": agent.get("image") or "",
|
||||
"image": (
|
||||
generate_image_url(
|
||||
agent["image"], agent["id"], agent.get("user_id")
|
||||
)
|
||||
if agent.get("image")
|
||||
else ""
|
||||
),
|
||||
}
|
||||
for agent in template_rows
|
||||
]
|
||||
|
||||
@@ -29,7 +29,11 @@ def _serialize_agent_basic(agent: dict) -> dict:
|
||||
"user": agent.get("user_id", ""),
|
||||
"name": agent.get("name", ""),
|
||||
"image": (
|
||||
generate_image_url(agent["image"]) if agent.get("image") else ""
|
||||
generate_image_url(
|
||||
agent["image"], agent["id"], agent.get("user_id")
|
||||
)
|
||||
if agent.get("image")
|
||||
else ""
|
||||
),
|
||||
"description": agent.get("description", ""),
|
||||
"source": str(source_id) if source_id else "",
|
||||
@@ -158,7 +162,11 @@ class SharedAgents(Resource):
|
||||
"name": agent.get("name", ""),
|
||||
"description": agent.get("description", ""),
|
||||
"image": (
|
||||
generate_image_url(agent["image"]) if agent.get("image") else ""
|
||||
generate_image_url(
|
||||
agent["image"], agent["id"], agent.get("user_id")
|
||||
)
|
||||
if agent.get("image")
|
||||
else ""
|
||||
),
|
||||
"tools": agent.get("tools", []) or [],
|
||||
"tool_details": resolve_tool_details(
|
||||
|
||||
@@ -6,12 +6,13 @@ from pathlib import Path
|
||||
|
||||
import uuid
|
||||
|
||||
from flask import current_app, jsonify, make_response, request
|
||||
from flask import Response, current_app, jsonify, make_response, redirect, request
|
||||
from flask_restx import fields, Namespace, Resource
|
||||
|
||||
from application.api import api
|
||||
from application.cache import get_redis_instance
|
||||
from application.core.settings import settings
|
||||
from application.storage.db.base_repository import looks_like_uuid
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.session import db_readonly
|
||||
from application.stt.constants import (
|
||||
@@ -35,13 +36,37 @@ from application.stt.live_session import (
|
||||
)
|
||||
from application.stt.stt_creator import STTCreator
|
||||
from application.tts.tts_creator import TTSCreator
|
||||
from application.utils import safe_filename
|
||||
from application.utils import (
|
||||
get_agent_image_content_type,
|
||||
is_external_image_url,
|
||||
is_safe_agent_image_path,
|
||||
safe_filename,
|
||||
verify_agent_image_capability,
|
||||
)
|
||||
|
||||
|
||||
attachments_ns = Namespace(
|
||||
"attachments", description="File attachments and media operations", path="/api"
|
||||
)
|
||||
|
||||
_AGENT_IMAGE_CHUNK_BYTES = 64 * 1024
|
||||
_AGENT_IMAGE_PRESIGNED_TTL_SECONDS = 300
|
||||
_AGENT_IMAGE_REDIRECT_CACHE_SECONDS = _AGENT_IMAGE_PRESIGNED_TTL_SECONDS - 60
|
||||
|
||||
|
||||
def _stream_file(file_obj, size_bytes: int):
|
||||
"""Yield exactly the size-checked file range and always close the handle."""
|
||||
remaining = size_bytes
|
||||
try:
|
||||
while remaining > 0:
|
||||
chunk = file_obj.read(min(_AGENT_IMAGE_CHUNK_BYTES, remaining))
|
||||
if not chunk:
|
||||
break
|
||||
remaining -= len(chunk)
|
||||
yield chunk
|
||||
finally:
|
||||
file_obj.close()
|
||||
|
||||
|
||||
def _resolve_authenticated_user():
|
||||
decoded_token = getattr(request, "decoded_token", None)
|
||||
@@ -630,25 +655,71 @@ class LiveSpeechToTextFinish(Resource):
|
||||
)
|
||||
|
||||
|
||||
@attachments_ns.route("/images/<path:image_path>")
|
||||
@attachments_ns.route("/images/<string:agent_id>/<string:capability>")
|
||||
class ServeImage(Resource):
|
||||
@api.doc(description="Serve an image from storage")
|
||||
def get(self, image_path):
|
||||
if ".." in image_path or image_path.startswith("/") or "\x00" in image_path:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Invalid image path"}), 400
|
||||
)
|
||||
@api.doc(description="Serve an agent image using an opaque capability URL")
|
||||
def get(self, agent_id, capability):
|
||||
try:
|
||||
if not looks_like_uuid(agent_id):
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Image not found"}), 404
|
||||
)
|
||||
with db_readonly() as conn:
|
||||
agent = AgentsRepository(conn).find_image_record(agent_id)
|
||||
if not agent:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Image not found"}), 404
|
||||
)
|
||||
|
||||
image_path = agent.get("image")
|
||||
user_id = agent.get("user_id")
|
||||
content_type = get_agent_image_content_type(image_path)
|
||||
if (
|
||||
is_external_image_url(image_path)
|
||||
or not content_type
|
||||
or not is_safe_agent_image_path(image_path, user_id)
|
||||
or not verify_agent_image_capability(
|
||||
capability, agent_id, image_path, user_id
|
||||
)
|
||||
):
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Image not found"}), 404
|
||||
)
|
||||
|
||||
from application.api.user.base import storage
|
||||
|
||||
size_bytes = storage.get_file_size(image_path)
|
||||
if size_bytes < 0 or size_bytes > settings.AGENT_IMAGE_MAX_BYTES:
|
||||
current_app.logger.warning(
|
||||
"Rejected oversized legacy agent image for agent %s", agent_id
|
||||
)
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Image not found"}), 404
|
||||
)
|
||||
|
||||
if settings.STORAGE_TYPE == "s3":
|
||||
image_url = storage.generate_presigned_url(
|
||||
image_path,
|
||||
expires_in=_AGENT_IMAGE_PRESIGNED_TTL_SECONDS,
|
||||
content_type=content_type,
|
||||
)
|
||||
response = redirect(image_url, code=302)
|
||||
response.headers.set(
|
||||
"Cache-Control",
|
||||
f"private, max-age={_AGENT_IMAGE_REDIRECT_CACHE_SECONDS}",
|
||||
)
|
||||
response.headers.set("X-Content-Type-Options", "nosniff")
|
||||
return response
|
||||
|
||||
file_obj = storage.get_file(image_path)
|
||||
extension = image_path.split(".")[-1].lower()
|
||||
content_type = f"image/{extension}"
|
||||
if extension == "jpg":
|
||||
content_type = "image/jpeg"
|
||||
response = make_response(file_obj.read())
|
||||
response.headers.set("Content-Type", content_type)
|
||||
response.headers.set("Cache-Control", "max-age=86400")
|
||||
response = Response(
|
||||
_stream_file(file_obj, size_bytes),
|
||||
content_type=content_type,
|
||||
direct_passthrough=True,
|
||||
)
|
||||
response.content_length = size_bytes
|
||||
response.headers.set("Cache-Control", "public, max-age=86400, immutable")
|
||||
response.headers.set("X-Content-Type-Options", "nosniff")
|
||||
|
||||
return response
|
||||
except FileNotFoundError:
|
||||
@@ -657,7 +728,7 @@ class ServeImage(Resource):
|
||||
)
|
||||
except ValueError:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Invalid image path"}), 400
|
||||
jsonify({"success": False, "message": "Image not found"}), 404
|
||||
)
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error serving image: {e}")
|
||||
|
||||
@@ -6,9 +6,11 @@ import datetime
|
||||
import os
|
||||
import uuid
|
||||
from functools import wraps
|
||||
from pathlib import PurePosixPath
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from flask import current_app, jsonify, make_response, Response
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
from werkzeug.utils import secure_filename
|
||||
|
||||
from sqlalchemy import text as _sql_text
|
||||
@@ -18,6 +20,7 @@ from application.storage.db.base_repository import looks_like_uuid, row_to_dict
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
from application.storage.db.session import db_readonly, db_session
|
||||
from application.storage.storage_creator import StorageCreator
|
||||
from application.utils import AGENT_IMAGE_FORMATS, safe_user_storage_component
|
||||
from application.vectorstore.vector_creator import VectorCreator
|
||||
|
||||
|
||||
@@ -28,7 +31,6 @@ current_dir = os.path.dirname(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
)
|
||||
|
||||
|
||||
def generate_minute_range(start_date, end_date):
|
||||
"""Generate a dictionary with minute-level time ranges."""
|
||||
return {
|
||||
@@ -237,8 +239,63 @@ def get_vector_store(source_id):
|
||||
return store
|
||||
|
||||
|
||||
def _validate_agent_image_upload(file, filename: str) -> None:
|
||||
"""Validate an uploaded agent avatar and rewind it for storage.
|
||||
|
||||
Args:
|
||||
file: Werkzeug file upload object.
|
||||
filename: Sanitized upload filename.
|
||||
|
||||
Raises:
|
||||
OSError: If the upload stream cannot be inspected.
|
||||
ValueError: If the file exceeds a limit or is not an allowed raster image.
|
||||
"""
|
||||
extension = PurePosixPath(filename).suffix.lower()
|
||||
image_policy = AGENT_IMAGE_FORMATS.get(extension)
|
||||
if not filename or image_policy is None:
|
||||
raise ValueError("Unsupported image extension")
|
||||
|
||||
stream = file.stream
|
||||
stream.seek(0, os.SEEK_END)
|
||||
size_bytes = stream.tell()
|
||||
stream.seek(0)
|
||||
if size_bytes > settings.AGENT_IMAGE_MAX_BYTES:
|
||||
raise ValueError("Image exceeds the upload size limit")
|
||||
|
||||
with Image.open(stream) as image:
|
||||
image_format = (image.format or "").upper()
|
||||
if image_format != image_policy[0]:
|
||||
raise ValueError("Image content does not match its extension")
|
||||
if image.width * image.height > settings.AGENT_IMAGE_MAX_PIXELS:
|
||||
raise ValueError("Image exceeds the pixel limit")
|
||||
image.verify()
|
||||
stream.seek(0)
|
||||
|
||||
|
||||
def _safe_agent_image_filename(filename: str) -> str:
|
||||
"""Build an ASCII storage filename without discarding a Unicode file's suffix.
|
||||
|
||||
Args:
|
||||
filename: Original multipart upload filename.
|
||||
|
||||
Returns:
|
||||
A traversal-safe filename with a normalized, allow-listed extension.
|
||||
|
||||
Raises:
|
||||
ValueError: If the original filename has no supported image extension.
|
||||
"""
|
||||
normalized_path = filename.replace("\\", "/")
|
||||
original_path = PurePosixPath(normalized_path)
|
||||
extension = original_path.suffix.lower()
|
||||
if extension not in AGENT_IMAGE_FORMATS:
|
||||
raise ValueError("Unsupported image extension")
|
||||
|
||||
safe_stem = secure_filename(original_path.stem).strip("._") or "avatar"
|
||||
return f"{safe_stem}{extension}"
|
||||
|
||||
|
||||
def handle_image_upload(
|
||||
request, existing_url: str, user: str, storage, base_path: str = "attachments/"
|
||||
request, existing_url: str, user: str, storage
|
||||
) -> Tuple[str, Optional[Response]]:
|
||||
"""
|
||||
Handle image file upload from request.
|
||||
@@ -248,7 +305,6 @@ def handle_image_upload(
|
||||
existing_url: Existing image URL (fallback)
|
||||
user: User ID
|
||||
storage: Storage instance
|
||||
base_path: Base path for upload
|
||||
|
||||
Returns:
|
||||
Tuple of (image_url, error_response)
|
||||
@@ -258,8 +314,26 @@ def handle_image_upload(
|
||||
if "image" in request.files:
|
||||
file = request.files["image"]
|
||||
if file.filename != "":
|
||||
filename = secure_filename(file.filename)
|
||||
upload_path = f"{settings.UPLOAD_FOLDER.rstrip('/')}/{user}/{base_path.rstrip('/')}/{uuid.uuid4()}_{filename}"
|
||||
try:
|
||||
filename = _safe_agent_image_filename(file.filename)
|
||||
_validate_agent_image_upload(file, filename)
|
||||
except (
|
||||
Image.DecompressionBombError,
|
||||
UnidentifiedImageError,
|
||||
ValueError,
|
||||
OSError,
|
||||
) as e:
|
||||
current_app.logger.warning(f"Invalid agent image upload: {e}")
|
||||
return None, make_response(
|
||||
jsonify({"success": False, "message": "Invalid image upload"}),
|
||||
400,
|
||||
)
|
||||
|
||||
owner_component = safe_user_storage_component(user)
|
||||
upload_path = (
|
||||
f"{settings.UPLOAD_FOLDER.rstrip('/')}/{owner_component}/"
|
||||
f"attachments/{uuid.uuid4()}_{filename}"
|
||||
)
|
||||
try:
|
||||
storage.save_file(file, upload_path, storage_class="STANDARD")
|
||||
image_url = upload_path
|
||||
|
||||
+8
-12
@@ -28,6 +28,7 @@ from application.api.user.routes import user # noqa: E402
|
||||
from application.api.connector.routes import connector # noqa: E402
|
||||
from application.api.v1 import v1_bp # noqa: E402
|
||||
from application.celery_init import celery # noqa: E402
|
||||
from application.core.secret_key import resolve_jwt_secret_key # noqa: E402
|
||||
from application.core.settings import settings # noqa: E402
|
||||
from application.storage.db.bootstrap import ensure_database_ready # noqa: E402
|
||||
from application.stt.upload_limits import ( # noqa: E402
|
||||
@@ -84,18 +85,13 @@ app.config.update(
|
||||
celery.config_from_object("application.celeryconfig")
|
||||
api.init_app(app)
|
||||
|
||||
if settings.AUTH_TYPE in ("simple_jwt", "session_jwt", "oidc") and not settings.JWT_SECRET_KEY:
|
||||
key_file = ".jwt_secret_key"
|
||||
try:
|
||||
with open(key_file, "r") as f:
|
||||
settings.JWT_SECRET_KEY = f.read().strip()
|
||||
except FileNotFoundError:
|
||||
new_key = os.urandom(32).hex()
|
||||
with open(key_file, "w") as f:
|
||||
f.write(new_key)
|
||||
settings.JWT_SECRET_KEY = new_key
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to setup JWT_SECRET_KEY: {e}")
|
||||
# The same stable secret also signs opaque agent-image capabilities, including
|
||||
# in no-auth mode. Production replicas must receive one shared configured key;
|
||||
# only local development may use the atomic filesystem fallback.
|
||||
settings.JWT_SECRET_KEY = resolve_jwt_secret_key(
|
||||
settings.JWT_SECRET_KEY,
|
||||
os.getenv("DEPLOYMENT_TYPE"),
|
||||
)
|
||||
if settings.AUTH_TYPE == "oidc":
|
||||
_missing_oidc = [
|
||||
name
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Resolve stable application signing secrets for local and production use."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
_SHARED_SECRET_DEPLOYMENTS = {"cloud", "production"}
|
||||
|
||||
|
||||
def _read_secret_file(key_path: Path) -> str:
|
||||
"""Read and validate a generated local signing secret."""
|
||||
secret = key_path.read_text(encoding="utf-8").strip()
|
||||
if not secret:
|
||||
raise RuntimeError(f"Signing secret file is empty: {key_path}")
|
||||
return secret
|
||||
|
||||
|
||||
def _create_secret_file_atomically(key_path: Path) -> str:
|
||||
"""Create a complete mode-0600 secret before atomically publishing it."""
|
||||
new_secret = os.urandom(32).hex()
|
||||
descriptor, temporary_name = tempfile.mkstemp(
|
||||
dir=key_path.parent,
|
||||
prefix=f".{key_path.name}.",
|
||||
text=True,
|
||||
)
|
||||
temporary_path = Path(temporary_name)
|
||||
try:
|
||||
if hasattr(os, "fchmod"):
|
||||
os.fchmod(descriptor, 0o600)
|
||||
with os.fdopen(descriptor, "w", encoding="utf-8") as temporary_file:
|
||||
descriptor = -1
|
||||
temporary_file.write(new_secret)
|
||||
temporary_file.flush()
|
||||
os.fsync(temporary_file.fileno())
|
||||
|
||||
try:
|
||||
# A hard link publishes the fully written file without replacing a
|
||||
# secret another process may have won the race to create.
|
||||
os.link(temporary_path, key_path)
|
||||
except FileExistsError:
|
||||
return _read_secret_file(key_path)
|
||||
return new_secret
|
||||
finally:
|
||||
if descriptor >= 0:
|
||||
os.close(descriptor)
|
||||
temporary_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def resolve_jwt_secret_key(
|
||||
configured_secret: str | None,
|
||||
deployment_type: str | None,
|
||||
key_file: str | Path = ".jwt_secret_key",
|
||||
) -> str:
|
||||
"""Return a shared configured secret or a stable local-development secret.
|
||||
|
||||
Args:
|
||||
configured_secret: Operator-supplied signing secret.
|
||||
deployment_type: Deployment class, such as ``cloud`` or ``production``.
|
||||
key_file: Local fallback file used outside production deployments.
|
||||
|
||||
Returns:
|
||||
The configured or locally persisted signing secret.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If production lacks a shared secret or local persistence fails.
|
||||
"""
|
||||
if configured_secret and configured_secret.strip():
|
||||
return configured_secret
|
||||
|
||||
normalized_deployment = (deployment_type or "").strip().lower()
|
||||
if normalized_deployment in _SHARED_SECRET_DEPLOYMENTS:
|
||||
raise RuntimeError(
|
||||
"JWT_SECRET_KEY must be set to the same strong random value on every "
|
||||
f"{normalized_deployment} API and worker replica"
|
||||
)
|
||||
|
||||
key_path = Path(key_file)
|
||||
try:
|
||||
return _read_secret_file(key_path)
|
||||
except FileNotFoundError:
|
||||
try:
|
||||
return _create_secret_file_atomically(key_path)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to create signing secret: {exc}") from exc
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to read signing secret: {exc}") from exc
|
||||
@@ -99,6 +99,8 @@ class Settings(BaseSettings):
|
||||
DOCLING_TABULAR_MAX_BYTES: int = 2_000_000
|
||||
DOCLING_MARKUP_MAX_BYTES: int = 8_000_000
|
||||
ATTACHMENT_TEXT_MAX_BYTES: int = 5_000_000
|
||||
AGENT_IMAGE_MAX_BYTES: int = 5_000_000
|
||||
AGENT_IMAGE_MAX_PIXELS: int = 16_777_216
|
||||
VECTOR_STORE: str = "faiss" # "faiss" or "elasticsearch" or "qdrant" or "milvus" or "lancedb" or "pgvector"
|
||||
# Allow-list of retriever keys an agent may use. Values must match the
|
||||
# ``RetrieverCreator.retrievers`` registry keys (``classic`` / ``default``),
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Base storage class for file system abstraction."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import BinaryIO, List, Callable
|
||||
from typing import BinaryIO, Callable, List, Optional
|
||||
|
||||
|
||||
class BaseStorage(ABC):
|
||||
@@ -37,12 +37,41 @@ class BaseStorage(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
def generate_presigned_url(self, path: str, expires_in: int = 300) -> str:
|
||||
def get_file_size(self, path: str) -> int:
|
||||
"""Return a file's encoded size without retaining its bytes in memory.
|
||||
|
||||
Storage backends should override this with a metadata-only lookup. The
|
||||
compatibility fallback seeks an already-open file and is suitable for
|
||||
file-like backends that do not provide object metadata.
|
||||
|
||||
Args:
|
||||
path: Path to the file.
|
||||
|
||||
Returns:
|
||||
File size in bytes.
|
||||
"""
|
||||
file_obj = self.get_file(path)
|
||||
try:
|
||||
current = file_obj.tell()
|
||||
file_obj.seek(0, 2)
|
||||
size = file_obj.tell()
|
||||
file_obj.seek(current)
|
||||
return size
|
||||
finally:
|
||||
file_obj.close()
|
||||
|
||||
def generate_presigned_url(
|
||||
self,
|
||||
path: str,
|
||||
expires_in: int = 300,
|
||||
content_type: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Return a short-lived presigned download URL; not all backends support it.
|
||||
|
||||
Args:
|
||||
path: Path to the file
|
||||
expires_in: TTL of the signed URL in seconds
|
||||
content_type: Optional response Content-Type override
|
||||
|
||||
Returns:
|
||||
str: A presigned URL granting time-limited read access
|
||||
|
||||
@@ -176,6 +176,25 @@ class AgentsRepository:
|
||||
row = result.fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def find_image_record(self, agent_id: str) -> Optional[dict]:
|
||||
"""Fetch only the fields needed to verify a public image capability.
|
||||
|
||||
This lookup is intentionally not owner-scoped because the resulting
|
||||
row is not authorization: callers must verify its HMAC capability
|
||||
before reading or returning any storage bytes.
|
||||
"""
|
||||
if not looks_like_uuid(agent_id):
|
||||
return None
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT id, user_id, image FROM agents "
|
||||
"WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": agent_id},
|
||||
)
|
||||
row = result.fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def list_by_ids(self, agent_ids) -> list[dict]:
|
||||
"""Fetch agents whose id is in ``agent_ids`` (team-shared listing path)."""
|
||||
ids = [str(a) for a in agent_ids if looks_like_uuid(str(a))]
|
||||
|
||||
@@ -60,6 +60,14 @@ class LocalStorage(BaseStorage):
|
||||
|
||||
return open(full_path, 'rb')
|
||||
|
||||
def get_file_size(self, path: str) -> int:
|
||||
"""Return the size of a local file without opening and buffering it."""
|
||||
full_path = self._get_full_path(path)
|
||||
try:
|
||||
return os.path.getsize(full_path)
|
||||
except FileNotFoundError as exc:
|
||||
raise FileNotFoundError(f"File not found: {full_path}") from exc
|
||||
|
||||
def delete_file(self, path: str) -> bool:
|
||||
"""Delete a file from local storage."""
|
||||
full_path = self._get_full_path(path)
|
||||
|
||||
@@ -121,13 +121,33 @@ class S3Storage(BaseStorage):
|
||||
file_obj.seek(0)
|
||||
return file_obj
|
||||
|
||||
def generate_presigned_url(self, path: str, expires_in: int = 300) -> str:
|
||||
def get_file_size(self, path: str) -> int:
|
||||
"""Return an S3 object's size using HEAD without downloading its body."""
|
||||
path = self._validate_path(path)
|
||||
try:
|
||||
metadata = self.s3.head_object(Bucket=self.bucket_name, Key=path)
|
||||
except ClientError as exc:
|
||||
error_code = str(exc.response.get("Error", {}).get("Code", ""))
|
||||
if error_code in {"404", "NoSuchKey", "NotFound"}:
|
||||
raise FileNotFoundError(f"File not found: {path}") from exc
|
||||
raise
|
||||
return int(metadata["ContentLength"])
|
||||
|
||||
def generate_presigned_url(
|
||||
self,
|
||||
path: str,
|
||||
expires_in: int = 300,
|
||||
content_type: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Return a short-lived presigned GET URL for a private object (TTL <= 1h)."""
|
||||
path = self._validate_path(path)
|
||||
expires_in = min(expires_in, 3600)
|
||||
params = {"Bucket": self.bucket_name, "Key": path}
|
||||
if content_type:
|
||||
params["ResponseContentType"] = content_type
|
||||
return self.s3.generate_presigned_url(
|
||||
"get_object",
|
||||
Params={"Bucket": self.bucket_name, "Key": path},
|
||||
Params=params,
|
||||
ExpiresIn=expires_in,
|
||||
)
|
||||
|
||||
|
||||
+109
-22
@@ -1,12 +1,13 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from pathlib import PurePosixPath
|
||||
from typing import List
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import tiktoken
|
||||
from flask import jsonify, make_response
|
||||
@@ -237,31 +238,117 @@ def validate_function_name(function_name):
|
||||
return True
|
||||
|
||||
|
||||
def generate_image_url(image_path):
|
||||
AGENT_IMAGE_FORMATS = {
|
||||
".gif": ("GIF", "image/gif"),
|
||||
".jpeg": ("JPEG", "image/jpeg"),
|
||||
".jpg": ("JPEG", "image/jpeg"),
|
||||
".png": ("PNG", "image/png"),
|
||||
".webp": ("WEBP", "image/webp"),
|
||||
}
|
||||
|
||||
|
||||
def is_external_image_url(image_path: object) -> bool:
|
||||
"""Return whether an image value is an externally hosted HTTP(S) URL."""
|
||||
return isinstance(image_path, str) and image_path.startswith(("http://", "https://"))
|
||||
|
||||
|
||||
def safe_user_storage_component(user_id: object) -> str:
|
||||
"""Return a deterministic, traversal-safe directory component for a user ID."""
|
||||
raw_user_id = str(user_id or "")
|
||||
component = secure_filename(raw_user_id)[:80] or "user"
|
||||
digest = hashlib.sha256(raw_user_id.encode("utf-8")).hexdigest()[:16]
|
||||
return f"{component}-{digest}"
|
||||
|
||||
|
||||
def get_agent_image_content_type(image_path: object) -> str | None:
|
||||
"""Return an allow-listed raster MIME type for an internal image path."""
|
||||
if not isinstance(image_path, str):
|
||||
return None
|
||||
policy = AGENT_IMAGE_FORMATS.get(PurePosixPath(image_path).suffix.lower())
|
||||
return policy[1] if policy else None
|
||||
|
||||
|
||||
def is_safe_agent_image_path(image_path: object, user_id: object) -> bool:
|
||||
"""Validate that a path is an agent avatar under its owner's upload directory."""
|
||||
if not isinstance(image_path, str) or not image_path or not user_id:
|
||||
return False
|
||||
if is_external_image_url(image_path) or "\\" in image_path or "\x00" in image_path:
|
||||
return False
|
||||
|
||||
candidate = PurePosixPath(image_path)
|
||||
upload_root = PurePosixPath(str(settings.UPLOAD_FOLDER).rstrip("/"))
|
||||
if candidate.is_absolute() or upload_root.is_absolute() or ".." in candidate.parts:
|
||||
return False
|
||||
if get_agent_image_content_type(image_path) is None:
|
||||
return False
|
||||
|
||||
owner_components = {safe_user_storage_component(user_id)}
|
||||
raw_user_id = str(user_id)
|
||||
raw_owner = PurePosixPath(raw_user_id)
|
||||
if (
|
||||
len(raw_owner.parts) == 1
|
||||
and raw_user_id not in {"", ".", ".."}
|
||||
and "\\" not in raw_user_id
|
||||
and "\x00" not in raw_user_id
|
||||
):
|
||||
# Compatibility for avatars written before user IDs were sanitized.
|
||||
owner_components.add(raw_user_id)
|
||||
|
||||
for owner_component in owner_components:
|
||||
expected_parent = upload_root / owner_component / "attachments"
|
||||
if candidate.parent == expected_parent and candidate.name:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def generate_agent_image_capability(
|
||||
agent_id: object, image_path: object, user_id: object
|
||||
) -> str:
|
||||
"""Create an HMAC capability for one agent's current internal image."""
|
||||
secret = getattr(settings, "JWT_SECRET_KEY", "")
|
||||
if not isinstance(secret, str) or not secret:
|
||||
return ""
|
||||
try:
|
||||
canonical_agent_id = str(uuid.UUID(str(agent_id)))
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
return ""
|
||||
if not isinstance(image_path, str) or not user_id:
|
||||
return ""
|
||||
payload = (
|
||||
f"docsgpt-agent-image-v1\0{canonical_agent_id}\0{user_id}\0{image_path}"
|
||||
).encode("utf-8")
|
||||
return hmac.new(secret.encode("utf-8"), payload, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
def verify_agent_image_capability(
|
||||
capability: object, agent_id: object, image_path: object, user_id: object
|
||||
) -> bool:
|
||||
"""Verify an agent image capability without timing-leaky string comparison."""
|
||||
if not isinstance(capability, str) or not re.fullmatch(r"[0-9a-f]{64}", capability):
|
||||
return False
|
||||
expected = generate_agent_image_capability(agent_id, image_path, user_id)
|
||||
return bool(expected) and hmac.compare_digest(capability, expected)
|
||||
|
||||
|
||||
def generate_image_url(image_path, agent_id=None, user_id=None):
|
||||
"""Return an external URL or an opaque capability URL for an agent image.
|
||||
|
||||
Internal storage paths are never included in the returned URL. Invalid or
|
||||
unscoped paths fail closed so a poisoned database row cannot become a file
|
||||
read capability.
|
||||
"""
|
||||
if isinstance(image_path, str) and (
|
||||
image_path.startswith("http://") or image_path.startswith("https://")
|
||||
):
|
||||
return image_path
|
||||
strategy = getattr(settings, "URL_STRATEGY", "backend")
|
||||
if strategy == "s3":
|
||||
bucket_name = settings.S3_BUCKET_NAME
|
||||
endpoint_url = settings.S3_ENDPOINT_URL
|
||||
if endpoint_url:
|
||||
# S3-compatible service (MinIO, R2, B2, Spaces, ...).
|
||||
base = endpoint_url.rstrip("/")
|
||||
if settings.S3_PATH_STYLE:
|
||||
return f"{base}/{bucket_name}/{image_path}"
|
||||
parsed = urlparse(base)
|
||||
return f"{parsed.scheme}://{bucket_name}.{parsed.netloc}/{image_path}"
|
||||
region_name = (
|
||||
settings.S3_REGION
|
||||
or getattr(settings, "SAGEMAKER_REGION", None)
|
||||
or "eu-central-1"
|
||||
)
|
||||
return f"https://{bucket_name}.s3.{region_name}.amazonaws.com/{image_path}"
|
||||
else:
|
||||
base_url = getattr(settings, "API_URL", "http://localhost:7091")
|
||||
return f"{base_url}/api/images/{image_path}"
|
||||
if not is_safe_agent_image_path(image_path, user_id):
|
||||
return ""
|
||||
capability = generate_agent_image_capability(agent_id, image_path, user_id)
|
||||
if not capability:
|
||||
return ""
|
||||
canonical_agent_id = str(uuid.UUID(str(agent_id)))
|
||||
base_url = getattr(settings, "API_URL", "http://localhost:7091").rstrip("/")
|
||||
return f"{base_url}/api/images/{canonical_agent_id}/{capability}"
|
||||
|
||||
|
||||
def calculate_compression_threshold(
|
||||
|
||||
Reference in new issue
Block a user