Zergling_man/moment
Zergling_man/moment
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Reject files too large before attempting to upload

Browse Source
This commit is contained in:
miruka 2020-03-09 10:23:08 -04:00
parent 382235875c
commit 048c1ddabf
1 changed files with 16 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
src/backend/matrix_client.py
View File

@ -32,7 +32,8 @@ from nio.crypto import async_generator_from_data
from . import __app_name__, __display_name__, utils from . import __app_name__, __display_name__, utils
from .errors import ( from .errors import (
BadMimeType, InvalidUserId, InvalidUserInContext, MatrixError, BadMimeType, InvalidUserId, InvalidUserInContext, MatrixError,
MatrixNotFound, UneededThumbnail, UserFromOtherServerDisallowed, MatrixNotFound, MatrixTooLarge, UneededThumbnail,
UserFromOtherServerDisallowed,
) )
from .html_markdown import HTML_PROCESSOR as HTML from .html_markdown import HTML_PROCESSOR as HTML
from .media_cache import Media, Thumbnail from .media_cache import Media, Thumbnail
@ -384,6 +385,7 @@ class MatrixClient(nio.AsyncClient):
url, mime, crypt_dict = await self.upload( url, mime, crypt_dict = await self.upload(
lambda *_: path, lambda *_: path,
filename = path.name, filename = path.name,
filesize = size,
encrypt = encrypt, encrypt = encrypt,
monitor = monitor, monitor = monitor,
) )
@ -466,6 +468,7 @@ class MatrixClient(nio.AsyncClient):
lambda *_: thumb_data, lambda *_: thumb_data,
filename = filename =
f"{path.stem}_sample{''.join(path.suffixes)}", f"{path.stem}_sample{''.join(path.suffixes)}",
filesize = thumb_info.size,
encrypt = encrypt, encrypt = encrypt,
monitor = monitor, monitor = monitor,
) )
@ -886,13 +889,17 @@ class MatrixClient(nio.AsyncClient):
async def upload( async def upload(
self, self,
data_provider: nio.DataProvider, data_provider: nio.DataProvider,
mime: Optional[str] = None,
filename: Optional[str] = None, filename: Optional[str] = None,
filesize: Optional[int] = None,
mime: Optional[str] = None,
encrypt: bool = False, encrypt: bool = False,
monitor: Optional[nio.TransferMonitor] = None, monitor: Optional[nio.TransferMonitor] = None,
) -> UploadReturn: ) -> UploadReturn:
"""Upload a file to the matrix homeserver.""" """Upload a file to the matrix homeserver."""
if filesize > self.models["accounts"][self.user_id].max_upload_size:
raise MatrixTooLarge()
mime = mime or await utils.guess_mime(data_provider(0, 0)) mime = mime or await utils.guess_mime(data_provider(0, 0))
response, decryption_dict = await super().upload( response, decryption_dict = await super().upload(
@ -909,12 +916,18 @@ class MatrixClient(nio.AsyncClient):
async def set_avatar_from_file(self, path: Union[Path, str]) -> None: async def set_avatar_from_file(self, path: Union[Path, str]) -> None:
"""Upload an image to the homeserver and set it as our avatar.""" """Upload an image to the homeserver and set it as our avatar."""
path = Path(path)
mime = await utils.guess_mime(path) mime = await utils.guess_mime(path)
if mime.split("/")[0] != "image": if mime.split("/")[0] != "image":
raise BadMimeType(wanted="image/*", got=mime) raise BadMimeType(wanted="image/*", got=mime)
mxc, *_ = await self.upload(lambda *_: path, mime, Path(path).name) mxc, *_ = await self.upload(
data_provider = lambda *_: path,
filename = path.name,
filesize = path.resolve().stat().st_size,
mime = mime,
)
await self.set_avatar(mxc) await self.set_avatar(mxc)