Coverage for slidge/core/attachment_upload.py: 83%
229 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-18 04:30 +0000
1"""Upload attachments to an HTTP upload service (:xep:`0363`)."""
3from __future__ import annotations
5import contextlib
6import io
7import logging
8import os
9import shutil
10import stat
11import tempfile
12import warnings
13from collections.abc import AsyncIterator
14from mimetypes import guess_type
15from pathlib import Path
16from typing import IO, TYPE_CHECKING, cast
17from urllib.parse import quote as urlquote
18from uuid import uuid4
20import aiohttp
21from slixmpp import JID, Iq
23from ..db.models import Attachment
24from ..util.types import LegacyAttachment
25from ..util.util import fix_namespaces, fix_suffix
26from . import config
28if TYPE_CHECKING:
29 from .gateway import BaseGateway
30 from .session import BaseSession
33class AttachmentUploader:
34 """Turns a :class:`.LegacyAttachment` into a URL that XMPP clients can fetch."""
36 xmpp: BaseGateway
37 session: BaseSession | None
38 """
39 The session uplaoding attachments, or :const:`None` for the gateway component,
40 which is not bound to a :term:`User`.
41 """
43 def __init__(self, xmpp: BaseGateway, session: BaseSession | None) -> None:
44 self.xmpp = xmpp
45 self.session = session
47 @contextlib.asynccontextmanager
48 async def dedup_lock(
49 self,
50 attachment: LegacyAttachment | Path | str,
51 ) -> AsyncIterator[None]:
52 """Take a lock on an attachment with a given name.
54 Prevents races which download the same attachment several times."""
55 session = self.session
56 if (
57 session is None
58 or not isinstance(attachment, LegacyAttachment)
59 or attachment.legacy_file_id is None
60 ):
61 yield
62 else:
63 async with session.lock(("attachment", attachment.legacy_file_id)):
64 yield
66 async def get_stored(self, attachment: LegacyAttachment) -> Attachment:
67 """
68 Fetch the :class:`.Attachment` already uploaded for this user, if any,
69 or a new (transient) one.
70 """
71 session = self.session
72 if attachment.legacy_file_id is not None and session is not None:
73 with self.xmpp.store.session() as orm:
74 stored = (
75 orm.query(Attachment)
76 .filter_by(
77 legacy_file_id=str(attachment.legacy_file_id),
78 user_account_id=session.user_pk,
79 )
80 .one_or_none()
81 )
82 if stored is not None:
83 if not await self.__valid_url(session, stored.url):
84 stored.url = None # type:ignore
85 return stored
86 return Attachment(
87 user_account_id=None if session is None else session.user_pk,
88 legacy_file_id=None
89 if attachment.legacy_file_id is None
90 else str(attachment.legacy_file_id),
91 url=attachment.url if config.USE_ATTACHMENT_ORIGINAL_URLS else None,
92 )
94 def record(self, stored: Attachment) -> None:
95 """Remember an uploaded attachment, so that it is not uploaded twice.
97 No-op without a session, since attachments are stored per user.
98 """
99 # TODO: we need a separate mechanism to record gateway component attachments.
100 if self.session is None:
101 return
102 with self.xmpp.store.session(expire_on_commit=False) as orm:
103 orm.add(stored)
104 orm.commit()
106 @staticmethod
107 async def __valid_url(session: BaseSession, url: str) -> bool:
108 async with session.http.head(url) as r:
109 return r.status < 400
111 async def get_url(self, att: LegacyAttachment, stored: Attachment) -> str:
112 att = _ensure_name(att)
114 if len(att.name) > config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH:
115 log.debug("Trimming long filename: %s", att.name)
116 base, ext = os.path.splitext(att.name)
117 att.name = (
118 base[: config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH - len(ext)] + ext
119 )
121 if config.FIX_FILENAME_SUFFIX_MIME_TYPE and isinstance(att.path, Path):
122 att.name, att.content_type = fix_suffix(
123 att.path, att.content_type, att.name
124 )
126 att.legacy_file_id = stored.legacy_file_id
128 if config.NO_UPLOAD_PATH:
129 att.path, new_url = await self.__no_upload(att, stored.legacy_file_id)
130 new_url = (config.NO_UPLOAD_URL_PREFIX or "") + "/" + urlquote(new_url)
131 else:
132 new_url = await self.__upload(att)
134 if stored.legacy_file_id:
135 stored.url = new_url
137 return new_url
139 async def __upload(self, att: _AttachmentWithName) -> str:
140 assert config.UPLOAD_SERVICE
141 att = await _ensure_metadata(att)
142 iq_slot = await self.__request_upload_slot(
143 config.UPLOAD_SERVICE,
144 att.name,
145 att.size,
146 att.content_type,
147 )
148 if iq_slot["type"] == "error":
149 # COMPAT: In theory, __request_upload_slot() raises IqError, but
150 # there is a bug in prosody's mod_privilege where the outer
151 # IQ type is (illegally) not set to error when the inner IQ
152 # has type='error'.
153 raise RuntimeError(f"Error while requesting upload slot: {iq_slot}")
154 slot = iq_slot.get_plugin("http_upload_slot", check=True)
155 if slot is None:
156 raise RuntimeError(f"No upload slot in this IQ: {iq_slot}")
157 put = slot["put"]["url"]
158 assert isinstance(put, str)
159 if not put:
160 raise RuntimeError(f"Cannot find a PUT URL in: {slot}")
161 get = slot["get"]["url"]
162 assert isinstance(get, str)
163 if not get:
164 raise RuntimeError(f"Cannot find a GET URL in: {slot}")
165 headers = {
166 "Content-Length": str(att.size),
167 "Content-Type": att.content_type,
168 **{header["name"]: header["value"] for header in slot["put"]["headers"]},
169 }
171 async with (
172 aiohttp.ClientSession() as http,
173 _get_data(att, http) as data,
174 http.put(slot["put"]["url"], data=data, headers=headers) as resp,
175 ):
176 resp.raise_for_status()
178 return get
180 async def __request_upload_slot(
181 self,
182 upload_service: JID | str,
183 filename: str,
184 size: int,
185 content_type: str,
186 ) -> Iq:
187 iq_request = self.xmpp.make_iq_get(ito=upload_service)
188 request = iq_request["http_upload_request"]
189 request["filename"] = filename
190 request["size"] = str(size)
191 request["content-type"] = content_type
192 session = self.session
193 if session is not None:
194 iq_request.set_from(session.user_jid)
195 try:
196 return await self.xmpp.plugin["xep_0356"].send_privileged_iq(iq_request)
197 except Exception as e: # noqa: BLE001
198 warnings.warn(
199 "Could not request upload slot on behalf of "
200 f"{session.user_jid}: {e}."
201 "Falling back to not using privileges."
202 )
203 fix_namespaces(iq_request.xml, "jabber:client", "jabber:component:accept")
204 iq_request.set_from(config.UPLOAD_REQUESTER or self.xmpp.boundjid)
205 return await iq_request.send() # type:ignore[no-any-return]
207 @staticmethod
208 async def __no_upload(
209 att: _AttachmentWithName, legacy_file_id: str | None
210 ) -> tuple[Path, str]:
211 file_id = uuid4().hex if legacy_file_id is None else legacy_file_id
212 assert config.NO_UPLOAD_PATH is not None
213 assert config.NO_UPLOAD_URL_PREFIX is not None
214 destination_dir = Path(config.NO_UPLOAD_PATH) / file_id
216 if destination_dir.exists():
217 log.debug("Dest dir exists: %s", destination_dir)
218 files = [f for f in destination_dir.glob("**/*") if f.is_file()]
219 if len(files) == 1:
220 log.debug(
221 "Found the legacy attachment '%s' at '%s'",
222 legacy_file_id,
223 files[0],
224 )
225 name = files[0].name
226 uu = files[0].parent.name # anti-obvious url trick, see below
227 return files[0], f"{file_id}/{uu}/{name}"
228 else:
229 log.warning(
230 (
231 "There are several or zero files in %s, "
232 "slidge doesn't know which one to pick among %s. "
233 "Removing the dir."
234 ),
235 destination_dir,
236 files,
237 )
238 shutil.rmtree(destination_dir)
240 log.debug("Did not find a file in: %s", destination_dir)
241 # let's use a UUID to avoid URLs being too obvious
242 uu = str(uuid4())
243 destination_dir = destination_dir / uu
244 destination_dir.mkdir(parents=True)
246 assert att.name
247 destination = destination_dir / att.name
248 if att.path:
249 assert isinstance(att.path, Path)
250 try:
251 destination.hardlink_to(att.path)
252 except OSError as e:
253 if is_temp_path(att.path):
254 shutil.copy2(att.path, destination)
255 else:
256 log.debug("Could not hardlink: %s, attempting symlink", e)
257 try:
258 destination.symlink_to(att.path)
259 except OSError as e:
260 log.debug("Could not symlink: %s, copying data", e)
261 shutil.copy2(att.path, destination)
262 elif att.data:
263 destination.write_bytes(att.data)
264 else:
265 with destination.open("wb") as f:
266 if att.aio_stream:
267 async for chunk in att.aio_stream:
268 f.write(chunk)
269 elif att.stream:
270 shutil.copyfileobj(att.stream, f)
271 elif att.url:
272 async with (
273 aiohttp.ClientSession() as http,
274 http.get(att.url) as resp,
275 ):
276 resp.raise_for_status()
277 async for chunk in resp.content.iter_chunked(64 * 1024):
278 f.write(chunk)
279 else:
280 raise RuntimeError
282 if config.NO_UPLOAD_FILE_READ_OTHERS:
283 log.debug("Changing perms of %s", destination)
284 destination.chmod(destination.stat().st_mode | stat.S_IROTH)
286 url = f"{file_id}/{uu}/{att.name}"
287 return destination, url
290class _AttachmentWithName(LegacyAttachment):
291 name: str
292 path: Path | None
295class _AttachmentWithMetadata(_AttachmentWithName):
296 size: int
297 content_type: str
300def _ensure_name(att: LegacyAttachment) -> _AttachmentWithName:
301 if not att.name:
302 if att.path:
303 att.name = Path(att.path).name
304 elif att.url:
305 att.name = att.url.split("/")[-1]
306 else:
307 att.name = "unnamed-file"
308 return cast(_AttachmentWithName, att)
311async def _ensure_metadata(att: _AttachmentWithName) -> _AttachmentWithMetadata:
312 if att.size is None:
313 if att.data:
314 att.size = len(att.data)
315 elif att.stream:
316 att.stream.seek(0, io.SEEK_END)
317 att.size = att.stream.tell()
318 att.stream.seek(0)
319 elif att.path:
320 assert isinstance(att.path, Path)
321 att.size = att.path.stat().st_size
322 elif att.url:
323 async with (
324 aiohttp.ClientSession() as http,
325 http.head(att.url) as resp,
326 ):
327 att.size = resp.content_length
328 elif att.aio_stream:
329 warnings.warn("A size should be passed with async iterators")
330 tmp_dir = Path(tempfile.mkdtemp(prefix=_TEMP_PREFIX))
331 with (tmp_dir / att.name).open("wb") as fp:
332 async for chunk in att.aio_stream:
333 fp.write(chunk)
334 att.path = Path(fp.name)
335 att.size = att.path.stat().st_size
337 if not att.content_type:
338 att.content_type, _encoding = guess_type(att.name)
339 if not att.content_type:
340 att.content_type = "application/octet-stream"
342 return cast(_AttachmentWithMetadata, att)
345@contextlib.asynccontextmanager
346async def _get_data(
347 att: LegacyAttachment, http: aiohttp.ClientSession
348) -> AsyncIterator[bytes | IO[bytes] | AsyncIterator[bytes]]:
349 if (data := att.data or att.stream) is not None:
350 yield data
351 elif att.path is not None:
352 assert isinstance(att.path, Path)
353 with att.path.open("rb") as fp:
354 yield fp
355 elif att.aio_stream is not None:
356 # The aiostream may already have been consumed if a size wasn't passed.
357 # But in this case the `path`attribute is not Nonce, cf above.
358 yield att.aio_stream
359 elif att.url is not None:
360 async with http.get(att.url) as resp_get:
361 resp_get.raise_for_status()
362 yield resp_get.content.iter_any()
363 else:
364 raise RuntimeError("NEVER")
367def is_temp_path(path: Path, async_iterator_download_only: bool = False) -> bool:
368 try:
369 rel = path.relative_to(_TEMP_ROOT)
370 except ValueError:
371 return False
372 if async_iterator_download_only:
373 return rel.parts[0].startswith(_TEMP_PREFIX)
374 else:
375 return True
378_TEMP_ROOT = Path(tempfile.gettempdir())
379_TEMP_PREFIX = "slidge-async-iterator-download"
381log = logging.getLogger(__name__)