Merge pull request #3832 from c020627/fix-linkai-bounded-download

fix(linkai): cap the file download in the LinkAI sender

models/linkai/link_ai_bot.py

@@ -16,6 +16,7 @@ from common.log import logger
 from config import conf, pconf
 import threading
 from common import memory, state_dir, utils
+from common.media_download import MAX_FILE_BYTES, download_to_file
 import base64
 import os
 from urllib.parse import urlparse
@@ -515,10 +516,11 @@ def _download_file(url: str):
         # the process CWD, which the packaged desktop app does not control and may
         # not be able to write, and state_dir owns the layout anyway.
         file_path = state_dir.tmp_dir() / file_name
-        response = requests.get(url, timeout=(5, 60))
-        # An expired signed link answers with an HTML error page, not the file.
-        response.raise_for_status()
-        file_path.write_bytes(response.content)
+        # Stream it under the shared cap, the way the LinkAI voice download
+        # already does: the vendor decides the size, and an unbounded body
+        # would be buffered whole before anything could object. A refused
+        # download leaves no partial file behind.
+        download_to_file(url, str(file_path), MAX_FILE_BYTES, timeout=(5, 60))
         return str(file_path)
     except Exception as e:
         logger.warn(e)

tests/test_linkai_bounded_download.py

@@ -0,0 +1,89 @@
+# encoding:utf-8
+"""LinkAI's file sender must cap its download.
+
+``_download_file`` did a bare requests.get and wrote resp.content, so the
+whole body was buffered before anything could object -- ``media_download``'s
+docstring spells out why a socket timeout alone never fires on a server that
+keeps trickling bytes. The LinkAI voice download already routes through the
+shared helper; this file download was the one that did not.
+"""
+
+import pytest
+
+
+class _Stream:
+    """A minimal streamed response: chunks, headers, close."""
+
+    def __init__(self, chunks, headers=None):
+        self._chunks = list(chunks)
+        self.headers = headers or {}
+        # The pre-fix code path read the whole body instead of streaming.
+        self.content = b"".join(self._chunks)
+
+    def raise_for_status(self):
+        return None
+
+    def iter_content(self, chunk_size=1):
+        for chunk in self._chunks:
+            yield chunk
+
+    def close(self):
+        return None
+
+
+def _oversized(chunks=8, size=512):
+    return [b"x" * size] * chunks
+
+
+def _cap(monkeypatch, module, size=1024):
+    """Shrink the shared cap so an oversized body is a few chunks, not 100 MB.
+
+    ``raising=False`` keeps this working on the unfixed code, where the
+    module does not import the constant yet: there the real 100 MB cap
+    applies and the small body sails through -- which is the bug.
+    """
+    monkeypatch.setattr(module, "MAX_FILE_BYTES", size, raising=False)
+
+
+def test_linkai_file_download_refuses_an_oversized_body(tmp_path, monkeypatch):
+    from models.linkai import link_ai_bot
+
+    _cap(monkeypatch, link_ai_bot)
+    monkeypatch.setattr(link_ai_bot.state_dir, "tmp_dir", lambda: tmp_path)
+    monkeypatch.setattr(link_ai_bot.requests, "get",
+                        lambda *a, **k: _Stream(_oversized()))
+
+    assert link_ai_bot._download_file("https://files.example.com/report.pdf") is None
+    # A refused download leaves no partial file behind.
+    assert list(tmp_path.iterdir()) == []
+
+
+def test_linkai_file_download_still_writes_ordinary_bodies(tmp_path, monkeypatch):
+    from models.linkai import link_ai_bot
+
+    monkeypatch.setattr(link_ai_bot.state_dir, "tmp_dir", lambda: tmp_path)
+    monkeypatch.setattr(link_ai_bot.requests, "get",
+                        lambda *a, **k: _Stream([b"%PDF-1.7"]))
+
+    result = link_ai_bot._download_file("https://files.example.com/report.pdf")
+
+    assert result == str(tmp_path / "report.pdf")
+    assert (tmp_path / "report.pdf").read_bytes() == b"%PDF-1.7"
+
+
+@pytest.mark.parametrize("status_code", [403, 404])
+def test_linkai_file_download_refuses_an_error_page(tmp_path, monkeypatch, status_code):
+    """An expired signed link answers with an HTML page, not the file."""
+    from models.linkai import link_ai_bot
+
+    class _Error(_Stream):
+        def raise_for_status(self):
+            raise link_ai_bot.requests.exceptions.HTTPError(
+                f"{status_code} Client Error")
+
+    monkeypatch.setattr(link_ai_bot.state_dir, "tmp_dir", lambda: tmp_path)
+    monkeypatch.setattr(link_ai_bot.requests, "get",
+                        lambda *a, **k: _Error([b"<html>expired</html>"]))
+
+    assert link_ai_bot._download_file("https://files.example.com/report.pdf") is None
+    assert list(tmp_path.iterdir()) == []

tests/test_linkai_download_status.py

@@ -19,6 +19,9 @@ def _download(tmp_path, status_code, content):
     response = requests.Response()
     response.status_code = status_code
     response._content = content
+    # The download is streamed now, so the response has to hand its body out
+    # in chunks; a Response whose content was already read can do that.
+    response._content_consumed = True
     with patch.object(link_ai_bot.state_dir, "tmp_dir", return_value=tmp_path), \
             patch.object(link_ai_bot.requests, "get", return_value=response):
         return link_ai_bot._download_file(URL)

tests/test_linkai_media_urls.py

@@ -53,7 +53,8 @@ def test_downloaded_file_name_excludes_query(monkeypatch, tmp_path):
     from models.linkai import link_ai_bot
 
     # The download checks the status before writing, so the stand-in response
-    # has to carry one the way a real requests.Response does.
+    # has to carry one the way a real requests.Response does. It is streamed
+    # now, so it also has to hand its body out in chunks.
     monkeypatch.setattr(
         link_ai_bot.requests,
         "get",
@@ -61,7 +62,10 @@ def test_downloaded_file_name_excludes_query(monkeypatch, tmp_path):
             "R", (), {
                 "content": b"x",
                 "status_code": 200,
+                "headers": {},
                 "raise_for_status": lambda self: None,
+                "iter_content": lambda self, chunk_size=1: iter([b"x"]),
+                "close": lambda self: None,
             }
         )(),
     )