Harden TikTok downloads with caps and URL checks

lumen_tiktok.py

@@ -121,11 +121,19 @@ def _tiktok_video_candidates(media_data: dict) -> list[dict[str, Any]]:
         raw_url = media_data.get(url_key)
         if not raw_url:
             continue
+        if not isinstance(raw_url, str):
+            log.warning("[tiktok] Skipping non-string URL for %s: %r.", url_key, raw_url)
+            continue
         if raw_url.startswith("//"):
             # Протокол-относительный URL чиним как https, иначе ниже склеилось бы в битый "https://www.tikwm.com//host/path".
             raw_url = "https:" + raw_url
         elif not raw_url.startswith("http"):
-            raw_url = "https://www.tikwm.com" + raw_url
+            # Склейка давала userinfo-путаницу ("@evil.com/x" ехал хостом):
+            # только относительный путь к базовому хосту.
+            raw_url = urllib.parse.urljoin("https://www.tikwm.com/", raw_url)
+        if urllib.parse.urlsplit(raw_url).scheme.lower() not in ("http", "https"):
+            log.warning("[tiktok] Skipping non-HTTP(S) URL for %s: %r.", url_key, raw_url)
+            continue
         try:
             size_bytes = int(media_data.get(size_key) or 0)
         except (TypeError, ValueError):
@@ -157,14 +165,19 @@ async def _tikwm_throttle() -> None:
         _tikwm_last_request_ts = now
 
 
-async def _read_capped_body(resp: aiohttp.ClientResponse, url: str) -> bytes | None:
+# JSON TikWM обязан быть маленьким: свой кап, а не общий видео-лимит.
+TIKWM_JSON_MAX_BYTES = 1 * 1024 * 1024
+
+
+async def _read_capped_body(resp: aiohttp.ClientResponse, url: str, *, cap_bytes: int | None = None) -> bytes | None:
     # JSON TikWM обязан быть маленьким: читаем с тем же капом, что бинарные
     # скачивания, иначе скомпрометированный прокси льёт в память без края.
+    cap = cap_bytes if cap_bytes is not None else TIKTOK_DOWNLOAD_MAX_BYTES
     content_length = resp.headers.get("Content-Length")
     if content_length is not None:
         try:
-            if int(content_length) > TIKTOK_DOWNLOAD_MAX_BYTES:
-                log.warning("[tikwm] Refusing %s: Content-Length %s exceeds the %d byte cap.", url, content_length, TIKTOK_DOWNLOAD_MAX_BYTES)
+            if int(content_length) > cap:
+                log.warning("[tikwm] Refusing %s: Content-Length %s exceeds the %d byte cap.", url, content_length, cap)
                 return None
         except ValueError:
             pass
@@ -172,8 +185,8 @@ async def _read_capped_body(resp: aiohttp.ClientResponse, url: str) -> bytes | N
     total = 0
     async for chunk in resp.content.iter_chunked(65536):
         total += len(chunk)
-        if total > TIKTOK_DOWNLOAD_MAX_BYTES:
-            log.warning("[tikwm] Aborting %s: exceeded the %d byte cap mid-stream.", url, TIKTOK_DOWNLOAD_MAX_BYTES)
+        if total > cap:
+            log.warning("[tikwm] Aborting %s: exceeded the %d byte cap mid-stream.", url, cap)
             return None
         chunks.append(chunk)
     return b"".join(chunks)
