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