mirror of
https://github.com/langgenius/dify.git
synced 2026-05-03 17:08:03 +08:00
refactor(storage): remove signer, using general file storage
- Removed unused app asset download and upload endpoints, along with sandbox archive and file download endpoints. - Updated imports in the file controller to reflect the removal of these endpoints. - Simplified the generator.py file by consolidating the code context field definition. - Enhanced the storage layer with a unified presign wrapper for better handling of presigned URLs.
This commit is contained in:
Harry
@ -1,25 +1,31 @@
|
||||
"""App assets storage layer.
|
||||
|
||||
This module provides storage abstractions for app assets (draft files, build zips,
|
||||
resolved assets, skill bundles, source zips, bundle exports/imports).
|
||||
|
||||
Key components:
|
||||
- AssetPath: Factory for creating typed storage paths
|
||||
- AppAssetStorage: High-level storage operations with presign support
|
||||
|
||||
All presign operations use the unified FilePresignStorage wrapper, which automatically
|
||||
falls back to Dify's file proxy when the underlying storage doesn't support presigned URLs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
import time
|
||||
import urllib.parse
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Callable, Iterable
|
||||
from collections.abc import Generator, Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, ClassVar
|
||||
from uuid import UUID
|
||||
|
||||
from configs import dify_config
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
from extensions.storage.cached_presign_storage import CachedPresignStorage
|
||||
from libs import rsa
|
||||
from extensions.storage.file_presign_storage import FilePresignStorage
|
||||
|
||||
_ASSET_BASE = "app_assets"
|
||||
_SILENT_STORAGE_NOT_FOUND = b"File Not Found"
|
||||
_ASSET_PATH_REGISTRY: dict[str, tuple[bool, Callable[..., SignedAssetPath]]] = {}
|
||||
_ASSET_PATH_REGISTRY: dict[str, tuple[bool, Any]] = {}
|
||||
|
||||
|
||||
def _require_uuid(value: str, field_name: str) -> None:
|
||||
@ -29,12 +35,14 @@ def _require_uuid(value: str, field_name: str) -> None:
|
||||
raise ValueError(f"{field_name} must be a UUID") from exc
|
||||
|
||||
|
||||
def register_asset_path(asset_type: str, *, requires_node: bool, factory: Callable[..., SignedAssetPath]) -> None:
|
||||
def register_asset_path(asset_type: str, *, requires_node: bool, factory: Any) -> None:
|
||||
_ASSET_PATH_REGISTRY[asset_type] = (requires_node, factory)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AssetPathBase(ABC):
|
||||
"""Base class for all asset paths."""
|
||||
|
||||
asset_type: ClassVar[str]
|
||||
tenant_id: str
|
||||
app_id: str
|
||||
@ -50,49 +58,24 @@ class AssetPathBase(ABC):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class SignedAssetPath(AssetPathBase, ABC):
|
||||
@abstractmethod
|
||||
def signature_parts(self) -> tuple[str, str | None]:
|
||||
"""Return (resource_id, sub_resource_id) used for signing.
|
||||
|
||||
sub_resource_id should be None when not applicable.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def proxy_path_parts(self) -> list[str]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _DraftAssetPath(SignedAssetPath):
|
||||
class _DraftAssetPath(AssetPathBase):
|
||||
asset_type: ClassVar[str] = "draft"
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"{_ASSET_BASE}/{self.tenant_id}/{self.app_id}/draft/{self.resource_id}"
|
||||
|
||||
def signature_parts(self) -> tuple[str, str | None]:
|
||||
return (self.resource_id, None)
|
||||
|
||||
def proxy_path_parts(self) -> list[str]:
|
||||
return [self.asset_type, self.tenant_id, self.app_id, self.resource_id]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _BuildZipAssetPath(SignedAssetPath):
|
||||
class _BuildZipAssetPath(AssetPathBase):
|
||||
asset_type: ClassVar[str] = "build-zip"
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"{_ASSET_BASE}/{self.tenant_id}/{self.app_id}/artifacts/{self.resource_id}.zip"
|
||||
|
||||
def signature_parts(self) -> tuple[str, str | None]:
|
||||
return (self.resource_id, None)
|
||||
|
||||
def proxy_path_parts(self) -> list[str]:
|
||||
return [self.asset_type, self.tenant_id, self.app_id, self.resource_id]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ResolvedAssetPath(SignedAssetPath):
|
||||
class _ResolvedAssetPath(AssetPathBase):
|
||||
asset_type: ClassVar[str] = "resolved"
|
||||
node_id: str
|
||||
|
||||
@ -103,58 +86,34 @@ class _ResolvedAssetPath(SignedAssetPath):
|
||||
def get_storage_key(self) -> str:
|
||||
return f"{_ASSET_BASE}/{self.tenant_id}/{self.app_id}/artifacts/{self.resource_id}/resolved/{self.node_id}"
|
||||
|
||||
def signature_parts(self) -> tuple[str, str | None]:
|
||||
return (self.resource_id, self.node_id)
|
||||
|
||||
def proxy_path_parts(self) -> list[str]:
|
||||
return [self.asset_type, self.tenant_id, self.app_id, self.resource_id, self.node_id]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _SkillBundleAssetPath(SignedAssetPath):
|
||||
class _SkillBundleAssetPath(AssetPathBase):
|
||||
asset_type: ClassVar[str] = "skill-bundle"
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"{_ASSET_BASE}/{self.tenant_id}/{self.app_id}/artifacts/{self.resource_id}/skill_artifact_set.json"
|
||||
|
||||
def signature_parts(self) -> tuple[str, str | None]:
|
||||
return (self.resource_id, None)
|
||||
|
||||
def proxy_path_parts(self) -> list[str]:
|
||||
return [self.asset_type, self.tenant_id, self.app_id, self.resource_id]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _SourceZipAssetPath(SignedAssetPath):
|
||||
class _SourceZipAssetPath(AssetPathBase):
|
||||
asset_type: ClassVar[str] = "source-zip"
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"{_ASSET_BASE}/{self.tenant_id}/{self.app_id}/sources/{self.resource_id}.zip"
|
||||
|
||||
def signature_parts(self) -> tuple[str, str | None]:
|
||||
return (self.resource_id, None)
|
||||
|
||||
def proxy_path_parts(self) -> list[str]:
|
||||
return [self.asset_type, self.tenant_id, self.app_id, self.resource_id]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _BundleExportZipAssetPath(SignedAssetPath):
|
||||
class _BundleExportZipAssetPath(AssetPathBase):
|
||||
asset_type: ClassVar[str] = "bundle-export-zip"
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"{_ASSET_BASE}/{self.tenant_id}/{self.app_id}/bundle_exports/{self.resource_id}.zip"
|
||||
|
||||
def signature_parts(self) -> tuple[str, str | None]:
|
||||
return (self.resource_id, None)
|
||||
|
||||
def proxy_path_parts(self) -> list[str]:
|
||||
return [self.asset_type, self.tenant_id, self.app_id, self.resource_id]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BundleImportZipPath:
|
||||
"""Path for temporary import zip files. Not signed, uses direct presign URLs only."""
|
||||
"""Path for temporary import zip files."""
|
||||
|
||||
tenant_id: str
|
||||
import_id: str
|
||||
@ -167,28 +126,30 @@ class BundleImportZipPath:
|
||||
|
||||
|
||||
class AssetPath:
|
||||
"""Factory for creating typed asset paths."""
|
||||
|
||||
@staticmethod
|
||||
def draft(tenant_id: str, app_id: str, node_id: str) -> SignedAssetPath:
|
||||
def draft(tenant_id: str, app_id: str, node_id: str) -> AssetPathBase:
|
||||
return _DraftAssetPath(tenant_id=tenant_id, app_id=app_id, resource_id=node_id)
|
||||
|
||||
@staticmethod
|
||||
def build_zip(tenant_id: str, app_id: str, assets_id: str) -> SignedAssetPath:
|
||||
def build_zip(tenant_id: str, app_id: str, assets_id: str) -> AssetPathBase:
|
||||
return _BuildZipAssetPath(tenant_id=tenant_id, app_id=app_id, resource_id=assets_id)
|
||||
|
||||
@staticmethod
|
||||
def resolved(tenant_id: str, app_id: str, assets_id: str, node_id: str) -> SignedAssetPath:
|
||||
def resolved(tenant_id: str, app_id: str, assets_id: str, node_id: str) -> AssetPathBase:
|
||||
return _ResolvedAssetPath(tenant_id=tenant_id, app_id=app_id, resource_id=assets_id, node_id=node_id)
|
||||
|
||||
@staticmethod
|
||||
def skill_bundle(tenant_id: str, app_id: str, assets_id: str) -> SignedAssetPath:
|
||||
def skill_bundle(tenant_id: str, app_id: str, assets_id: str) -> AssetPathBase:
|
||||
return _SkillBundleAssetPath(tenant_id=tenant_id, app_id=app_id, resource_id=assets_id)
|
||||
|
||||
@staticmethod
|
||||
def source_zip(tenant_id: str, app_id: str, workflow_id: str) -> SignedAssetPath:
|
||||
def source_zip(tenant_id: str, app_id: str, workflow_id: str) -> AssetPathBase:
|
||||
return _SourceZipAssetPath(tenant_id=tenant_id, app_id=app_id, resource_id=workflow_id)
|
||||
|
||||
@staticmethod
|
||||
def bundle_export_zip(tenant_id: str, app_id: str, export_id: str) -> SignedAssetPath:
|
||||
def bundle_export_zip(tenant_id: str, app_id: str, export_id: str) -> AssetPathBase:
|
||||
return _BundleExportZipAssetPath(tenant_id=tenant_id, app_id=app_id, resource_id=export_id)
|
||||
|
||||
@staticmethod
|
||||
@ -202,7 +163,7 @@ class AssetPath:
|
||||
app_id: str,
|
||||
resource_id: str,
|
||||
sub_resource_id: str | None = None,
|
||||
) -> SignedAssetPath:
|
||||
) -> AssetPathBase:
|
||||
entry = _ASSET_PATH_REGISTRY.get(asset_type)
|
||||
if not entry:
|
||||
raise ValueError(f"Unsupported asset type: {asset_type}")
|
||||
@ -224,120 +185,26 @@ register_asset_path("source-zip", requires_node=False, factory=AssetPath.source_
|
||||
register_asset_path("bundle-export-zip", requires_node=False, factory=AssetPath.bundle_export_zip)
|
||||
|
||||
|
||||
class AppAssetSigner:
|
||||
SIGNATURE_PREFIX = "app-asset"
|
||||
SIGNATURE_VERSION = "v1"
|
||||
OPERATION_DOWNLOAD = "download"
|
||||
OPERATION_UPLOAD = "upload"
|
||||
|
||||
@classmethod
|
||||
def create_download_signature(cls, asset_path: SignedAssetPath, expires_at: int, nonce: str) -> str:
|
||||
return cls._create_signature(
|
||||
asset_path=asset_path,
|
||||
operation=cls.OPERATION_DOWNLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def create_upload_signature(cls, asset_path: SignedAssetPath, expires_at: int, nonce: str) -> str:
|
||||
return cls._create_signature(
|
||||
asset_path=asset_path,
|
||||
operation=cls.OPERATION_UPLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def verify_download_signature(cls, asset_path: SignedAssetPath, expires_at: int, nonce: str, sign: str) -> bool:
|
||||
return cls._verify_signature(
|
||||
asset_path=asset_path,
|
||||
operation=cls.OPERATION_DOWNLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
sign=sign,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def verify_upload_signature(cls, asset_path: SignedAssetPath, expires_at: int, nonce: str, sign: str) -> bool:
|
||||
return cls._verify_signature(
|
||||
asset_path=asset_path,
|
||||
operation=cls.OPERATION_UPLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
sign=sign,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _verify_signature(
|
||||
cls,
|
||||
*,
|
||||
asset_path: SignedAssetPath,
|
||||
operation: str,
|
||||
expires_at: int,
|
||||
nonce: str,
|
||||
sign: str,
|
||||
) -> bool:
|
||||
if expires_at <= 0:
|
||||
return False
|
||||
|
||||
expected_sign = cls._create_signature(
|
||||
asset_path=asset_path,
|
||||
operation=operation,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
if not hmac.compare_digest(sign, expected_sign):
|
||||
return False
|
||||
|
||||
current_time = int(time.time())
|
||||
if expires_at < current_time:
|
||||
return False
|
||||
|
||||
if expires_at - current_time > dify_config.FILES_ACCESS_TIMEOUT:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def _create_signature(cls, *, asset_path: SignedAssetPath, operation: str, expires_at: int, nonce: str) -> str:
|
||||
key = cls._tenant_key(asset_path.tenant_id)
|
||||
message = cls._signature_message(
|
||||
asset_path=asset_path,
|
||||
operation=operation,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
sign = hmac.new(key, message.encode(), hashlib.sha256).digest()
|
||||
return base64.urlsafe_b64encode(sign).decode()
|
||||
|
||||
@classmethod
|
||||
def _signature_message(cls, *, asset_path: SignedAssetPath, operation: str, expires_at: int, nonce: str) -> str:
|
||||
resource_id, sub_resource_id = asset_path.signature_parts()
|
||||
return (
|
||||
f"{cls.SIGNATURE_PREFIX}|{cls.SIGNATURE_VERSION}|{operation}|"
|
||||
f"{asset_path.asset_type}|{asset_path.tenant_id}|{asset_path.app_id}|"
|
||||
f"{resource_id}|{sub_resource_id or ''}|{expires_at}|{nonce}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _tenant_key(cls, tenant_id: str) -> bytes:
|
||||
try:
|
||||
rsa_key, _ = rsa.get_decrypt_decoding(tenant_id)
|
||||
except rsa.PrivkeyNotFoundError as exc:
|
||||
raise ValueError(f"Tenant private key missing for tenant_id={tenant_id}") from exc
|
||||
private_key = rsa_key.export_key()
|
||||
return hashlib.sha256(private_key).digest()
|
||||
|
||||
|
||||
class AppAssetStorage:
|
||||
_base_storage: BaseStorage
|
||||
"""High-level storage operations for app assets.
|
||||
|
||||
Wraps BaseStorage with:
|
||||
- FilePresignStorage for presign fallback support
|
||||
- CachedPresignStorage for URL caching
|
||||
|
||||
Usage:
|
||||
storage = AppAssetStorage(base_storage, redis_client=redis)
|
||||
storage.save(asset_path, content)
|
||||
url = storage.get_download_url(asset_path)
|
||||
"""
|
||||
|
||||
_storage: CachedPresignStorage
|
||||
|
||||
def __init__(self, storage: BaseStorage, *, redis_client: Any, cache_key_prefix: str = "app_assets") -> None:
|
||||
self._base_storage = storage
|
||||
# Wrap with FilePresignStorage for fallback support, then CachedPresignStorage for caching
|
||||
presign_storage = FilePresignStorage(storage)
|
||||
self._storage = CachedPresignStorage(
|
||||
storage=storage,
|
||||
storage=presign_storage,
|
||||
redis_client=redis_client,
|
||||
cache_key_prefix=cache_key_prefix,
|
||||
)
|
||||
@ -347,69 +214,44 @@ class AppAssetStorage:
|
||||
return self._storage
|
||||
|
||||
def save(self, asset_path: AssetPathBase, content: bytes) -> None:
|
||||
self._storage.save(self.get_storage_key(asset_path), content)
|
||||
self._storage.save(asset_path.get_storage_key(), content)
|
||||
|
||||
def load(self, asset_path: AssetPathBase) -> bytes:
|
||||
return self._storage.load_once(self.get_storage_key(asset_path))
|
||||
return self._storage.load_once(asset_path.get_storage_key())
|
||||
|
||||
def load_stream(self, asset_path: AssetPathBase) -> Generator[bytes, None, None]:
|
||||
return self._storage.load_stream(asset_path.get_storage_key())
|
||||
|
||||
def load_or_none(self, asset_path: AssetPathBase) -> bytes | None:
|
||||
try:
|
||||
data = self._storage.load_once(self.get_storage_key(asset_path))
|
||||
data = self._storage.load_once(asset_path.get_storage_key())
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
if data == _SILENT_STORAGE_NOT_FOUND:
|
||||
return None
|
||||
return data
|
||||
|
||||
def exists(self, asset_path: AssetPathBase) -> bool:
|
||||
return self._storage.exists(asset_path.get_storage_key())
|
||||
|
||||
def delete(self, asset_path: AssetPathBase) -> None:
|
||||
self._storage.delete(self.get_storage_key(asset_path))
|
||||
self._storage.delete(asset_path.get_storage_key())
|
||||
|
||||
def get_storage_key(self, asset_path: AssetPathBase) -> str:
|
||||
return asset_path.get_storage_key()
|
||||
def get_download_url(self, asset_path: AssetPathBase, expires_in: int = 3600) -> str:
|
||||
return self._storage.get_download_url(asset_path.get_storage_key(), expires_in)
|
||||
|
||||
def get_download_url(self, asset_path: SignedAssetPath, expires_in: int = 3600) -> str:
|
||||
storage_key = self.get_storage_key(asset_path)
|
||||
try:
|
||||
return self._storage.get_download_url(storage_key, expires_in)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
def get_download_urls(self, asset_paths: Iterable[AssetPathBase], expires_in: int = 3600) -> list[str]:
|
||||
storage_keys = [p.get_storage_key() for p in asset_paths]
|
||||
return self._storage.get_download_urls(storage_keys, expires_in)
|
||||
|
||||
return self._generate_signed_proxy_download_url(asset_path, expires_in)
|
||||
|
||||
def get_download_urls(
|
||||
self,
|
||||
asset_paths: Iterable[SignedAssetPath],
|
||||
expires_in: int = 3600,
|
||||
) -> list[str]:
|
||||
asset_paths_list = list(asset_paths)
|
||||
storage_keys = [self.get_storage_key(asset_path) for asset_path in asset_paths_list]
|
||||
|
||||
try:
|
||||
return self._storage.get_download_urls(storage_keys, expires_in)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
|
||||
return [self._generate_signed_proxy_download_url(asset_path, expires_in) for asset_path in asset_paths_list]
|
||||
|
||||
def get_upload_url(
|
||||
self,
|
||||
asset_path: SignedAssetPath,
|
||||
expires_in: int = 3600,
|
||||
) -> str:
|
||||
storage_key = self.get_storage_key(asset_path)
|
||||
try:
|
||||
return self._storage.get_upload_url(storage_key, expires_in)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
|
||||
return self._generate_signed_proxy_upload_url(asset_path, expires_in)
|
||||
def get_upload_url(self, asset_path: AssetPathBase, expires_in: int = 3600) -> str:
|
||||
return self._storage.get_upload_url(asset_path.get_storage_key(), expires_in)
|
||||
|
||||
# Bundle import convenience methods
|
||||
def get_import_upload_url(self, path: BundleImportZipPath, expires_in: int = 3600) -> str:
|
||||
"""Get upload URL for import zip (direct presign, no proxy fallback)."""
|
||||
return self._storage.get_upload_url(path.get_storage_key(), expires_in)
|
||||
|
||||
def get_import_download_url(self, path: BundleImportZipPath, expires_in: int = 3600) -> str:
|
||||
"""Get download URL for import zip (direct presign, no proxy fallback)."""
|
||||
return self._storage.get_download_url(path.get_storage_key(), expires_in)
|
||||
|
||||
def delete_import_zip(self, path: BundleImportZipPath) -> None:
|
||||
@ -420,31 +262,3 @@ class AppAssetStorage:
|
||||
import logging
|
||||
|
||||
logging.getLogger(__name__).debug("Failed to delete import zip: %s", path.get_storage_key())
|
||||
|
||||
def _generate_signed_proxy_download_url(self, asset_path: SignedAssetPath, expires_in: int) -> str:
|
||||
expires_in = min(expires_in, dify_config.FILES_ACCESS_TIMEOUT)
|
||||
expires_at = int(time.time()) + max(expires_in, 1)
|
||||
nonce = os.urandom(16).hex()
|
||||
sign = AppAssetSigner.create_download_signature(asset_path=asset_path, expires_at=expires_at, nonce=nonce)
|
||||
|
||||
base_url = dify_config.FILES_URL
|
||||
url = self._build_proxy_url(base_url=base_url, asset_path=asset_path, action="download")
|
||||
query = urllib.parse.urlencode({"expires_at": expires_at, "nonce": nonce, "sign": sign})
|
||||
return f"{url}?{query}"
|
||||
|
||||
def _generate_signed_proxy_upload_url(self, asset_path: SignedAssetPath, expires_in: int) -> str:
|
||||
expires_in = min(expires_in, dify_config.FILES_ACCESS_TIMEOUT)
|
||||
expires_at = int(time.time()) + max(expires_in, 1)
|
||||
nonce = os.urandom(16).hex()
|
||||
sign = AppAssetSigner.create_upload_signature(asset_path=asset_path, expires_at=expires_at, nonce=nonce)
|
||||
|
||||
base_url = dify_config.FILES_URL
|
||||
url = self._build_proxy_url(base_url=base_url, asset_path=asset_path, action="upload")
|
||||
query = urllib.parse.urlencode({"expires_at": expires_at, "nonce": nonce, "sign": sign})
|
||||
return f"{url}?{query}"
|
||||
|
||||
@staticmethod
|
||||
def _build_proxy_url(*, base_url: str, asset_path: SignedAssetPath, action: str) -> str:
|
||||
encoded_parts = [urllib.parse.quote(part, safe="") for part in asset_path.proxy_path_parts()]
|
||||
path = "/".join(encoded_parts)
|
||||
return f"{base_url}/files/app-assets/{path}/{action}"
|
||||
|
||||
@ -7,9 +7,9 @@ from uuid import UUID, uuid4
|
||||
|
||||
from core.sandbox.entities.files import SandboxFileDownloadTicket, SandboxFileNode
|
||||
from core.sandbox.inspector.base import SandboxFileSource
|
||||
from core.sandbox.security.archive_signer import SandboxArchivePath, SandboxArchiveSigner
|
||||
from core.sandbox.security.sandbox_file_signer import SandboxFileDownloadPath
|
||||
from core.sandbox.storage import sandbox_file_storage
|
||||
from core.sandbox.storage.archive_storage import SandboxArchivePath
|
||||
from core.sandbox.storage.sandbox_file_storage import SandboxFileDownloadPath
|
||||
from core.virtual_environment.__base.exec import CommandExecutionError
|
||||
from core.virtual_environment.__base.helpers import execute
|
||||
from extensions.ext_storage import storage
|
||||
@ -68,15 +68,14 @@ print(json.dumps(entries))
|
||||
|
||||
def _get_archive_download_url(self) -> str:
|
||||
"""Get a pre-signed download URL for the sandbox archive."""
|
||||
from extensions.storage.file_presign_storage import FilePresignStorage
|
||||
|
||||
archive_path = SandboxArchivePath(tenant_id=UUID(self._tenant_id), sandbox_id=UUID(self._sandbox_id))
|
||||
storage_key = archive_path.get_storage_key()
|
||||
if not storage.exists(storage_key):
|
||||
raise ValueError("Sandbox archive not found")
|
||||
return SandboxArchiveSigner.build_signed_url(
|
||||
archive_path=archive_path,
|
||||
expires_in=self._EXPORT_EXPIRES_IN_SECONDS,
|
||||
action=SandboxArchiveSigner.OPERATION_DOWNLOAD,
|
||||
)
|
||||
presign_storage = FilePresignStorage(storage.storage_runner)
|
||||
return presign_storage.get_download_url(storage_key, self._EXPORT_EXPIRES_IN_SECONDS)
|
||||
|
||||
def _create_zip_sandbox(self) -> ZipSandbox:
|
||||
"""Create a ZipSandbox instance for archive operations."""
|
||||
|
||||
@ -7,8 +7,8 @@ from uuid import UUID, uuid4
|
||||
|
||||
from core.sandbox.entities.files import SandboxFileDownloadTicket, SandboxFileNode
|
||||
from core.sandbox.inspector.base import SandboxFileSource
|
||||
from core.sandbox.security.sandbox_file_signer import SandboxFileDownloadPath
|
||||
from core.sandbox.storage import sandbox_file_storage
|
||||
from core.sandbox.storage.sandbox_file_storage import SandboxFileDownloadPath
|
||||
from core.virtual_environment.__base.exec import CommandExecutionError
|
||||
from core.virtual_environment.__base.helpers import execute
|
||||
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
|
||||
|
||||
@ -1 +0,0 @@
|
||||
"""Sandbox security helpers."""
|
||||
@ -1,152 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
import time
|
||||
import urllib.parse
|
||||
from dataclasses import dataclass
|
||||
from uuid import UUID
|
||||
|
||||
from configs import dify_config
|
||||
from libs import rsa
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SandboxArchivePath:
|
||||
tenant_id: UUID
|
||||
sandbox_id: UUID
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"sandbox/{self.tenant_id}/{self.sandbox_id}.tar.gz"
|
||||
|
||||
def proxy_path(self) -> str:
|
||||
return f"{self.tenant_id}/{self.sandbox_id}"
|
||||
|
||||
|
||||
class SandboxArchiveSigner:
|
||||
SIGNATURE_PREFIX = "sandbox-archive"
|
||||
SIGNATURE_VERSION = "v1"
|
||||
OPERATION_DOWNLOAD = "download"
|
||||
OPERATION_UPLOAD = "upload"
|
||||
|
||||
@classmethod
|
||||
def create_download_signature(cls, archive_path: SandboxArchivePath, expires_at: int, nonce: str) -> str:
|
||||
return cls._create_signature(
|
||||
archive_path=archive_path,
|
||||
operation=cls.OPERATION_DOWNLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def create_upload_signature(cls, archive_path: SandboxArchivePath, expires_at: int, nonce: str) -> str:
|
||||
return cls._create_signature(
|
||||
archive_path=archive_path,
|
||||
operation=cls.OPERATION_UPLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def verify_download_signature(
|
||||
cls, archive_path: SandboxArchivePath, expires_at: int, nonce: str, sign: str
|
||||
) -> bool:
|
||||
return cls._verify_signature(
|
||||
archive_path=archive_path,
|
||||
operation=cls.OPERATION_DOWNLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
sign=sign,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def verify_upload_signature(cls, archive_path: SandboxArchivePath, expires_at: int, nonce: str, sign: str) -> bool:
|
||||
return cls._verify_signature(
|
||||
archive_path=archive_path,
|
||||
operation=cls.OPERATION_UPLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
sign=sign,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _verify_signature(
|
||||
cls,
|
||||
*,
|
||||
archive_path: SandboxArchivePath,
|
||||
operation: str,
|
||||
expires_at: int,
|
||||
nonce: str,
|
||||
sign: str,
|
||||
) -> bool:
|
||||
if expires_at <= 0:
|
||||
return False
|
||||
|
||||
expected_sign = cls._create_signature(
|
||||
archive_path=archive_path,
|
||||
operation=operation,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
if not hmac.compare_digest(sign, expected_sign):
|
||||
return False
|
||||
|
||||
current_time = int(time.time())
|
||||
if expires_at < current_time:
|
||||
return False
|
||||
|
||||
if expires_at - current_time > dify_config.FILES_ACCESS_TIMEOUT:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def build_signed_url(
|
||||
cls,
|
||||
*,
|
||||
archive_path: SandboxArchivePath,
|
||||
expires_in: int,
|
||||
action: str,
|
||||
) -> str:
|
||||
expires_in = min(expires_in, dify_config.FILES_ACCESS_TIMEOUT)
|
||||
expires_at = int(time.time()) + max(expires_in, 1)
|
||||
nonce = os.urandom(16).hex()
|
||||
sign = cls._create_signature(
|
||||
archive_path=archive_path,
|
||||
operation=action,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
|
||||
base_url = dify_config.FILES_URL
|
||||
url = f"{base_url}/files/sandbox-archives/{archive_path.proxy_path()}/{action}"
|
||||
query = urllib.parse.urlencode({"expires_at": expires_at, "nonce": nonce, "sign": sign})
|
||||
return f"{url}?{query}"
|
||||
|
||||
@classmethod
|
||||
def _create_signature(
|
||||
cls,
|
||||
*,
|
||||
archive_path: SandboxArchivePath,
|
||||
operation: str,
|
||||
expires_at: int,
|
||||
nonce: str,
|
||||
) -> str:
|
||||
key = cls._tenant_key(str(archive_path.tenant_id))
|
||||
message = (
|
||||
f"{cls.SIGNATURE_PREFIX}|{cls.SIGNATURE_VERSION}|{operation}|"
|
||||
f"{archive_path.tenant_id}|{archive_path.sandbox_id}|{expires_at}|{nonce}"
|
||||
)
|
||||
sign = hmac.new(key, message.encode(), hashlib.sha256).digest()
|
||||
return base64.urlsafe_b64encode(sign).decode()
|
||||
|
||||
@classmethod
|
||||
def _tenant_key(cls, tenant_id: str) -> bytes:
|
||||
try:
|
||||
rsa_key, _ = rsa.get_decrypt_decoding(tenant_id)
|
||||
except rsa.PrivkeyNotFoundError as exc:
|
||||
raise ValueError(f"Tenant private key missing for tenant_id={tenant_id}") from exc
|
||||
private_key = rsa_key.export_key()
|
||||
return hashlib.sha256(private_key).digest()
|
||||
@ -1,155 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
import time
|
||||
import urllib.parse
|
||||
from dataclasses import dataclass
|
||||
from uuid import UUID
|
||||
|
||||
from configs import dify_config
|
||||
from libs import rsa
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SandboxFileDownloadPath:
|
||||
tenant_id: UUID
|
||||
sandbox_id: UUID
|
||||
export_id: str
|
||||
filename: str
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"sandbox_file_downloads/{self.tenant_id}/{self.sandbox_id}/{self.export_id}/{self.filename}"
|
||||
|
||||
def proxy_path(self) -> str:
|
||||
encoded_parts = [
|
||||
urllib.parse.quote(str(self.tenant_id), safe=""),
|
||||
urllib.parse.quote(str(self.sandbox_id), safe=""),
|
||||
urllib.parse.quote(self.export_id, safe=""),
|
||||
urllib.parse.quote(self.filename, safe=""),
|
||||
]
|
||||
return "/".join(encoded_parts)
|
||||
|
||||
|
||||
class SandboxFileSigner:
|
||||
SIGNATURE_PREFIX = "sandbox-file-download"
|
||||
SIGNATURE_VERSION = "v1"
|
||||
OPERATION_DOWNLOAD = "download"
|
||||
OPERATION_UPLOAD = "upload"
|
||||
|
||||
@classmethod
|
||||
def build_signed_url(
|
||||
cls,
|
||||
*,
|
||||
export_path: SandboxFileDownloadPath,
|
||||
expires_in: int,
|
||||
action: str,
|
||||
) -> str:
|
||||
expires_in = min(expires_in, dify_config.FILES_ACCESS_TIMEOUT)
|
||||
expires_at = int(time.time()) + max(expires_in, 1)
|
||||
nonce = os.urandom(16).hex()
|
||||
sign = cls._create_signature(
|
||||
export_path=export_path,
|
||||
operation=action,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
|
||||
base_url = dify_config.FILES_URL
|
||||
url = f"{base_url}/files/sandbox-file-downloads/{export_path.proxy_path()}/{action}"
|
||||
query = urllib.parse.urlencode({"expires_at": expires_at, "nonce": nonce, "sign": sign})
|
||||
return f"{url}?{query}"
|
||||
|
||||
@classmethod
|
||||
def verify_download_signature(
|
||||
cls,
|
||||
*,
|
||||
export_path: SandboxFileDownloadPath,
|
||||
expires_at: int,
|
||||
nonce: str,
|
||||
sign: str,
|
||||
) -> bool:
|
||||
return cls._verify_signature(
|
||||
export_path=export_path,
|
||||
operation=cls.OPERATION_DOWNLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
sign=sign,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def verify_upload_signature(
|
||||
cls,
|
||||
*,
|
||||
export_path: SandboxFileDownloadPath,
|
||||
expires_at: int,
|
||||
nonce: str,
|
||||
sign: str,
|
||||
) -> bool:
|
||||
return cls._verify_signature(
|
||||
export_path=export_path,
|
||||
operation=cls.OPERATION_UPLOAD,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
sign=sign,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _verify_signature(
|
||||
cls,
|
||||
*,
|
||||
export_path: SandboxFileDownloadPath,
|
||||
operation: str,
|
||||
expires_at: int,
|
||||
nonce: str,
|
||||
sign: str,
|
||||
) -> bool:
|
||||
if expires_at <= 0:
|
||||
return False
|
||||
|
||||
expected_sign = cls._create_signature(
|
||||
export_path=export_path,
|
||||
operation=operation,
|
||||
expires_at=expires_at,
|
||||
nonce=nonce,
|
||||
)
|
||||
if not hmac.compare_digest(sign, expected_sign):
|
||||
return False
|
||||
|
||||
current_time = int(time.time())
|
||||
if expires_at < current_time:
|
||||
return False
|
||||
|
||||
if expires_at - current_time > dify_config.FILES_ACCESS_TIMEOUT:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def _create_signature(
|
||||
cls,
|
||||
*,
|
||||
export_path: SandboxFileDownloadPath,
|
||||
operation: str,
|
||||
expires_at: int,
|
||||
nonce: str,
|
||||
) -> str:
|
||||
key = cls._tenant_key(str(export_path.tenant_id))
|
||||
message = (
|
||||
f"{cls.SIGNATURE_PREFIX}|{cls.SIGNATURE_VERSION}|{operation}|"
|
||||
f"{export_path.tenant_id}|{export_path.sandbox_id}|{export_path.export_id}|{export_path.filename}|"
|
||||
f"{expires_at}|{nonce}"
|
||||
)
|
||||
digest = hmac.new(key, message.encode(), hashlib.sha256).digest()
|
||||
return base64.urlsafe_b64encode(digest).decode()
|
||||
|
||||
@classmethod
|
||||
def _tenant_key(cls, tenant_id: str) -> bytes:
|
||||
try:
|
||||
rsa_key, _ = rsa.get_decrypt_decoding(tenant_id)
|
||||
except rsa.PrivkeyNotFoundError as exc:
|
||||
raise ValueError(f"Tenant private key missing for tenant_id={tenant_id}") from exc
|
||||
private_key = rsa_key.export_key()
|
||||
return hashlib.sha256(private_key).digest()
|
||||
@ -1,11 +1,13 @@
|
||||
from .archive_storage import ArchiveSandboxStorage
|
||||
from .archive_storage import ArchiveSandboxStorage, SandboxArchivePath
|
||||
from .noop_storage import NoopSandboxStorage
|
||||
from .sandbox_file_storage import SandboxFileStorage, sandbox_file_storage
|
||||
from .sandbox_file_storage import SandboxFileDownloadPath, SandboxFileStorage, sandbox_file_storage
|
||||
from .sandbox_storage import SandboxStorage
|
||||
|
||||
__all__ = [
|
||||
"ArchiveSandboxStorage",
|
||||
"NoopSandboxStorage",
|
||||
"SandboxArchivePath",
|
||||
"SandboxFileDownloadPath",
|
||||
"SandboxFileStorage",
|
||||
"SandboxStorage",
|
||||
"sandbox_file_storage",
|
||||
|
||||
@ -1,18 +1,31 @@
|
||||
"""Archive-based sandbox storage for persisting sandbox state.
|
||||
|
||||
This module provides storage operations for sandbox workspace archives (tar.gz),
|
||||
enabling state persistence across sandbox sessions.
|
||||
|
||||
Storage key format: sandbox/{tenant_id}/{sandbox_id}.tar.gz
|
||||
|
||||
All presign operations use the unified FilePresignStorage wrapper, which automatically
|
||||
falls back to Dify's file proxy when the underlying storage doesn't support presigned URLs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from uuid import UUID
|
||||
|
||||
from core.sandbox.security.archive_signer import SandboxArchivePath, SandboxArchiveSigner
|
||||
from core.virtual_environment.__base.exec import PipelineExecutionError
|
||||
from core.virtual_environment.__base.helpers import pipeline
|
||||
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
|
||||
from extensions.ext_storage import storage
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
from extensions.storage.file_presign_storage import FilePresignStorage
|
||||
|
||||
from .sandbox_storage import SandboxStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
WORKSPACE_DIR = "."
|
||||
|
||||
ARCHIVE_DOWNLOAD_TIMEOUT = 60 * 5
|
||||
ARCHIVE_UPLOAD_TIMEOUT = 60 * 5
|
||||
|
||||
@ -21,40 +34,67 @@ def build_tar_exclude_args(patterns: list[str]) -> list[str]:
|
||||
return [f"--exclude={p}" for p in patterns]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SandboxArchivePath:
|
||||
"""Path for sandbox workspace archives."""
|
||||
|
||||
tenant_id: UUID
|
||||
sandbox_id: UUID
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"sandbox/{self.tenant_id}/{self.sandbox_id}.tar.gz"
|
||||
|
||||
|
||||
class ArchiveSandboxStorage(SandboxStorage):
|
||||
"""Archive-based storage for sandbox workspace persistence.
|
||||
|
||||
Uses tar.gz archives to save and restore sandbox workspace state.
|
||||
Requires a presign-capable storage wrapper for generating download/upload URLs.
|
||||
"""
|
||||
|
||||
_tenant_id: str
|
||||
_sandbox_id: str
|
||||
_exclude_patterns: list[str]
|
||||
_storage: FilePresignStorage
|
||||
|
||||
def __init__(self, tenant_id: str, sandbox_id: str, exclude_patterns: list[str] | None = None):
|
||||
def __init__(
|
||||
self,
|
||||
tenant_id: str,
|
||||
sandbox_id: str,
|
||||
storage: BaseStorage,
|
||||
exclude_patterns: list[str] | None = None,
|
||||
):
|
||||
self._tenant_id = tenant_id
|
||||
self._sandbox_id = sandbox_id
|
||||
self._exclude_patterns = exclude_patterns or []
|
||||
# Wrap with FilePresignStorage for presign fallback support
|
||||
self._storage = FilePresignStorage(storage)
|
||||
|
||||
@property
|
||||
def _archive_path(self) -> SandboxArchivePath:
|
||||
return SandboxArchivePath(UUID(self._tenant_id), UUID(self._sandbox_id))
|
||||
|
||||
@property
|
||||
def _storage_key(self) -> str:
|
||||
return SandboxArchivePath(UUID(self._tenant_id), UUID(self._sandbox_id)).get_storage_key()
|
||||
return self._archive_path.get_storage_key()
|
||||
|
||||
@property
|
||||
def _archive_name(self) -> str:
|
||||
return f"{self._sandbox_id}.tar.gz"
|
||||
|
||||
@property
|
||||
def _archive_path(self) -> str:
|
||||
def _archive_tmp_path(self) -> str:
|
||||
return f"/tmp/{self._archive_name}"
|
||||
|
||||
def mount(self, sandbox: VirtualEnvironment) -> bool:
|
||||
"""Load archive from storage into sandbox workspace."""
|
||||
if not self.exists():
|
||||
logger.debug("No archive found for sandbox %s, skipping mount", self._sandbox_id)
|
||||
return False
|
||||
|
||||
archive_path = SandboxArchivePath(UUID(self._tenant_id), UUID(self._sandbox_id))
|
||||
download_url = SandboxArchiveSigner.build_signed_url(
|
||||
archive_path=archive_path,
|
||||
expires_in=ARCHIVE_DOWNLOAD_TIMEOUT,
|
||||
action=SandboxArchiveSigner.OPERATION_DOWNLOAD,
|
||||
)
|
||||
download_url = self._storage.get_download_url(self._storage_key, ARCHIVE_DOWNLOAD_TIMEOUT)
|
||||
archive_name = self._archive_name
|
||||
|
||||
try:
|
||||
(
|
||||
pipeline(sandbox)
|
||||
@ -74,13 +114,10 @@ class ArchiveSandboxStorage(SandboxStorage):
|
||||
return True
|
||||
|
||||
def unmount(self, sandbox: VirtualEnvironment) -> bool:
|
||||
archive_path = SandboxArchivePath(UUID(self._tenant_id), UUID(self._sandbox_id))
|
||||
upload_url = SandboxArchiveSigner.build_signed_url(
|
||||
archive_path=archive_path,
|
||||
expires_in=ARCHIVE_UPLOAD_TIMEOUT,
|
||||
action=SandboxArchiveSigner.OPERATION_UPLOAD,
|
||||
)
|
||||
archive_path = self._archive_path
|
||||
"""Save sandbox workspace to storage as archive."""
|
||||
upload_url = self._storage.get_upload_url(self._storage_key, ARCHIVE_UPLOAD_TIMEOUT)
|
||||
archive_path = self._archive_tmp_path
|
||||
|
||||
(
|
||||
pipeline(sandbox)
|
||||
.add(
|
||||
@ -105,11 +142,13 @@ class ArchiveSandboxStorage(SandboxStorage):
|
||||
return True
|
||||
|
||||
def exists(self) -> bool:
|
||||
return storage.exists(self._storage_key)
|
||||
"""Check if archive exists in storage."""
|
||||
return self._storage.exists(self._storage_key)
|
||||
|
||||
def delete(self) -> None:
|
||||
"""Delete archive from storage."""
|
||||
try:
|
||||
storage.delete(self._storage_key)
|
||||
self._storage.delete(self._storage_key)
|
||||
logger.info("Deleted archive for sandbox %s", self._sandbox_id)
|
||||
except Exception:
|
||||
logger.exception("Failed to delete archive for sandbox %s", self._sandbox_id)
|
||||
|
||||
@ -1,23 +1,58 @@
|
||||
"""Sandbox file storage for exporting files from sandbox environments.
|
||||
|
||||
This module provides storage operations for files exported from sandbox environments,
|
||||
including download tickets for both runtime and archive-based file sources.
|
||||
|
||||
Storage key format: sandbox_file_downloads/{tenant_id}/{sandbox_id}/{export_id}/{filename}
|
||||
|
||||
All presign operations use the unified FilePresignStorage wrapper, which automatically
|
||||
falls back to Dify's file proxy when the underlying storage doesn't support presigned URLs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from core.sandbox.security.sandbox_file_signer import SandboxFileDownloadPath, SandboxFileSigner
|
||||
from extensions.ext_redis import redis_client
|
||||
from extensions.ext_storage import storage
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
from extensions.storage.cached_presign_storage import CachedPresignStorage
|
||||
from extensions.storage.silent_storage import SilentStorage
|
||||
from extensions.storage.file_presign_storage import FilePresignStorage
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SandboxFileDownloadPath:
|
||||
"""Path for sandbox file exports."""
|
||||
|
||||
tenant_id: UUID
|
||||
sandbox_id: UUID
|
||||
export_id: str
|
||||
filename: str
|
||||
|
||||
def get_storage_key(self) -> str:
|
||||
return f"sandbox_file_downloads/{self.tenant_id}/{self.sandbox_id}/{self.export_id}/{self.filename}"
|
||||
|
||||
|
||||
class SandboxFileStorage:
|
||||
_base_storage: BaseStorage
|
||||
"""Storage operations for sandbox file exports.
|
||||
|
||||
Wraps BaseStorage with:
|
||||
- FilePresignStorage for presign fallback support
|
||||
- CachedPresignStorage for URL caching
|
||||
|
||||
Usage:
|
||||
storage = SandboxFileStorage(base_storage, redis_client=redis)
|
||||
storage.save(download_path, content)
|
||||
url = storage.get_download_url(download_path)
|
||||
"""
|
||||
|
||||
_storage: CachedPresignStorage
|
||||
|
||||
def __init__(self, storage: BaseStorage, *, redis_client: Any) -> None:
|
||||
self._base_storage = storage
|
||||
# Wrap with FilePresignStorage for fallback support, then CachedPresignStorage for caching
|
||||
presign_storage = FilePresignStorage(storage)
|
||||
self._storage = CachedPresignStorage(
|
||||
storage=storage,
|
||||
storage=presign_storage,
|
||||
redis_client=redis_client,
|
||||
cache_key_prefix="sandbox_file_downloads",
|
||||
)
|
||||
@ -26,29 +61,19 @@ class SandboxFileStorage:
|
||||
self._storage.save(download_path.get_storage_key(), content)
|
||||
|
||||
def get_download_url(self, download_path: SandboxFileDownloadPath, expires_in: int = 3600) -> str:
|
||||
storage_key = download_path.get_storage_key()
|
||||
try:
|
||||
return self._storage.get_download_url(storage_key, expires_in)
|
||||
except NotImplementedError:
|
||||
return SandboxFileSigner.build_signed_url(
|
||||
export_path=download_path,
|
||||
expires_in=expires_in,
|
||||
action=SandboxFileSigner.OPERATION_DOWNLOAD,
|
||||
)
|
||||
return self._storage.get_download_url(download_path.get_storage_key(), expires_in)
|
||||
|
||||
def get_upload_url(self, download_path: SandboxFileDownloadPath, expires_in: int = 3600) -> str:
|
||||
storage_key = download_path.get_storage_key()
|
||||
try:
|
||||
return self._storage.get_upload_url(storage_key, expires_in)
|
||||
except NotImplementedError:
|
||||
return SandboxFileSigner.build_signed_url(
|
||||
export_path=download_path,
|
||||
expires_in=expires_in,
|
||||
action=SandboxFileSigner.OPERATION_UPLOAD,
|
||||
)
|
||||
return self._storage.get_upload_url(download_path.get_storage_key(), expires_in)
|
||||
|
||||
|
||||
class _LazySandboxFileStorage:
|
||||
"""Lazy initializer for singleton SandboxFileStorage.
|
||||
|
||||
Delays storage initialization until first access, ensuring Flask app
|
||||
context is available.
|
||||
"""
|
||||
|
||||
_instance: SandboxFileStorage | None
|
||||
|
||||
def __init__(self) -> None:
|
||||
@ -56,12 +81,16 @@ class _LazySandboxFileStorage:
|
||||
|
||||
def _get_instance(self) -> SandboxFileStorage:
|
||||
if self._instance is None:
|
||||
from extensions.ext_redis import redis_client
|
||||
from extensions.ext_storage import storage
|
||||
|
||||
if not hasattr(storage, "storage_runner"):
|
||||
raise RuntimeError(
|
||||
"Storage is not initialized; call storage.init_app before using sandbox_file_storage"
|
||||
)
|
||||
self._instance = SandboxFileStorage(
|
||||
storage=SilentStorage(storage.storage_runner), redis_client=redis_client
|
||||
storage=storage.storage_runner,
|
||||
redis_client=redis_client,
|
||||
)
|
||||
return self._instance
|
||||
|
||||
@ -69,4 +98,4 @@ class _LazySandboxFileStorage:
|
||||
return getattr(self._get_instance(), name)
|
||||
|
||||
|
||||
sandbox_file_storage = _LazySandboxFileStorage()
|
||||
sandbox_file_storage: SandboxFileStorage = _LazySandboxFileStorage() # type: ignore[assignment]
|
||||
|
||||
Reference in New Issue
Block a user