@@ -211,7 +224,7 @@ async def _fetch_tikwm_media_data(
                 async with session.get(api_url, timeout=12, headers=api_headers) as r:
                     if r.status == 200:
                         all_attempts_403_empty = False
-                        raw_body = await _read_capped_body(r, api_url)
+                        raw_body = await _read_capped_body(r, api_url, cap_bytes=TIKWM_JSON_MAX_BYTES)
                         if raw_body is None:
                             continue
                         try:
@@ -224,7 +237,8 @@ async def _fetch_tikwm_media_data(
                             return res.get("data")
                         log.warning("[tikwm] Endpoint %s returned code %s: %s", api_url, res.get("code"), res.get("msg"))
                     else:
-                        body = await r.read()
+                        raw_error = await _read_capped_body(r, api_url, cap_bytes=TIKWM_JSON_MAX_BYTES)
+                        body = raw_error if raw_error is not None else b""
                         if r.status == 403:
                             saw_403 = True
                         if r.status != 403 or body:

lumen_tiktok_flow.py

@@ -20,6 +20,7 @@ from aiogram.types import (
 )
 
 from lumen_tiktok import (
+    TIKTOK_DOWNLOAD_MAX_BYTES,
     _original_sound_label,
     _GENERIC_ORIGINAL_SOUND_PHRASES,
     _chunk_tiktok_media_items,
@@ -210,7 +211,12 @@ async def _fetch_tikwm_media_data_with_proxy_fallback(session: aiohttp.ClientSes
     import bot
     result = None
     for candidate in bot._tikwm_proxy_candidates():
-        result = await bot._fetch_tikwm_media_data(session, resolved_url, headers, proxy_base_url=candidate)
+        try:
+            result = await bot._fetch_tikwm_media_data(session, resolved_url, headers, proxy_base_url=candidate)
+        except Exception as exc:
+            # Упавший кандидат не отменяет остальных: пробуем дальше.
+            log.warning("[tiktok] TikWM proxy candidate %r failed: %s", candidate, exc)
+            continue
         if result is not None:
             return result
     return result
@@ -240,7 +246,11 @@ async def _try_send_tiktok_slideshow(
               async with bot._tiktok_slide_download_semaphore:
                    if post_budget[0] <= 0:
                         return None
-                   data = await bot._download_url_bin(session, slide_url, headers=headers)
+                   # Резерв до скачивания: кап слайда — остаток бюджета.
+                   data = await bot._download_url_bin(
+                        session, slide_url, headers=headers,
+                        cap_bytes=min(TIKTOK_DOWNLOAD_MAX_BYTES, post_budget[0]),
+                   )
                    if data:
                         post_budget[0] -= len(data)
                         if post_budget[0] < 0:

tests/bot_test_helpers.py

@@ -345,16 +345,15 @@ class _FakeTikTokSession:
 
 
 class _FakeTikwmApiResponse:
-    def __init__(self, status=200, json_body=None, body_bytes=b""):
+    def __init__(self, status=200, json_body=None, body_bytes: bytes | None = None):
         import json as _json
         self.status = status
         self._json_body = json_body or {}
-        self._body_bytes = body_bytes
-        # Честный фейк (см. _FakeTikTokResponse): контент из json, если байты не заданы.
+        # Явно заданные байты (включая пустые) едут и в read(), и в стрим,
+        # как у настоящего ответа.
+        self._body_bytes = body_bytes if body_bytes is not None else _json.dumps(self._json_body).encode("utf-8")
         self.headers = {}
-        self.content = _FakeDownloadContent(
-            [body_bytes] if body_bytes else [_json.dumps(self._json_body).encode("utf-8")]
-        )
+        self.content = _FakeDownloadContent([self._body_bytes])
 
     async def __aenter__(self):
         return self

tests/test_bot_tiktok.py

@@ -1667,6 +1667,71 @@ def test_fetch_tikwm_media_data_skips_huge_content_length():
     assert result is None
 
 
+def test_fetch_tikwm_media_data_json_capped_below_video_limit():
+    # JSON капнут своим мегабайтом, а не общим видео-лимитом 75МБ.
+    import lumen_tiktok
+    from tests.bot_test_helpers import _FakeTikwmApiResponse, _FakeTikwmApiSession
+    assert lumen_tiktok.TIKWM_JSON_MAX_BYTES <= 2 * 1024 * 1024
+    big = b'{"code": 0, "data": {}, "pad": "' + b"x" * (lumen_tiktok.TIKWM_JSON_MAX_BYTES + 100) + b'"}'
+    resp = _FakeTikwmApiResponse(status=200, body_bytes=big)
+    result = asyncio.run(bot._fetch_tikwm_media_data(
+        _FakeTikwmApiSession([resp]), "https://www.tiktok.com/@u/video/1", {}))
+    assert result is None
+
+
+def test_fetch_tikwm_media_data_error_branch_reads_capped():
+    # Ветка ошибки тоже читает с капом, а не голым read().
+    from tests.bot_test_helpers import _FakeDownloadContent, _FakeTikwmApiSession
+    calls = []
+
+    class _NoUncappedReadResponse:
+        status = 500
+        headers = {"Content-Length": str(10 ** 12)}
+
+        def __init__(self):
+            self.content = _FakeDownloadContent([b"e" * 100])
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *args):
+            return False
+
+        async def read(self):
+            calls.append("read")
+            return b"e" * 100
+
+    result = asyncio.run(bot._fetch_tikwm_media_data(
+        _FakeTikwmApiSession([_NoUncappedReadResponse()]), "https://www.tiktok.com/@u/video/1", {}))
+    assert result is None
+    assert calls == []
+
+
+def test_tikwm_proxy_fallback_tries_next_candidate_on_exception(monkeypatch):
+    # Исключение первого кандидата не отменяет остальных.
+    import lumen_tiktok_flow
+    calls = []
+
+    async def flaky_fetch(session, resolved_url, headers, proxy_base_url=""):
+        calls.append(proxy_base_url)
+        if len(calls) == 1:
+            raise RuntimeError("proxy down")
+        return {"play": "https://tikwm.com/v.mp4"}
+
+    async def always_boom(session, resolved_url, headers, proxy_base_url=""):
+        raise RuntimeError("all down")
+
+    monkeypatch.setattr(bot, "_tikwm_proxy_candidates", lambda: ["http://p1", "http://p2"])
+    monkeypatch.setattr(bot, "_fetch_tikwm_media_data", flaky_fetch)
+    result = asyncio.run(lumen_tiktok_flow._fetch_tikwm_media_data_with_proxy_fallback(
+        None, "https://www.tiktok.com/@u/video/1", {}))
+    assert result == {"play": "https://tikwm.com/v.mp4"}
+    assert calls == ["http://p1", "http://p2"]
+    monkeypatch.setattr(bot, "_fetch_tikwm_media_data", always_boom)
+    assert asyncio.run(lumen_tiktok_flow._fetch_tikwm_media_data_with_proxy_fallback(
+        None, "https://www.tiktok.com/@u/video/1", {})) is None
+
+
 def test_probe_video_dimensions_survives_na_width(monkeypatch):
     # A7-6: width=N/A раньше ронял float() и вместе с ним валидный duration.
     class _FakeFFprobe:
@@ -1726,6 +1791,17 @@ def test_tiktok_video_candidates_protocol_relative_url():
     assert candidates[0]["url"] == "https://cdn.tikwm.com/v.mp4"
 
 
+def test_tiktok_video_candidates_rejects_userinfo_and_non_string():
+    # Склейка давала userinfo-путаницу: "@evil.com/x" ехал хостом.
+    import urllib.parse
+    candidates = bot._tiktok_video_candidates({"play": "@evil.com/x", "size": 1})
+    assert candidates[0]["url"] == "https://www.tikwm.com/@evil.com/x"
+    assert urllib.parse.urlsplit(candidates[0]["url"]).hostname == "www.tikwm.com"
+    assert bot._tiktok_video_candidates({"play": "javascript:alert(1)", "size": 1}) == []
+    assert bot._tiktok_video_candidates({"play": 12345, "size": 1}) == []
+    assert bot._tiktok_video_candidates({"play": {"u": 1}, "size": 1}) == []
+
+
 def test_download_url_bin_honors_caller_cap():
     # Аудит A7-5: видео режем лимитом отправки (50МБ), а не общим 75МБ.
     mapping = {"http://8.8.8.8/v.mp4": _FakeDownloadResponse([b"123456789012345"])}
@@ -1739,7 +1815,7 @@ def test_slideshow_stops_at_post_budget(monkeypatch):
     # Аудит D1/A7-1: общий кап RAM на пост — остаток слайдов пропускается.
     calls = {}
 
-    async def fake_download(session, url, headers=None):
+    async def fake_download(session, url, headers=None, cap_bytes=None):
         return b"k" * 1024
 
     class _FakeTgBot:
@@ -1776,3 +1852,45 @@ def test_slideshow_stops_at_post_budget(monkeypatch):
     finally:
         bot.chat_state.pop(999985, None)
 
+
+class _SlideshowCapBot:
+    async def send_photo(self, **kwargs):
+        return SimpleNamespace()
+
+    async def send_video(self, **kwargs):
+        return SimpleNamespace()
+
+    async def send_media_group(self, **kwargs):
+        return SimpleNamespace()
+
+
+def test_slideshow_slide_download_reserves_post_budget(monkeypatch):
+    # Кап слайда — остаток общего бюджета до скачивания, а не после.
+    import lumen_tiktok
+    seen_caps = {}
+
+    async def fake_download(session, url, headers=None, cap_bytes=None):
+        seen_caps[url] = cap_bytes
+        return b"k" * 1024
+
+    async def fake_edit(*args, **kwargs):
+        return True
+
+    monkeypatch.setattr(bot, "TIKTOK_SLIDESHOW_MAX_BYTES", 1500)
+    monkeypatch.setattr(bot, "_download_url_bin", fake_download)
+    monkeypatch.setattr(bot, "_edit_message_quietly", fake_edit)
+    monkeypatch.setattr(bot, "_delete_message_quietly", fake_edit)
+    monkeypatch.setattr(bot, "bot", _SlideshowCapBot())
+    incoming = _FakeIncomingMessage(999984)
+    incoming.message_id = 7
+    media_data = {"images": ["http://x/1.jpg", "http://x/2.jpg", "http://x/3.jpg"]}
+    try:
+        asyncio.run(bot._try_send_tiktok_slideshow(
+            None, media_data, incoming, None, "author", {},
+        ))
+        assert seen_caps
+        assert all(cap is not None and cap <= 1500 for cap in seen_caps.values())
+        assert all(cap <= lumen_tiktok.TIKTOK_DOWNLOAD_MAX_BYTES for cap in seen_caps.values())
+    finally:
+        bot.chat_state.pop(999984, None)
+