964 lines
34 KiB
Python
964 lines
34 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import io
|
|
import os
|
|
import stat
|
|
import struct
|
|
import zipfile
|
|
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import httpx
|
|
import pytest
|
|
from pypdf import PdfWriter
|
|
from pypdf.generic import (
|
|
DecodedStreamObject,
|
|
DictionaryObject,
|
|
NameObject,
|
|
)
|
|
|
|
from services.runtime.ingestion_files import (
|
|
ClamdScanner,
|
|
FileIngestionError,
|
|
FileIngestionService,
|
|
FileInput,
|
|
FilePolicy,
|
|
PinnedNetworkBackend,
|
|
_AsyncioNetworkStream,
|
|
scan_text_source,
|
|
validate_download_url,
|
|
)
|
|
|
|
Resolver = Callable[[str], Awaitable[Sequence[str]]]
|
|
|
|
|
|
class SequenceResolver:
|
|
def __init__(self, *answers: Sequence[str]) -> None:
|
|
self.answers = list(answers)
|
|
self.hosts: list[str] = []
|
|
|
|
async def __call__(self, host: str) -> Sequence[str]:
|
|
self.hosts.append(host)
|
|
if len(self.answers) == 1:
|
|
return self.answers[0]
|
|
return self.answers.pop(0)
|
|
|
|
|
|
class MappingResolver:
|
|
def __init__(self, records: dict[str, Sequence[str]]) -> None:
|
|
self.records = records
|
|
self.hosts: list[str] = []
|
|
|
|
async def __call__(self, host: str) -> Sequence[str]:
|
|
self.hosts.append(host)
|
|
return self.records.get(host, ("93.184.216.34",))
|
|
|
|
|
|
class FakeScanner:
|
|
def __init__(self, error: Exception | None = None) -> None:
|
|
self.error = error
|
|
self.scanned: list[tuple[Path, bytes, float, int]] = []
|
|
|
|
async def require_clean(self, path: Path, *, deadline: float) -> None:
|
|
self.scanned.append((path, path.read_bytes(), deadline, stat.S_IMODE(path.stat().st_mode)))
|
|
if self.error is not None:
|
|
raise self.error
|
|
|
|
|
|
class RoutingTransport(httpx.AsyncBaseTransport):
|
|
def __init__(self, responses: dict[str, httpx.Response]) -> None:
|
|
self.responses = responses
|
|
self.urls: list[str] = []
|
|
|
|
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
|
url = str(request.url)
|
|
self.urls.append(url)
|
|
response = self.responses[url]
|
|
return httpx.Response(
|
|
response.status_code,
|
|
headers=response.headers,
|
|
content=await response.aread(),
|
|
request=request,
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
|
|
def run(coro: Awaitable[Any]) -> Any:
|
|
return asyncio.run(coro)
|
|
|
|
|
|
def file_input(
|
|
url: str = "https://files.example.test/document.txt", **overrides: str
|
|
) -> dict[str, str]:
|
|
value = {
|
|
"download_url": url,
|
|
"file_id": "file-abc123",
|
|
"mime_type": "text/plain",
|
|
"file_name": "document.txt",
|
|
}
|
|
value.update(overrides)
|
|
return value
|
|
|
|
|
|
def text_response(
|
|
body: bytes, *, content_type: str = "text/plain; charset=utf-8"
|
|
) -> httpx.Response:
|
|
return httpx.Response(200, headers={"content-type": content_type}, content=body)
|
|
|
|
|
|
def make_pdf(text: str, pages: int = 1) -> bytes:
|
|
writer = PdfWriter()
|
|
page = writer.add_blank_page(width=300, height=300)
|
|
font = DictionaryObject(
|
|
{
|
|
NameObject("/Type"): NameObject("/Font"),
|
|
NameObject("/Subtype"): NameObject("/Type1"),
|
|
NameObject("/BaseFont"): NameObject("/Helvetica"),
|
|
}
|
|
)
|
|
page[NameObject("/Resources")] = DictionaryObject(
|
|
{
|
|
NameObject("/Font"): DictionaryObject({NameObject("/F1"): font}),
|
|
}
|
|
)
|
|
content = DecodedStreamObject()
|
|
content.set_data(f"BT /F1 12 Tf 20 200 Td ({text}) Tj ET".encode("ascii"))
|
|
page[NameObject("/Contents")] = content
|
|
for _ in range(pages - 1):
|
|
writer.add_blank_page(width=300, height=300)
|
|
buffer = io.BytesIO()
|
|
writer.write(buffer)
|
|
return buffer.getvalue()
|
|
|
|
|
|
def make_docx(text: str = "", *, document_xml: bytes | None = None) -> bytes:
|
|
if document_xml is None:
|
|
document_xml = (
|
|
"<?xml version='1.0' encoding='UTF-8'?>"
|
|
"<w:document xmlns:w='http://schemas.openxmlformats.org/wordprocessingml/2006/main'>"
|
|
f"<w:body><w:p><w:r><w:t>{text}</w:t></w:r></w:p></w:body></w:document>"
|
|
).encode()
|
|
content_types = (
|
|
b"<?xml version='1.0' encoding='UTF-8'?>"
|
|
b"<Types xmlns='http://schemas.openxmlformats.org/package/2006/content-types'>"
|
|
b"<Override PartName='/word/document.xml' "
|
|
b"ContentType='application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml'/>"
|
|
b"</Types>"
|
|
)
|
|
buffer = io.BytesIO()
|
|
with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive:
|
|
archive.writestr("[Content_Types].xml", content_types)
|
|
archive.writestr("word/document.xml", document_xml)
|
|
return buffer.getvalue()
|
|
|
|
|
|
def make_transport(url: str, body: bytes, *, content_type: str = "text/plain") -> RoutingTransport:
|
|
return RoutingTransport(
|
|
{url: httpx.Response(200, headers={"content-type": content_type}, content=body)}
|
|
)
|
|
|
|
|
|
def make_policy(
|
|
tmp_path: Path,
|
|
scanner: FakeScanner | None = None,
|
|
**overrides: Any,
|
|
) -> FilePolicy:
|
|
return FilePolicy(scanner=scanner or FakeScanner(), temp_directory=tmp_path, **overrides)
|
|
|
|
|
|
def test_rejects_insecure_or_credentialed_download_urls() -> None:
|
|
async def scenario() -> None:
|
|
resolver = MappingResolver({})
|
|
for value in (
|
|
"http://files.example.test/a.txt",
|
|
"https://user:password@files.example.test/a.txt",
|
|
"https://files.example.test/a.txt#fragment",
|
|
"https://files.example.test:8443/a.txt",
|
|
"//files.example.test/a.txt",
|
|
):
|
|
with pytest.raises(FileIngestionError):
|
|
await validate_download_url(value, resolver=resolver)
|
|
assert resolver.hosts == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_text_source_scanner_returns_bounded_digest_and_securely_removes_temporary_copy(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
async def scenario() -> None:
|
|
scanner = FakeScanner()
|
|
policy = make_policy(tmp_path, scanner, max_text_chars_per_file=128)
|
|
source = await scan_text_source("Distinctive repo text", "guide.md", policy=policy)
|
|
|
|
assert source.text == "Distinctive repo text"
|
|
assert source.content == b"Distinctive repo text"
|
|
assert source.sha256 == hashlib.sha256(source.content).hexdigest()
|
|
assert scanner.scanned[0][1] == source.content
|
|
if os.name != "nt":
|
|
assert scanner.scanned[0][3] == 0o600
|
|
assert not scanner.scanned[0][0].exists()
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_text_source_scanner_fails_closed_for_missing_scanner_and_bounds(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
with pytest.raises(FileIngestionError, match="Malware scanning is unavailable"):
|
|
await scan_text_source(
|
|
"safe", "guide.md", policy=FilePolicy(scanner=None, temp_directory=tmp_path)
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await scan_text_source(
|
|
"too long",
|
|
"guide.md",
|
|
policy=make_policy(tmp_path, max_text_chars_per_file=4),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await scan_text_source("bad\x00text", "guide.md", policy=make_policy(tmp_path))
|
|
|
|
run(scenario())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"address",
|
|
[
|
|
"127.0.0.1",
|
|
"198.19.5.6",
|
|
"169.254.169.254",
|
|
"192.0.2.4",
|
|
"224.0.0.1",
|
|
"::1",
|
|
"fe80::1",
|
|
"fc00::1",
|
|
"2001:db8::1",
|
|
],
|
|
)
|
|
def test_rejects_non_public_dns_answers(address: str) -> None:
|
|
async def scenario() -> None:
|
|
with pytest.raises(FileIngestionError):
|
|
await validate_download_url(
|
|
"https://files.example.test/a.txt",
|
|
resolver=MappingResolver({"files.example.test": (address,)}),
|
|
)
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_pinned_network_backend_rejects_dns_rebinding_before_connect() -> None:
|
|
async def scenario() -> None:
|
|
resolver = SequenceResolver(("93.184.216.34",), ("198.19.0.7",))
|
|
await validate_download_url("https://files.example.test/a.txt", resolver=resolver)
|
|
backend = PinnedNetworkBackend(resolver=resolver)
|
|
with pytest.raises(FileIngestionError):
|
|
await backend.connect_tcp("files.example.test", 443, timeout=0.1)
|
|
assert resolver.hosts == ["files.example.test", "files.example.test"]
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_tls_close_aborts_when_peer_never_completes_shutdown() -> None:
|
|
async def scenario() -> None:
|
|
class FakeTransport:
|
|
def __init__(self) -> None:
|
|
self.aborted = False
|
|
|
|
def abort(self) -> None:
|
|
self.aborted = True
|
|
|
|
class StalledWriter:
|
|
def __init__(self) -> None:
|
|
self.transport = FakeTransport()
|
|
self.closed = False
|
|
self.wait_started = asyncio.Event()
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
async def wait_closed(self) -> None:
|
|
self.wait_started.set()
|
|
await asyncio.Future()
|
|
|
|
writer = StalledWriter()
|
|
stream = _AsyncioNetworkStream(None, writer) # type: ignore[arg-type]
|
|
await asyncio.wait_for(stream.aclose(), timeout=2.0)
|
|
assert writer.wait_started.is_set()
|
|
assert writer.closed
|
|
assert writer.transport.aborted
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_tls_close_preserves_fast_peer_shutdown() -> None:
|
|
async def scenario() -> None:
|
|
class FakeTransport:
|
|
def __init__(self) -> None:
|
|
self.aborted = False
|
|
|
|
def abort(self) -> None:
|
|
self.aborted = True
|
|
|
|
class FastWriter:
|
|
def __init__(self) -> None:
|
|
self.transport = FakeTransport()
|
|
self.closed = False
|
|
self.waited = False
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
async def wait_closed(self) -> None:
|
|
self.waited = True
|
|
|
|
writer = FastWriter()
|
|
stream = _AsyncioNetworkStream(None, writer) # type: ignore[arg-type]
|
|
await stream.aclose()
|
|
assert writer.closed
|
|
assert writer.waited
|
|
assert not writer.transport.aborted
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_tls_close_aborts_and_reraises_cancellation() -> None:
|
|
async def scenario() -> None:
|
|
class FakeTransport:
|
|
def __init__(self) -> None:
|
|
self.aborted = False
|
|
|
|
def abort(self) -> None:
|
|
self.aborted = True
|
|
|
|
class StalledWriter:
|
|
def __init__(self) -> None:
|
|
self.transport = FakeTransport()
|
|
self.wait_started = asyncio.Event()
|
|
|
|
def close(self) -> None:
|
|
return None
|
|
|
|
async def wait_closed(self) -> None:
|
|
self.wait_started.set()
|
|
await asyncio.Future()
|
|
|
|
writer = StalledWriter()
|
|
stream = _AsyncioNetworkStream(None, writer) # type: ignore[arg-type]
|
|
closing = asyncio.create_task(stream.aclose())
|
|
await writer.wait_started.wait()
|
|
closing.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await closing
|
|
assert writer.transport.aborted
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_tls_close_aborts_when_wait_closed_fails() -> None:
|
|
async def scenario() -> None:
|
|
class FakeTransport:
|
|
def __init__(self) -> None:
|
|
self.aborted = False
|
|
|
|
def abort(self) -> None:
|
|
self.aborted = True
|
|
|
|
class BrokenWriter:
|
|
def __init__(self) -> None:
|
|
self.transport = FakeTransport()
|
|
|
|
def close(self) -> None:
|
|
return None
|
|
|
|
async def wait_closed(self) -> None:
|
|
raise OSError("simulated close failure")
|
|
|
|
writer = BrokenWriter()
|
|
stream = _AsyncioNetworkStream(None, writer) # type: ignore[arg-type]
|
|
await stream.aclose()
|
|
assert writer.transport.aborted
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_stalled_tls_close_keeps_downloaded_body_on_scanned_path(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
payload = b"Complete body before the stalled TLS shutdown."
|
|
scanner = FakeScanner()
|
|
|
|
class FakeTransport:
|
|
def __init__(self) -> None:
|
|
self.aborted = False
|
|
|
|
def abort(self) -> None:
|
|
self.aborted = True
|
|
|
|
class StalledWriter:
|
|
def __init__(self) -> None:
|
|
self.transport = FakeTransport()
|
|
self.wait_started = asyncio.Event()
|
|
|
|
def close(self) -> None:
|
|
return None
|
|
|
|
async def wait_closed(self) -> None:
|
|
self.wait_started.set()
|
|
await asyncio.Future()
|
|
|
|
writer = StalledWriter()
|
|
network_stream = _AsyncioNetworkStream(None, writer) # type: ignore[arg-type]
|
|
|
|
class ClosingBody(httpx.AsyncByteStream):
|
|
async def __aiter__(self) -> AsyncIterator[bytes]:
|
|
yield payload
|
|
|
|
async def aclose(self) -> None:
|
|
await network_stream.aclose()
|
|
|
|
class ResponseTransport(httpx.AsyncBaseTransport):
|
|
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
headers={"content-type": "text/plain"},
|
|
stream=ClosingBody(),
|
|
request=request,
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, scanner=scanner),
|
|
resolver=MappingResolver({}),
|
|
transport=ResponseTransport(),
|
|
)
|
|
processed = await service.acquire([cast(FileInput, file_input())], actor=object())
|
|
assert len(processed) == 1
|
|
assert processed[0].content == payload
|
|
assert processed[0].text == payload.decode()
|
|
assert [entry[1] for entry in scanner.scanned] == [payload]
|
|
assert writer.transport.aborted
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_revalidates_each_redirect_before_following_it(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
first = "https://files.example.test/a"
|
|
private = "https://127.0.0.1/private"
|
|
transport = RoutingTransport({first: httpx.Response(302, headers={"location": private})})
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path),
|
|
resolver=MappingResolver({}),
|
|
transport=transport,
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await service.acquire([file_input(first)], actor=object())
|
|
assert transport.urls == [first]
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_rejects_redirects_beyond_the_configured_limit(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
urls = [f"https://files{i}.example.test/a" for i in range(3)]
|
|
transport = RoutingTransport(
|
|
{
|
|
urls[0]: httpx.Response(302, headers={"location": urls[1]}),
|
|
urls[1]: httpx.Response(302, headers={"location": urls[2]}),
|
|
urls[2]: text_response(b"never fetched"),
|
|
}
|
|
)
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, max_redirects=1),
|
|
resolver=MappingResolver({}),
|
|
transport=transport,
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await service.acquire([file_input(urls[0])], actor=object())
|
|
assert transport.urls == urls[:2]
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_streams_within_byte_limits_and_returns_scanned_original_bytes(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.txt"
|
|
body = b"source bytes\n"
|
|
scanner = FakeScanner()
|
|
transport = make_transport(url, body)
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, scanner, max_file_bytes=100, max_total_bytes=100),
|
|
resolver=MappingResolver({}),
|
|
transport=transport,
|
|
)
|
|
(source,) = await service.acquire([file_input(url)], actor=object())
|
|
assert source.content == body
|
|
assert source.byte_size == len(body)
|
|
assert source.sha256 == hashlib.sha256(body).hexdigest()
|
|
assert source.text == "source bytes\n"
|
|
scanned_path, scanned_content, deadline, mode = scanner.scanned[0]
|
|
assert scanned_content == body
|
|
assert deadline > 0
|
|
if os.name != "nt":
|
|
assert mode == 0o600
|
|
assert scanned_path.parent.parent == tmp_path
|
|
assert not scanned_path.exists()
|
|
assert b"source bytes" not in repr(source).encode()
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_cumulative_size_limit_stops_before_downloading_another_source(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
first = "https://files.example.test/one.txt"
|
|
second = "https://files.example.test/two.txt"
|
|
transport = RoutingTransport(
|
|
{first: text_response(b"1234"), second: text_response(b"5678")}
|
|
)
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, max_file_bytes=10, max_total_bytes=6),
|
|
resolver=MappingResolver({}),
|
|
transport=transport,
|
|
)
|
|
files = [file_input(first), file_input(second, file_id="file-def456")]
|
|
with pytest.raises(FileIngestionError):
|
|
await service.acquire(files, actor=object())
|
|
assert transport.urls == [first, second]
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_rejects_oversize_stream_even_when_content_length_is_missing(tmp_path: Path) -> None:
|
|
class Chunks(httpx.AsyncByteStream):
|
|
async def __aiter__(self):
|
|
yield b"12345"
|
|
yield b"67890"
|
|
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.txt"
|
|
|
|
class StreamingTransport(httpx.AsyncBaseTransport):
|
|
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
headers={"content-type": "text/plain"},
|
|
stream=Chunks(),
|
|
request=request,
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, max_file_bytes=8),
|
|
resolver=MappingResolver({}),
|
|
transport=StreamingTransport(),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await service.acquire([file_input(url)], actor=object())
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_scan_is_required_and_failure_removes_ephemeral_file(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.txt"
|
|
scanner = FakeScanner(error=RuntimeError("scanner offline"))
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, scanner),
|
|
resolver=MappingResolver({}),
|
|
transport=make_transport(url, b"private source text"),
|
|
)
|
|
with pytest.raises(FileIngestionError) as raised:
|
|
await service.acquire([file_input(url)], actor=object())
|
|
assert "private source text" not in str(raised.value)
|
|
assert "files.example.test" not in str(raised.value)
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
missing = FileIngestionService(
|
|
FilePolicy(scanner=None, temp_directory=tmp_path),
|
|
resolver=MappingResolver({}),
|
|
transport=make_transport(url, b"source"),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await missing.acquire([file_input(url)], actor=object())
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_rejects_claimed_mime_mismatch_and_arbitrary_archives(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.txt"
|
|
valid_pdf = make_pdf("text")
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path),
|
|
resolver=MappingResolver({}),
|
|
transport=make_transport(url, valid_pdf, content_type="application/pdf"),
|
|
)
|
|
with pytest.raises(FileIngestionError) as raised:
|
|
await service.acquire([file_input(url, mime_type="text/plain")], actor=object())
|
|
assert raised.value.code == "mime_mismatch"
|
|
|
|
zip_buffer = io.BytesIO()
|
|
with zipfile.ZipFile(zip_buffer, "w") as archive:
|
|
archive.writestr("payload.txt", "not a docx")
|
|
archive_service = FileIngestionService(
|
|
make_policy(tmp_path),
|
|
resolver=MappingResolver({}),
|
|
transport=make_transport(url, zip_buffer.getvalue(), content_type="application/zip"),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await archive_service.acquire(
|
|
[file_input(url, mime_type="application/zip")], actor=object()
|
|
)
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"document_xml",
|
|
[
|
|
b"<w:document>",
|
|
(
|
|
b"<!DOCTYPE w:document [<!ENTITY xxe SYSTEM 'file:///private/source'>]>"
|
|
b"<w:document xmlns:w='http://schemas.openxmlformats.org/wordprocessingml/2006/main'>"
|
|
b"<w:body><w:p><w:r><w:t>&xxe;</w:t></w:r></w:p></w:body></w:document>"
|
|
),
|
|
],
|
|
)
|
|
def test_rejects_malformed_and_entity_bearing_docx_xml(
|
|
tmp_path: Path,
|
|
document_xml: bytes,
|
|
) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.docx"
|
|
content = make_docx(document_xml=document_xml)
|
|
mime = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path),
|
|
resolver=MappingResolver({}),
|
|
transport=make_transport(url, content, content_type=mime),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await service.acquire(
|
|
[file_input(url, mime_type=mime, file_name="document.docx")],
|
|
actor=object(),
|
|
)
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_extracts_bounded_utf8_text_pdf_and_docx(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
inputs = [
|
|
(
|
|
"https://files.example.test/a.txt",
|
|
file_input("https://files.example.test/a.txt"),
|
|
b"plain text",
|
|
"text/plain",
|
|
),
|
|
(
|
|
"https://files.example.test/b.pdf",
|
|
file_input(
|
|
"https://files.example.test/b.pdf",
|
|
file_id="file-pdf",
|
|
mime_type="application/pdf",
|
|
file_name="b.pdf",
|
|
),
|
|
make_pdf("pdf text"),
|
|
"application/pdf",
|
|
),
|
|
(
|
|
"https://files.example.test/c.docx",
|
|
file_input(
|
|
"https://files.example.test/c.docx",
|
|
file_id="file-docx",
|
|
mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
file_name="c.docx",
|
|
),
|
|
make_docx("docx text"),
|
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
),
|
|
]
|
|
transport = RoutingTransport(
|
|
{url: text_response(body, content_type=mime) for url, _, body, mime in inputs}
|
|
)
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path), resolver=MappingResolver({}), transport=transport
|
|
)
|
|
sources = await service.acquire([item for _, item, _, _ in inputs], actor=object())
|
|
assert [source.text.strip() for source in sources] == [
|
|
"plain text",
|
|
"pdf text",
|
|
"docx text",
|
|
]
|
|
assert [source.content for source in sources] == [body for _, _, body, _ in inputs]
|
|
assert [source.byte_size for source in sources] == [len(body) for _, _, body, _ in inputs]
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_image_only_pdf_without_extractable_text_is_rejected(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/image-only.pdf"
|
|
scanner = FakeScanner()
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, scanner),
|
|
resolver=MappingResolver({"files.example.test": ("93.184.216.34",)}),
|
|
transport=make_transport(url, make_pdf(""), content_type="application/pdf"),
|
|
)
|
|
with pytest.raises(FileIngestionError) as caught:
|
|
await service.acquire(
|
|
[file_input(url, mime_type="application/pdf", file_name="image-only.pdf")],
|
|
actor=object(),
|
|
)
|
|
assert caught.value.code == "extraction_failed"
|
|
assert scanner.scanned
|
|
assert not scanner.scanned[0][0].exists()
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_enforces_pdf_page_and_extracted_text_limits(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/two.pdf"
|
|
two_page_pdf = make_pdf("page", pages=2)
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, max_pdf_pages=1),
|
|
resolver=MappingResolver({}),
|
|
transport=make_transport(url, two_page_pdf, content_type="application/pdf"),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await service.acquire(
|
|
[file_input(url, mime_type="application/pdf", file_name="two.pdf")],
|
|
actor=object(),
|
|
)
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
text_url = "https://files.example.test/text.txt"
|
|
text_service = FileIngestionService(
|
|
make_policy(tmp_path, max_text_chars_per_file=4),
|
|
resolver=MappingResolver({}),
|
|
transport=make_transport(text_url, b"12345"),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await text_service.acquire([file_input(text_url)], actor=object())
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_docx_limits_reject_zip_bombs_and_relationship_overflow(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.docx"
|
|
compressed = make_docx("A" * 50_000)
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, max_docx_uncompressed_bytes=1024),
|
|
resolver=MappingResolver({}),
|
|
transport=make_transport(
|
|
url,
|
|
compressed,
|
|
content_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await service.acquire(
|
|
[
|
|
file_input(
|
|
url,
|
|
mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
)
|
|
],
|
|
actor=object(),
|
|
)
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_exact_file_input_shape_duplicate_ids_and_count_are_enforced(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.txt"
|
|
transport = make_transport(url, b"body")
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, max_files=2),
|
|
resolver=MappingResolver({}),
|
|
transport=transport,
|
|
)
|
|
with pytest.raises(FileIngestionError) as unknown_fields:
|
|
await service.acquire([file_input(url, unexpected="field")], actor=object())
|
|
assert unknown_fields.value.code == "invalid_input"
|
|
assert unknown_fields.value.input_reason == "attachment_unknown_fields"
|
|
assert "unexpected" not in str(unknown_fields.value)
|
|
|
|
opaque_prefix = "reference with spaces (sample caf" + chr(0x00E9) + " #1)/"
|
|
opaque_file_id = opaque_prefix + "x" * (128 - len(opaque_prefix))
|
|
sources = await service.acquire([file_input(url, file_id=opaque_file_id)], actor=object())
|
|
assert sources[0].file_id == opaque_file_id
|
|
|
|
with pytest.raises(FileIngestionError) as duplicate_ids:
|
|
await service.acquire([file_input(url), file_input(url)], actor=object())
|
|
assert duplicate_ids.value.code == "invalid_input"
|
|
assert duplicate_ids.value.input_reason == "duplicate_file_ids"
|
|
|
|
with pytest.raises(FileIngestionError) as too_many:
|
|
await service.acquire(
|
|
[
|
|
file_input(url, file_id="file-first"),
|
|
file_input(url, file_id="file-second"),
|
|
file_input(url, file_id="file-third"),
|
|
],
|
|
actor=object(),
|
|
)
|
|
assert too_many.value.code == "source_limit"
|
|
assert transport.urls == [url]
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_file_id_validation_rejects_empty_overlong_and_non_printable_values(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.txt"
|
|
transport = RoutingTransport({})
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path),
|
|
resolver=MappingResolver({}),
|
|
transport=transport,
|
|
)
|
|
invalid_ids = (
|
|
("", "attachment_file_id_empty"),
|
|
(" ", "attachment_file_id_empty"),
|
|
("x" * 129, "attachment_file_id_too_long"),
|
|
("reference\x00with control", "attachment_file_id_non_printable"),
|
|
("reference\nwith control", "attachment_file_id_non_printable"),
|
|
("reference\x85with control", "attachment_file_id_non_printable"),
|
|
)
|
|
for file_id, expected_reason in invalid_ids:
|
|
with pytest.raises(FileIngestionError) as invalid:
|
|
await service.acquire([file_input(url, file_id=file_id)], actor=object())
|
|
assert invalid.value.code == "invalid_input"
|
|
assert invalid.value.input_reason == expected_reason
|
|
assert str(invalid.value) == "File input is invalid."
|
|
assert transport.urls == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_enforces_download_deadline_and_removes_temporary_files(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
url = "https://files.example.test/document.txt"
|
|
|
|
class SlowTransport(httpx.AsyncBaseTransport):
|
|
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
|
await asyncio.sleep(0.25)
|
|
return httpx.Response(
|
|
200,
|
|
headers={"content-type": "text/plain"},
|
|
content=b"late",
|
|
request=request,
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
service = FileIngestionService(
|
|
make_policy(tmp_path, download_timeout_seconds=0.1),
|
|
resolver=MappingResolver({}),
|
|
transport=SlowTransport(),
|
|
)
|
|
with pytest.raises(FileIngestionError):
|
|
await service.acquire([file_input(url)], actor=object())
|
|
assert list(tmp_path.iterdir()) == []
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_clamd_scanner_accepts_only_exact_clean_response(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
path = tmp_path / "source.txt"
|
|
path.write_bytes(b"clamav payload")
|
|
received = bytearray()
|
|
|
|
async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
assert await reader.readexactly(10) == b"zINSTREAM\0"
|
|
while True:
|
|
length = struct.unpack("!I", await reader.readexactly(4))[0]
|
|
if length == 0:
|
|
break
|
|
received.extend(await reader.readexactly(length))
|
|
writer.write(b"stream: OK\0")
|
|
await writer.drain()
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
server = await asyncio.start_server(handle, "127.0.0.1", 0)
|
|
try:
|
|
port = server.sockets[0].getsockname()[1]
|
|
await ClamdScanner("127.0.0.1", port=port).require_clean(
|
|
path, deadline=asyncio.get_running_loop().time() + 1
|
|
)
|
|
finally:
|
|
server.close()
|
|
await server.wait_closed()
|
|
assert bytes(received) == b"clamav payload"
|
|
|
|
run(scenario())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"answer", [b"stream: Eicar-Test-Signature FOUND\0", b"ERROR\0", b"stream: OK with extra data\0"]
|
|
)
|
|
def test_clamd_scanner_fails_closed_on_hit_or_malformed_response(
|
|
tmp_path: Path, answer: bytes
|
|
) -> None:
|
|
async def scenario() -> None:
|
|
path = tmp_path / "source.txt"
|
|
path.write_bytes(b"payload")
|
|
|
|
async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
await reader.readexactly(10)
|
|
while True:
|
|
length = struct.unpack("!I", await reader.readexactly(4))[0]
|
|
if length == 0:
|
|
break
|
|
await reader.readexactly(length)
|
|
writer.write(answer)
|
|
await writer.drain()
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
server = await asyncio.start_server(handle, "127.0.0.1", 0)
|
|
try:
|
|
port = server.sockets[0].getsockname()[1]
|
|
with pytest.raises(FileIngestionError):
|
|
await ClamdScanner("127.0.0.1", port=port).require_clean(
|
|
path, deadline=asyncio.get_running_loop().time() + 1
|
|
)
|
|
finally:
|
|
server.close()
|
|
await server.wait_closed()
|
|
|
|
run(scenario())
|
|
|
|
|
|
def test_environment_policy_rejects_unbounded_and_invalid_limits(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("INGESTION_MAX_FILES", "1000000")
|
|
with pytest.raises(ValueError):
|
|
FilePolicy.from_environment(scanner=FakeScanner())
|
|
monkeypatch.setenv("INGESTION_MAX_FILES", "not-an-int")
|
|
with pytest.raises(ValueError):
|
|
FilePolicy.from_environment(scanner=FakeScanner())
|