Coverage for slidge/core/attachment_upload.py: 84%
288 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-10 04:45 +0000
1"""Upload attachments to an HTTP upload service (:xep:`0363`)."""
3from __future__ import annotations
5import contextlib
6import hashlib
7import io
8import logging
9import os
10import shutil
11import stat
12import tempfile
13import warnings
14from collections.abc import AsyncIterator
15from mimetypes import guess_type
16from pathlib import Path
17from typing import IO, TYPE_CHECKING, Literal, cast
18from urllib.parse import quote as urlquote
19from uuid import uuid4
21import aiohttp
22from PIL.Image import Image
23from slixmpp import JID, Iq
25from ..db.models import Attachment
26from ..util.types import AvatarMetadata, LegacyAttachment
27from ..util.util import fix_namespaces, fix_suffix
28from . import config
30if TYPE_CHECKING:
31 from ..db.avatar import AvatarType
32 from .gateway import BaseGateway
33 from .session import BaseSession
36class AttachmentUploader:
37 """Turns a :class:`.LegacyAttachment` into a URL that XMPP clients can fetch."""
39 xmpp: BaseGateway
40 session: BaseSession | None
41 """
42 The session uplaoding attachments, or :const:`None` for the gateway component,
43 which is not bound to a :term:`User`.
44 """
46 def __init__(self, xmpp: BaseGateway, session: BaseSession | None) -> None:
47 self.xmpp = xmpp
48 self.session = session
50 @contextlib.asynccontextmanager
51 async def dedup_lock(
52 self,
53 attachment: LegacyAttachment | Path | str,
54 ) -> AsyncIterator[None]:
55 """Take a lock on an attachment with a given name.
57 Prevents races which download the same attachment several times."""
58 session = self.session
59 if (
60 session is None
61 or not isinstance(attachment, LegacyAttachment)
62 or attachment.legacy_file_id is None
63 ):
64 yield
65 else:
66 async with session.lock(("attachment", attachment.legacy_file_id)):
67 yield
69 async def get_stored(self, attachment: LegacyAttachment) -> Attachment:
70 """
71 Fetch the :class:`.Attachment` already uploaded for this user, if any,
72 or a new (transient) one.
73 """
74 session = self.session
75 if attachment.legacy_file_id is not None and session is not None:
76 with self.xmpp.store.session() as orm:
77 stored = (
78 orm.query(Attachment)
79 .filter_by(
80 legacy_file_id=str(attachment.legacy_file_id),
81 user_account_id=session.user_pk,
82 )
83 .one_or_none()
84 )
85 if stored is not None:
86 if not await self.__valid_url(session, stored.url):
87 stored.url = None # type:ignore
88 return stored
89 return Attachment(
90 user_account_id=None if session is None else session.user_pk,
91 legacy_file_id=None
92 if attachment.legacy_file_id is None
93 else str(attachment.legacy_file_id),
94 url=attachment.url if config.USE_ATTACHMENT_ORIGINAL_URLS else None,
95 )
97 def record(self, stored: Attachment) -> None:
98 """Remember an uploaded attachment, so that it is not uploaded twice.
100 No-op without a session, since attachments are stored per user.
101 """
102 # TODO: we need a separate mechanism to record gateway component attachments.
103 if self.session is None:
104 return
105 with self.xmpp.store.session(expire_on_commit=False) as orm:
106 orm.add(stored)
107 orm.commit()
109 @staticmethod
110 async def __valid_url(session: BaseSession, url: str) -> bool:
111 async with session.http.head(url) as r:
112 return r.status < 400
114 async def get_url(self, att: LegacyAttachment, stored: Attachment) -> str:
115 att = _ensure_name(att)
117 if len(att.name) > config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH:
118 log.debug("Trimming long filename: %s", att.name)
119 base, ext = os.path.splitext(att.name)
120 att.name = (
121 base[: config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH - len(ext)] + ext
122 )
124 if config.FIX_FILENAME_SUFFIX_MIME_TYPE and isinstance(att.path, Path):
125 att.name, att.content_type = fix_suffix(
126 att.path, att.content_type, att.name
127 )
129 att.legacy_file_id = stored.legacy_file_id
131 if config.NO_UPLOAD_PATH:
132 att.path, new_url = await self.__no_upload(att, stored.legacy_file_id)
133 new_url = (
134 (config.NO_UPLOAD_URL_PREFIX or "") + "/message/" + urlquote(new_url)
135 )
136 else:
137 new_url = await self.__upload(att)
139 if stored.legacy_file_id:
140 stored.url = new_url
142 return new_url
144 async def __upload(
145 self,
146 att: _AttachmentWithName,
147 purpose: Literal["message", "profile"] = "message",
148 hasher: hashlib._Hash | None = None,
149 ) -> str:
150 assert config.UPLOAD_SERVICE
151 att = await _ensure_metadata(att)
152 iq_slot = await self.__request_upload_slot(
153 config.UPLOAD_SERVICE,
154 att.name,
155 att.size,
156 att.content_type,
157 purpose=purpose,
158 )
159 if iq_slot["type"] == "error":
160 # COMPAT: In theory, __request_upload_slot() raises IqError, but
161 # there is a bug in prosody's mod_privilege where the outer
162 # IQ type is (illegally) not set to error when the inner IQ
163 # has type='error'.
164 raise RuntimeError(f"Error while requesting upload slot: {iq_slot}")
165 slot = iq_slot.get_plugin("http_upload_slot", check=True)
166 if slot is None:
167 raise RuntimeError(f"No upload slot in this IQ: {iq_slot}")
168 put = slot["put"]["url"]
169 assert isinstance(put, str)
170 if not put:
171 raise RuntimeError(f"Cannot find a PUT URL in: {slot}")
172 get = slot["get"]["url"]
173 assert isinstance(get, str)
174 if not get:
175 raise RuntimeError(f"Cannot find a GET URL in: {slot}")
176 headers = {
177 "Content-Length": str(att.size),
178 "Content-Type": att.content_type,
179 **{header["name"]: header["value"] for header in slot["put"]["headers"]},
180 }
182 async with (
183 aiohttp.ClientSession() as http,
184 _get_data(att, http, hasher) as data,
185 http.put(slot["put"]["url"], data=data, headers=headers) as resp,
186 ):
187 resp.raise_for_status()
189 return get
191 async def __request_upload_slot(
192 self,
193 upload_service: JID | str,
194 filename: str,
195 size: int,
196 content_type: str,
197 *,
198 purpose: Literal["message", "profile"] = "message",
199 ) -> Iq:
200 iq_request = self.xmpp.make_iq_get(ito=upload_service)
201 request = iq_request["http_upload_request"]
202 request["filename"] = filename
203 request["size"] = str(size)
204 request["content-type"] = content_type
205 if purpose != "message":
206 request.enable(purpose)
207 session = self.session
208 if session is not None:
209 iq_request.set_from(session.user_jid)
210 try:
211 return await self.xmpp.plugin["xep_0356"].send_privileged_iq(iq_request)
212 except Exception as e: # noqa: BLE001
213 warnings.warn(
214 "Could not request upload slot on behalf of "
215 f"{session.user_jid}: {e}."
216 "Falling back to not using privileges."
217 )
218 fix_namespaces(iq_request.xml, "jabber:client", "jabber:component:accept")
219 iq_request.set_from(config.UPLOAD_REQUESTER or self.xmpp.boundjid)
220 return await iq_request.send() # type:ignore[no-any-return]
222 @staticmethod
223 async def __no_upload(
224 att: _AttachmentWithName, legacy_file_id: str | None
225 ) -> tuple[Path, str]:
226 file_id = uuid4().hex if legacy_file_id is None else legacy_file_id
227 assert config.NO_UPLOAD_PATH is not None
228 assert config.NO_UPLOAD_URL_PREFIX is not None
229 destination_dir = Path(config.NO_UPLOAD_PATH) / "message" / file_id
231 if destination_dir.exists():
232 log.debug("Dest dir exists: %s", destination_dir)
233 files = [f for f in destination_dir.glob("**/*") if f.is_file()]
234 if len(files) == 1:
235 log.debug(
236 "Found the legacy attachment '%s' at '%s'",
237 legacy_file_id,
238 files[0],
239 )
240 name = files[0].name
241 uu = files[0].parent.name # anti-obvious url trick, see below
242 return files[0], f"{file_id}/{uu}/{name}"
243 else:
244 log.warning(
245 (
246 "There are several or zero files in %s, "
247 "slidge doesn't know which one to pick among %s. "
248 "Removing the dir."
249 ),
250 destination_dir,
251 files,
252 )
253 shutil.rmtree(destination_dir)
255 log.debug("Did not find a file in: %s", destination_dir)
256 # let's use a UUID to avoid URLs being too obvious
257 uu = str(uuid4())
258 destination_dir = destination_dir / uu
259 destination_dir.mkdir(parents=True)
261 assert att.name
262 destination = destination_dir / att.name
263 if att.path:
264 assert isinstance(att.path, Path)
265 try:
266 destination.hardlink_to(att.path)
267 except OSError as e:
268 if is_temp_path(att.path):
269 shutil.copy2(att.path, destination)
270 else:
271 log.debug("Could not hardlink: %s, attempting symlink", e)
272 try:
273 destination.symlink_to(att.path)
274 except OSError as e:
275 log.debug("Could not symlink: %s, copying data", e)
276 shutil.copy2(att.path, destination)
277 elif att.data:
278 destination.write_bytes(att.data)
279 else:
280 with destination.open("wb") as f:
281 if att.aio_stream:
282 async for chunk in att.aio_stream:
283 f.write(chunk)
284 elif att.stream:
285 shutil.copyfileobj(att.stream, f)
286 elif att.url:
287 async with (
288 aiohttp.ClientSession() as http,
289 http.get(att.url) as resp,
290 ):
291 resp.raise_for_status()
292 async for chunk in resp.content.iter_chunked(64 * 1024):
293 f.write(chunk)
294 else:
295 raise RuntimeError
297 if config.NO_UPLOAD_FILE_READ_OTHERS:
298 log.debug("Changing perms of %s", destination)
299 destination.chmod(destination.stat().st_mode | stat.S_IROTH)
301 url = f"{file_id}/{uu}/{att.name}"
302 return destination, url
304 async def upload_avatar(
305 self, avatar: AvatarType, img: Image, hash_: str
306 ) -> AvatarMetadata | None:
307 if config.NO_UPLOAD_PATH:
308 return await self.__no_upload_avatar(avatar, img, hash_)
309 else:
310 return await self.__upload_avatar(avatar, img, hash_)
312 async def __no_upload_avatar(
313 self, avatar: AvatarType, img: Image, hash_: str
314 ) -> AvatarMetadata | None:
315 assert config.NO_UPLOAD_PATH
316 assert img.format
317 format = img.format.lower()
318 dest = (Path(config.NO_UPLOAD_PATH) / "profile" / hash_ / "avatar").with_suffix(
319 "." + format
320 )
321 url = f"{config.NO_UPLOAD_URL_PREFIX}/profile/{hash_}/{dest.name}"
322 if not dest.exists():
323 dest.parent.mkdir(exist_ok=True, parents=True)
324 if avatar.path:
325 shutil.copy2(avatar.path, dest)
326 elif avatar.data:
327 dest.write_bytes(avatar.data)
328 else:
329 img.save(dest)
330 return AvatarMetadata(
331 url=url,
332 width=img.width,
333 height=img.height,
334 bytes=dest.stat().st_size,
335 id=hashlib.sha1(dest.read_bytes()).hexdigest(),
336 type=format,
337 )
339 async def __upload_avatar(
340 self, avatar: AvatarType, img: Image, hash_: str
341 ) -> AvatarMetadata | None:
342 assert config.UPLOAD_SERVICE is not None
343 iq_or_info = await self.xmpp.plugin["xep_0030"].get_info(
344 JID(config.UPLOAD_SERVICE), cached=True
345 )
346 if isinstance(iq_or_info, Iq):
347 features = iq_or_info["disco_info"].get_features()
348 else:
349 features = iq_or_info.get_features()
350 if (
351 self.xmpp.plugin["xep_0363"].stanza.ProfilePurpose.namespace + "#profile"
352 not in features
353 ):
354 warnings.warn(
355 f"The upload service {config.UPLOAD_SERVICE} does not support "
356 "the 'profile' purpose, avatar data can only be served in-band.",
357 UserWarning,
358 )
359 return None
361 att = await _ensure_metadata(
362 _AttachmentWithName(
363 data=avatar.data,
364 url=avatar.url,
365 path=avatar.path,
366 name=avatar.path.name if avatar.path else "avatar",
367 )
368 )
369 hasher = hashlib.sha1()
370 try:
371 url = await self.__upload(att, purpose="profile", hasher=hasher)
372 except Exception:
373 log.exception("Could not upload avatar")
374 return None
375 return AvatarMetadata(
376 url=url,
377 width=img.width,
378 height=img.height,
379 bytes=att.size,
380 id=hasher.hexdigest(),
381 type=att.content_type.removeprefix("image/"),
382 )
385class _AttachmentWithName(LegacyAttachment):
386 name: str
387 path: Path | None
390class _AttachmentWithMetadata(_AttachmentWithName):
391 size: int
392 content_type: str
395def _ensure_name(att: LegacyAttachment) -> _AttachmentWithName:
396 if not att.name:
397 if att.path:
398 att.name = Path(att.path).name
399 elif att.url:
400 att.name = att.url.split("/")[-1]
401 else:
402 att.name = "unnamed-file"
403 return cast(_AttachmentWithName, att)
406async def _ensure_metadata(att: _AttachmentWithName) -> _AttachmentWithMetadata:
407 if att.size is None:
408 if att.data:
409 att.size = len(att.data)
410 elif att.stream:
411 att.stream.seek(0, io.SEEK_END)
412 att.size = att.stream.tell()
413 att.stream.seek(0)
414 elif att.path:
415 assert isinstance(att.path, Path)
416 att.size = att.path.stat().st_size
417 elif att.url:
418 async with (
419 aiohttp.ClientSession() as http,
420 http.head(att.url) as resp,
421 ):
422 att.size = resp.content_length
423 elif att.aio_stream:
424 warnings.warn("A size should be passed with async iterators")
425 tmp_dir = Path(tempfile.mkdtemp(prefix=_TEMP_PREFIX))
426 with (tmp_dir / att.name).open("wb") as fp:
427 async for chunk in att.aio_stream:
428 fp.write(chunk)
429 att.path = Path(fp.name)
430 att.size = att.path.stat().st_size
432 if not att.content_type:
433 att.content_type, _encoding = guess_type(att.name)
434 if not att.content_type:
435 att.content_type = "application/octet-stream"
437 return cast(_AttachmentWithMetadata, att)
440@contextlib.asynccontextmanager
441async def _get_data(
442 att: LegacyAttachment,
443 http: aiohttp.ClientSession,
444 hasher: hashlib._Hash | None = None,
445) -> AsyncIterator[bytes | IO[bytes] | AsyncIterator[bytes]]:
447 if att.data is not None:
448 if hasher is not None:
449 hasher.update(att.data)
450 yield att.data
451 elif att.stream is not None:
452 _hash(att.stream, hasher)
453 yield att.stream
454 elif att.path is not None:
455 assert isinstance(att.path, Path)
456 with att.path.open("rb") as fp:
457 _hash(fp, hasher)
458 yield fp
459 elif att.aio_stream is not None:
460 # The aiostream may already have been consumed if a size wasn't passed.
461 # But in this case the `path` attribute is not None, cf above.
462 yield _hash_and_yield_async(att.aio_stream)
463 elif att.url is not None:
464 async with http.get(att.url) as resp_get:
465 resp_get.raise_for_status()
466 yield _hash_and_yield_async(resp_get.content.iter_any())
467 else:
468 raise RuntimeError("NEVER")
471def _hash(fp: IO[bytes], hasher: hashlib._Hash | None = None) -> None:
472 if hasher is not None:
473 fp.seek(0)
474 while True:
475 chunk = fp.read(65536)
476 if not chunk:
477 break
478 hasher.update(chunk)
479 fp.seek(0)
482async def _hash_and_yield_async(
483 source: AsyncIterator[bytes], hasher: hashlib._Hash | None = None
484) -> AsyncIterator[bytes]:
485 async for chunk in source:
486 if hasher is not None:
487 hasher.update(chunk)
488 yield chunk
491def is_temp_path(path: Path, async_iterator_download_only: bool = False) -> bool:
492 try:
493 rel = path.relative_to(_TEMP_ROOT)
494 except ValueError:
495 return False
496 if async_iterator_download_only:
497 return rel.parts[0].startswith(_TEMP_PREFIX)
498 else:
499 return True
502_TEMP_ROOT = Path(tempfile.gettempdir())
503_TEMP_PREFIX = "slidge-async-iterator-download"
505log = logging.getLogger(__name__)