mirror of
https://github.com/AstrBotDevs/AstrBot.git
synced 2026-10-16 15:06:04 +08:00
fix(dashboard): store chat attachments under unique names (#10356)
* fix(dashboard): store chat attachments under unique names Uploads were saved under their original filename, so two attachments with the same name (every pasted screenshot is image.png) overwrote each other, and deleting one session removed a file another session still used. Store each upload as <timestamp id>_<name> and return the original name as `filename` for display, with the on-disk name in `stored_filename`. Fixes #10352 * fix(dashboard): keep long-suffix attachment names within 255 bytes
This commit is contained in:
@@ -63,6 +63,18 @@ def sanitize_upload_filename(filename: str | None) -> str:
|
||||
return name
|
||||
|
||||
|
||||
def unique_attachment_filename(filename: str) -> str:
|
||||
# Pasted screenshots are all named "image.png", so prefix the stored name
|
||||
# to keep uploads apart, and stay within the 255-byte filename limit.
|
||||
prefix = f"{generate_timestamp_id()}_"
|
||||
budget = 255 - len(prefix.encode())
|
||||
stem, suffix = os.path.splitext(filename)
|
||||
if len(suffix.encode()) >= budget:
|
||||
stem, suffix = filename, ""
|
||||
stem = stem.encode()[: budget - len(suffix.encode())].decode(errors="ignore")
|
||||
return f"{prefix}{stem}{suffix}"
|
||||
|
||||
|
||||
class LocalUploadFile:
|
||||
"""Adapt a merged local file to the upload contract save_uploaded_file() consumes.
|
||||
|
||||
@@ -649,7 +661,9 @@ class ChatService:
|
||||
attach_type = "file"
|
||||
|
||||
attachments_dir = Path(self.attachments_dir).resolve(strict=False)
|
||||
file_path = (attachments_dir / filename).resolve(strict=False)
|
||||
file_path = (attachments_dir / unique_attachment_filename(filename)).resolve(
|
||||
strict=False
|
||||
)
|
||||
if not file_path.is_relative_to(attachments_dir):
|
||||
raise ChatServiceError("Invalid filename")
|
||||
|
||||
@@ -676,6 +690,7 @@ class ChatService:
|
||||
)
|
||||
await asyncio.to_thread(file_path.rename, target_path)
|
||||
file_path = target_path
|
||||
filename = os.path.splitext(filename)[0] + detected_suffix
|
||||
attachment = await self.db.insert_attachment(
|
||||
path=str(file_path),
|
||||
type=attach_type,
|
||||
@@ -687,7 +702,8 @@ class ChatService:
|
||||
|
||||
return {
|
||||
"attachment_id": attachment.attachment_id,
|
||||
"filename": os.path.basename(attachment.path),
|
||||
"filename": filename,
|
||||
"stored_filename": os.path.basename(attachment.path),
|
||||
"type": attach_type,
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
from io import BytesIO
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from PIL import Image as PILImage
|
||||
|
||||
from astrbot.dashboard.services.chat_service import (
|
||||
ChatService,
|
||||
unique_attachment_filename,
|
||||
)
|
||||
|
||||
|
||||
def _png(color):
|
||||
buffer = BytesIO()
|
||||
PILImage.new("RGB", (2, 2), color).save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class FakeUploadFile:
|
||||
content_type = "image/png"
|
||||
content_length = None
|
||||
|
||||
def __init__(self, filename, data):
|
||||
self.filename = filename
|
||||
self.data = data
|
||||
|
||||
async def save(self, destination, *, max_bytes=None):
|
||||
with open(destination, "wb") as output:
|
||||
output.write(self.data)
|
||||
|
||||
|
||||
class FakeDatabase:
|
||||
def __init__(self):
|
||||
self.attachments = {}
|
||||
|
||||
async def insert_attachment(self, path, type, mime_type):
|
||||
attachment_id = f"attachment-{len(self.attachments) + 1}"
|
||||
attachment = SimpleNamespace(attachment_id=attachment_id, path=path)
|
||||
self.attachments[attachment_id] = attachment
|
||||
return attachment
|
||||
|
||||
async def get_attachments(self, attachment_ids):
|
||||
return [self.attachments[i] for i in attachment_ids]
|
||||
|
||||
async def delete_attachments(self, attachment_ids):
|
||||
for attachment_id in attachment_ids:
|
||||
self.attachments.pop(attachment_id)
|
||||
|
||||
|
||||
def make_service(tmp_path):
|
||||
service = ChatService.__new__(ChatService)
|
||||
service.db = FakeDatabase()
|
||||
service.attachments_dir = str(tmp_path)
|
||||
return service
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_name_uploads_do_not_overwrite_each_other(tmp_path):
|
||||
service = make_service(tmp_path)
|
||||
red, blue = _png((255, 0, 0)), _png((0, 0, 255))
|
||||
|
||||
first = await service.save_uploaded_file(FakeUploadFile("image.png", red))
|
||||
second = await service.save_uploaded_file(FakeUploadFile("image.png", blue))
|
||||
|
||||
assert first["stored_filename"] != second["stored_filename"]
|
||||
assert (tmp_path / first["stored_filename"]).read_bytes() == red
|
||||
assert (tmp_path / second["stored_filename"]).read_bytes() == blue
|
||||
assert first["filename"] == second["filename"] == "image.png"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleting_an_attachment_keeps_same_name_files(tmp_path):
|
||||
service = make_service(tmp_path)
|
||||
first = await service.save_uploaded_file(
|
||||
FakeUploadFile("image.png", _png((1, 2, 3)))
|
||||
)
|
||||
second = await service.save_uploaded_file(
|
||||
FakeUploadFile("image.png", _png((4, 5, 6)))
|
||||
)
|
||||
|
||||
await service.delete_attachments([second["attachment_id"]])
|
||||
|
||||
assert (tmp_path / first["stored_filename"]).exists()
|
||||
assert not (tmp_path / second["stored_filename"]).exists()
|
||||
|
||||
|
||||
def test_unique_attachment_filename_keeps_original_name():
|
||||
assert unique_attachment_filename("报告.pdf").endswith("_报告.pdf")
|
||||
|
||||
|
||||
def test_unique_attachment_filename_fits_in_255_bytes():
|
||||
name = unique_attachment_filename("图" * 120 + ".png")
|
||||
|
||||
assert len(name.encode()) <= 255
|
||||
assert name.endswith(".png")
|
||||
|
||||
|
||||
def test_unique_attachment_filename_with_long_suffix_fits_in_255_bytes():
|
||||
name = unique_attachment_filename("v1." + "x" * 250)
|
||||
|
||||
assert len(name.encode()) <= 255
|
||||
@@ -43,11 +43,12 @@ async def test_webchat_upload_uses_detected_image_type(tmp_path):
|
||||
|
||||
result = await service.save_uploaded_file(FakeUploadFile())
|
||||
|
||||
assert result["filename"].endswith(".jpg")
|
||||
assert result["filename"] == "pasted.jpg"
|
||||
assert fake_db.inserted["mime_type"] == "image/jpeg"
|
||||
assert fake_db.inserted["type"] == "image"
|
||||
assert (tmp_path / result["filename"]).exists()
|
||||
assert not (tmp_path / "pasted.png").exists()
|
||||
assert result["stored_filename"].endswith("_pasted.jpg")
|
||||
assert (tmp_path / result["stored_filename"]).exists()
|
||||
assert not list(tmp_path.glob("*.png"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
||||
Reference in New Issue
Block a user