mirror of
https://github.com/langgenius/dify.git
synced 2026-04-24 21:05:48 +08:00
refactor(storage): replace storage proxy with ticket-based URL system
- Removed the storage proxy controller and its associated endpoints for file download and upload. - Updated the file controller to use the new storage ticket service for generating download and upload URLs. - Modified the file presign storage to fallback to ticket-based URLs instead of signed proxy URLs. - Enhanced unit tests to validate the new ticket generation and retrieval logic.
This commit is contained in:
Harry
@ -1,4 +1,4 @@
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
@ -6,7 +6,7 @@ import pytest
|
||||
from configs import dify_config
|
||||
from core.app_assets.storage import AppAssetStorage, AssetPath
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
from extensions.storage.file_presign_storage import FilePresignStorage
|
||||
from services.storage_ticket_service import StorageTicket, StorageTicketService
|
||||
|
||||
|
||||
class DummyStorage(BaseStorage):
|
||||
@ -70,96 +70,133 @@ def test_asset_path_validation():
|
||||
AssetPath.draft(tenant_id=tenant_id, app_id=app_id, node_id="not-a-uuid")
|
||||
|
||||
|
||||
def test_file_presign_signature_verification(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test FilePresignStorage signature creation and verification."""
|
||||
monkeypatch.setattr(dify_config, "SECRET_KEY", "test-secret-key", raising=False)
|
||||
monkeypatch.setattr(dify_config, "FILES_ACCESS_TIMEOUT", 300, raising=False)
|
||||
def test_storage_ticket_service(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test StorageTicketService creates and retrieves tickets."""
|
||||
monkeypatch.setattr(dify_config, "FILES_URL", "http://files.local", raising=False)
|
||||
|
||||
filename = "test/path/file.txt"
|
||||
timestamp = str(int(time.time()))
|
||||
nonce = "test-nonce"
|
||||
mock_redis = MagicMock()
|
||||
stored_data = {}
|
||||
|
||||
# Test download signature
|
||||
sign = FilePresignStorage._create_signature("download", filename, timestamp, nonce)
|
||||
assert FilePresignStorage.verify_signature(
|
||||
filename=filename,
|
||||
operation="download",
|
||||
timestamp=timestamp,
|
||||
nonce=nonce,
|
||||
sign=sign,
|
||||
)
|
||||
def mock_setex(key, ttl, value):
|
||||
stored_data[key] = value
|
||||
|
||||
# Test upload signature
|
||||
upload_sign = FilePresignStorage._create_signature("upload", filename, timestamp, nonce)
|
||||
assert FilePresignStorage.verify_signature(
|
||||
filename=filename,
|
||||
operation="upload",
|
||||
timestamp=timestamp,
|
||||
nonce=nonce,
|
||||
sign=upload_sign,
|
||||
)
|
||||
def mock_get(key):
|
||||
return stored_data.get(key)
|
||||
|
||||
# Test expired signature
|
||||
expired_timestamp = str(int(time.time()) - 400)
|
||||
expired_sign = FilePresignStorage._create_signature("download", filename, expired_timestamp, nonce)
|
||||
assert not FilePresignStorage.verify_signature(
|
||||
filename=filename,
|
||||
operation="download",
|
||||
timestamp=expired_timestamp,
|
||||
nonce=nonce,
|
||||
sign=expired_sign,
|
||||
)
|
||||
mock_redis.setex = mock_setex
|
||||
mock_redis.get = mock_get
|
||||
|
||||
# Test wrong signature
|
||||
assert not FilePresignStorage.verify_signature(
|
||||
filename=filename,
|
||||
operation="download",
|
||||
timestamp=timestamp,
|
||||
nonce=nonce,
|
||||
sign="wrong-signature",
|
||||
)
|
||||
with patch("services.storage_ticket_service.redis_client", mock_redis):
|
||||
# Test download URL creation
|
||||
url = StorageTicketService.create_download_url("test/path/file.txt", expires_in=300, filename="file.txt")
|
||||
|
||||
assert url.startswith("http://files.local/files/storage-files/")
|
||||
token = url.split("/")[-1]
|
||||
|
||||
# Verify ticket was stored
|
||||
ticket = StorageTicketService.get_ticket(token)
|
||||
assert ticket is not None
|
||||
assert ticket.op == "download"
|
||||
assert ticket.storage_key == "test/path/file.txt"
|
||||
assert ticket.filename == "file.txt"
|
||||
|
||||
# Test upload URL creation
|
||||
upload_url = StorageTicketService.create_upload_url("test/upload.txt", expires_in=300, max_bytes=1024)
|
||||
|
||||
upload_token = upload_url.split("/")[-1]
|
||||
upload_ticket = StorageTicketService.get_ticket(upload_token)
|
||||
assert upload_ticket is not None
|
||||
assert upload_ticket.op == "upload"
|
||||
assert upload_ticket.storage_key == "test/upload.txt"
|
||||
assert upload_ticket.max_bytes == 1024
|
||||
|
||||
|
||||
def test_signed_proxy_url_generation(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that AppAssetStorage generates correct proxy URLs when presign is not supported."""
|
||||
def test_storage_ticket_not_found(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test StorageTicketService returns None for invalid token."""
|
||||
mock_redis = MagicMock()
|
||||
mock_redis.get.return_value = None
|
||||
|
||||
with patch("services.storage_ticket_service.redis_client", mock_redis):
|
||||
ticket = StorageTicketService.get_ticket("invalid-token")
|
||||
assert ticket is None
|
||||
|
||||
|
||||
def test_ticket_url_generation(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that AppAssetStorage generates correct ticket URLs when presign is not supported."""
|
||||
tenant_id = str(uuid4())
|
||||
app_id = str(uuid4())
|
||||
resource_id = str(uuid4())
|
||||
asset_path = AssetPath.draft(tenant_id, app_id, resource_id)
|
||||
|
||||
monkeypatch.setattr(dify_config, "SECRET_KEY", "test-secret-key", raising=False)
|
||||
monkeypatch.setattr(dify_config, "FILES_ACCESS_TIMEOUT", 300, raising=False)
|
||||
monkeypatch.setattr(dify_config, "FILES_URL", "http://files.local", raising=False)
|
||||
|
||||
storage = AppAssetStorage(DummyStorage(), redis_client=DummyRedis())
|
||||
url = storage.get_download_url(asset_path, expires_in=120)
|
||||
mock_redis = MagicMock()
|
||||
mock_redis.setex = MagicMock()
|
||||
|
||||
# URL should be a proxy URL since DummyStorage doesn't support presign
|
||||
storage_key = asset_path.get_storage_key()
|
||||
assert url.startswith("http://files.local/files/storage/")
|
||||
assert "/download?" in url
|
||||
assert "timestamp=" in url
|
||||
assert "nonce=" in url
|
||||
assert "sign=" in url
|
||||
with patch("services.storage_ticket_service.redis_client", mock_redis):
|
||||
storage = AppAssetStorage(DummyStorage(), redis_client=DummyRedis())
|
||||
url = storage.get_download_url(asset_path, expires_in=120)
|
||||
|
||||
# URL should be a ticket URL since DummyStorage doesn't support presign
|
||||
assert url.startswith("http://files.local/files/storage-files/")
|
||||
# Token should be a UUID
|
||||
token = url.split("/")[-1]
|
||||
assert len(token) == 36 # UUID format
|
||||
|
||||
|
||||
def test_upload_url_generation(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that AppAssetStorage generates correct upload URLs."""
|
||||
def test_upload_ticket_url_generation(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that AppAssetStorage generates correct upload ticket URLs."""
|
||||
tenant_id = str(uuid4())
|
||||
app_id = str(uuid4())
|
||||
resource_id = str(uuid4())
|
||||
asset_path = AssetPath.draft(tenant_id, app_id, resource_id)
|
||||
|
||||
monkeypatch.setattr(dify_config, "SECRET_KEY", "test-secret-key", raising=False)
|
||||
monkeypatch.setattr(dify_config, "FILES_ACCESS_TIMEOUT", 300, raising=False)
|
||||
monkeypatch.setattr(dify_config, "FILES_URL", "http://files.local", raising=False)
|
||||
|
||||
storage = AppAssetStorage(DummyStorage(), redis_client=DummyRedis())
|
||||
url = storage.get_upload_url(asset_path, expires_in=120)
|
||||
mock_redis = MagicMock()
|
||||
mock_redis.setex = MagicMock()
|
||||
|
||||
# URL should be a proxy URL since DummyStorage doesn't support presign
|
||||
assert url.startswith("http://files.local/files/storage/")
|
||||
assert "/upload?" in url
|
||||
assert "timestamp=" in url
|
||||
assert "nonce=" in url
|
||||
assert "sign=" in url
|
||||
with patch("services.storage_ticket_service.redis_client", mock_redis):
|
||||
storage = AppAssetStorage(DummyStorage(), redis_client=DummyRedis())
|
||||
url = storage.get_upload_url(asset_path, expires_in=120)
|
||||
|
||||
# URL should be a ticket URL since DummyStorage doesn't support presign
|
||||
assert url.startswith("http://files.local/files/storage-files/")
|
||||
# Token should be a UUID
|
||||
token = url.split("/")[-1]
|
||||
assert len(token) == 36 # UUID format
|
||||
|
||||
|
||||
def test_storage_ticket_dataclass():
|
||||
"""Test StorageTicket serialization and deserialization."""
|
||||
ticket = StorageTicket(
|
||||
op="download",
|
||||
storage_key="path/to/file.txt",
|
||||
filename="file.txt",
|
||||
)
|
||||
|
||||
data = ticket.to_dict()
|
||||
assert data == {
|
||||
"op": "download",
|
||||
"storage_key": "path/to/file.txt",
|
||||
"filename": "file.txt",
|
||||
}
|
||||
|
||||
restored = StorageTicket.from_dict(data)
|
||||
assert restored.op == ticket.op
|
||||
assert restored.storage_key == ticket.storage_key
|
||||
assert restored.filename == ticket.filename
|
||||
assert restored.max_bytes is None
|
||||
|
||||
# Test upload ticket with max_bytes
|
||||
upload_ticket = StorageTicket(
|
||||
op="upload",
|
||||
storage_key="path/to/upload.txt",
|
||||
max_bytes=1024,
|
||||
)
|
||||
|
||||
upload_data = upload_ticket.to_dict()
|
||||
assert upload_data["max_bytes"] == 1024
|
||||
|
||||
restored_upload = StorageTicket.from_dict(upload_data)
|
||||
assert restored_upload.max_bytes == 1024
|
||||
|
||||
Reference in New Issue
Block a user