155 lines
4.7 KiB
Python
155 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
|
os.environ["KEFU_DB_TYPE"] = "sqlite"
|
|
os.environ["KEFU_DATABASE_URL"] = ""
|
|
os.environ["KEFU_DB_PATH"] = str(BACKEND_DIR / "kefu.db")
|
|
if str(BACKEND_DIR) not in sys.path:
|
|
sys.path.insert(0, str(BACKEND_DIR))
|
|
|
|
import main
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
final_host: str,
|
|
headers: dict[str, str] | None = None,
|
|
chunks: tuple[bytes, ...] = (),
|
|
) -> None:
|
|
self.url = SimpleNamespace(host=final_host)
|
|
self.headers = headers or {}
|
|
self._chunks = chunks
|
|
self.raise_for_status = MagicMock()
|
|
self.iteration_started = False
|
|
|
|
async def aiter_bytes(self):
|
|
self.iteration_started = True
|
|
for chunk in self._chunks:
|
|
yield chunk
|
|
|
|
|
|
class _AsyncContext:
|
|
def __init__(self, value) -> None:
|
|
self.value = value
|
|
|
|
async def __aenter__(self):
|
|
return self.value
|
|
|
|
async def __aexit__(self, exc_type, exc, traceback):
|
|
return False
|
|
|
|
|
|
class _FakeAsyncClient:
|
|
def __init__(self, response: _FakeResponse) -> None:
|
|
self.response = response
|
|
self.stream_calls: list[tuple[tuple, dict]] = []
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, traceback):
|
|
return False
|
|
|
|
def stream(self, *args, **kwargs):
|
|
self.stream_calls.append((args, kwargs))
|
|
return _AsyncContext(self.response)
|
|
|
|
|
|
class MediaProxyHostGuardTests(unittest.TestCase):
|
|
def test_allows_each_root_domain_and_its_subdomains(self):
|
|
for allowed in main._MEDIA_PROXY_HOSTS:
|
|
with self.subTest(host=allowed):
|
|
self.assertTrue(main._is_allowed_media_host(allowed))
|
|
self.assertTrue(main._is_allowed_media_host(f"cdn.images.{allowed}"))
|
|
|
|
self.assertTrue(main._is_allowed_media_host("CDN.DOUYINPIC.COM."))
|
|
|
|
def test_rejects_empty_suffix_tricks_and_similar_domains(self):
|
|
rejected = (
|
|
"",
|
|
".",
|
|
"douyin.com.evil",
|
|
"evildouyin.com",
|
|
"byteimg.com.evil.example",
|
|
"evilbyteimg.com",
|
|
"ibyteimg.comevil",
|
|
"douyin.co",
|
|
"douyinpic.co",
|
|
"douyin-static.com",
|
|
"snssdk.example",
|
|
)
|
|
|
|
for hostname in rejected:
|
|
with self.subTest(host=hostname):
|
|
self.assertFalse(main._is_allowed_media_host(hostname))
|
|
|
|
|
|
class MediaProxyResponseGuardTests(unittest.IsolatedAsyncioTestCase):
|
|
async def _call_proxy(self, response: _FakeResponse):
|
|
client = _FakeAsyncClient(response)
|
|
constructor = MagicMock(return_value=client)
|
|
with patch.object(main.httpx, "AsyncClient", constructor):
|
|
result = await main.proxy_media(
|
|
url="https://cdn.douyinpic.com/media/test.jpg"
|
|
)
|
|
return result, client, constructor
|
|
|
|
async def test_rejects_redirect_to_non_allowlisted_final_domain(self):
|
|
response = _FakeResponse(
|
|
final_host="attacker.example",
|
|
headers={"content-type": "image/jpeg"},
|
|
chunks=(b"not-read",),
|
|
)
|
|
|
|
with self.assertRaises(main.HTTPException) as caught:
|
|
await self._call_proxy(response)
|
|
|
|
self.assertEqual(caught.exception.status_code, 400)
|
|
self.assertFalse(response.iteration_started)
|
|
|
|
async def test_rejects_oversized_content_length_before_streaming(self):
|
|
response = _FakeResponse(
|
|
final_host="cdn.douyinpic.com",
|
|
headers={"content-length": "9", "content-type": "image/jpeg"},
|
|
chunks=(b"not-read",),
|
|
)
|
|
|
|
with (
|
|
patch.object(main, "_MEDIA_PROXY_MAX_BYTES", 8),
|
|
self.assertRaises(main.HTTPException) as caught,
|
|
):
|
|
await self._call_proxy(response)
|
|
|
|
self.assertEqual(caught.exception.status_code, 413)
|
|
self.assertFalse(response.iteration_started)
|
|
|
|
async def test_rejects_stream_when_actual_bytes_exceed_limit(self):
|
|
response = _FakeResponse(
|
|
final_host="cdn.douyinpic.com",
|
|
headers={"content-length": "0", "content-type": "image/jpeg"},
|
|
chunks=(b"12345", b"6789"),
|
|
)
|
|
|
|
with (
|
|
patch.object(main, "_MEDIA_PROXY_MAX_BYTES", 8),
|
|
self.assertRaises(main.HTTPException) as caught,
|
|
):
|
|
await self._call_proxy(response)
|
|
|
|
self.assertEqual(caught.exception.status_code, 413)
|
|
self.assertTrue(response.iteration_started)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|