shortener.py - cdnV2 -

shortener.py (2005B)
      1 import random
      2 import string
      3 import asyncio
      4 import re
      5 import httpx
      6 
      7 from .storage import HF_TOKEN, bucket_url, auth_headers, _do_upload_bytes
      8 
      9 # Reusable HTTP client for fast connection pooling & HTTP keep-alive
     10 _client = httpx.AsyncClient(follow_redirects=True, timeout=10.0)
     11 
     12 # Alphanumeric validator to guard against path traversal attempts
     13 _ID_REGEX = re.compile(r"^[a-z0-9]{4,16}$")
     14 
     15 def generate_id(length: int = 8) -> str:
     16     chars = string.ascii_lowercase + string.digits
     17     return "".join(random.choices(chars, k=length))
     18 
     19 async def shorten_url(destination_url: str) -> str:
     20     """Save a destination URL as a tiny object in the Hugging Face bucket."""
     21     if not HF_TOKEN:
     22         raise RuntimeError("HF_TOKEN is not configured.")
     23 
     24     destination_url = destination_url.strip()
     25 
     26     # Collision-check loop with a safeguard cutoff
     27     attempts = 0
     28     while True:
     29         short_id = generate_id()
     30         r = await _client.head(bucket_url(f"_shortened/{short_id}"), headers=auth_headers())
     31         if r.status_code == 404:
     32             break
     33         attempts += 1
     34         if attempts > 5:
     35             # Fallback to longer ID if collisions occur
     36             short_id = generate_id(length=12)
     37             break
     38 
     39     # Upload straight from memory without writing to disk
     40     await asyncio.to_thread(
     41         _do_upload_bytes,
     42         destination_url.encode("utf-8"),
     43         f"_shortened/{short_id}",
     44     )
     45     return short_id
     46 
     47 async def get_destination_url(short_id: str) -> str | None:
     48     """Look up the destination URL directly from the bucket's resolve endpoint."""
     49     short_id = short_id.strip().lower()
     50     if not _ID_REGEX.match(short_id):
     51         return None
     52 
     53     try:
     54         r = await _client.get(bucket_url(f"_shortened/{short_id}"), headers=auth_headers())
     55         if r.status_code == 200:
     56             return r.text.strip()
     57     except httpx.RequestError:
     58         return None
     59     return None
     60 
     61 def list_all_ids() -> list[str]:
     62     """Stub kept for router compatibility."""
     63     return []