Using short link like imgbb

Author: lunaticsm

app/config.py

@@ -17,3 +17,4 @@
 CACHE_MAX_AGE_SECONDS = int(os.getenv("CACHE_MAX_AGE_SECONDS", "3600"))
 ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "admin-dev-password")
 ADMIN_LOCK_STEP_SECONDS = int(os.getenv("ADMIN_LOCK_STEP_SECONDS", str(5 * 60)))
+FILE_ID_LENGTH = max(4, min(32, int(os.getenv("FILE_ID_LENGTH", "7"))))

app/storage.py

@@ -1,17 +1,36 @@
+from __future__ import annotations
+
 import os
-import uuid
+import secrets
+import string
 from datetime import datetime, timedelta
 from sqlmodel import Session, select
-from app.config import UPLOAD_DIR, DELETE_AFTER_HOURS
+from app.config import UPLOAD_DIR, DELETE_AFTER_HOURS, FILE_ID_LENGTH
 from app.models import File
 
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 
+_SLUG_ALPHABET = string.ascii_letters + string.digits
+_MAX_SLUG_ATTEMPTS = 5
+
+
+def _generate_file_id(length: int = FILE_ID_LENGTH) -> str:
+    return "".join(secrets.choice(_SLUG_ALPHABET) for _ in range(length))
+
+
+def _reserve_path(ext: str) -> tuple[str, str]:
+    for _ in range(_MAX_SLUG_ATTEMPTS):
+        file_id = _generate_file_id()
+        stored_name = f"{file_id}{ext}"
+        path = os.path.join(UPLOAD_DIR, stored_name)
+        if not os.path.exists(path):
+            return stored_name, path
+    raise RuntimeError("Unable to allocate a unique file slug")
+
+
 def save_file(file_bytes: bytes, original_name: str, content_type: str) -> str:
     ext = os.path.splitext(original_name)[1] or ".bin"
-    file_id = str(uuid.uuid4())
-    stored_name = f"{file_id}{ext}"
-    path = os.path.join(UPLOAD_DIR, stored_name)
+    stored_name, path = _reserve_path(ext)
     with open(path, "wb") as f:
         f.write(file_bytes)
     return stored_name, len(file_bytes)

tests/test_app.py

@@ -1,4 +1,5 @@
 import importlib
+import re
 import sys
 from datetime import datetime, timedelta
 from pathlib import Path
@@ -71,6 +72,18 @@ def test_upload_list_serve_and_cache_headers(client):
     assert serve_response.headers["Cache-Control"] == "public, max-age=120"
 
 
+def test_upload_slug_is_short(client):
+    response = client.post("/upload", files={"file": ("slug.txt", b"x", "text/plain")})
+    assert response.status_code == 200
+    payload = response.json()
+    from app import config as app_config
+
+    slug = payload["id"]
+    assert len(slug) == app_config.FILE_ID_LENGTH
+    assert re.fullmatch(rf"[A-Za-z0-9]{{{app_config.FILE_ID_LENGTH}}}", slug)
+    assert payload["url"].split("/")[-1].startswith(slug)
+
+
 def test_directory_traversal_blocked(client):
     outside_file = client.upload_dir.parent / "top_secret.txt"  # type: ignore[attr-defined]
     outside_file.write_text("nope")