dify/api/controllers/console/remote_files.py

import urllib.parse

import httpx
from pydantic import BaseModel, Field

import services
from controllers.common import helpers
from controllers.common.errors import (
    FileTooLargeError,
    RemoteFileUploadError,
    UnsupportedFileTypeError,
)
from controllers.fastopenapi import console_router
from core.file import helpers as file_helpers
from core.helper import ssrf_proxy
from extensions.ext_database import db
from fields.file_fields import FileWithSignedUrl, RemoteFileInfo
from libs.login import current_account_with_tenant
from services.file_service import FileService


class RemoteFileUploadPayload(BaseModel):
    url: str = Field(..., description="URL to fetch")


@console_router.get(
    "/remote-files/<path:url>",
    response_model=RemoteFileInfo,
    tags=["console"],
)
def get_remote_file_info(url: str) -> RemoteFileInfo:
    decoded_url = urllib.parse.unquote(url)
    resp = ssrf_proxy.head(decoded_url)
    if resp.status_code != httpx.codes.OK:
        resp = ssrf_proxy.get(decoded_url, timeout=3)
    resp.raise_for_status()
    return RemoteFileInfo(
        file_type=resp.headers.get("Content-Type", "application/octet-stream"),
        file_length=int(resp.headers.get("Content-Length", 0)),
    )


@console_router.post(
    "/remote-files/upload",
    response_model=FileWithSignedUrl,
    tags=["console"],
    status_code=201,
)
def upload_remote_file(payload: RemoteFileUploadPayload) -> FileWithSignedUrl:
    url = payload.url

    try:
        resp = ssrf_proxy.head(url=url)
        if resp.status_code != httpx.codes.OK:
            resp = ssrf_proxy.get(url=url, timeout=3, follow_redirects=True)
        if resp.status_code != httpx.codes.OK:
            raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}")
    except httpx.RequestError as e:
        raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(e)}")

    file_info = helpers.guess_file_info_from_response(resp)

    if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size):
        raise FileTooLargeError

    content = resp.content if resp.request.method == "GET" else ssrf_proxy.get(url).content

    try:
        user, _ = current_account_with_tenant()
        upload_file = FileService(db.engine).upload_file(
            filename=file_info.filename,
            content=content,
            mimetype=file_info.mimetype,
            user=user,
            source_url=url,
        )
    except services.errors.file.FileTooLargeError as file_too_large_error:
        raise FileTooLargeError(file_too_large_error.description)
    except services.errors.file.UnsupportedFileTypeError:
        raise UnsupportedFileTypeError()

    return FileWithSignedUrl(
        id=upload_file.id,
        name=upload_file.name,
        size=upload_file.size,
        extension=upload_file.extension,
        url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
        mime_type=upload_file.mime_type,
        created_by=upload_file.created_by,
        created_at=int(upload_file.created_at.timestamp()),
    )