Skip to content

Commit

Permalink
Fix bug in '.push_file()'
Browse files Browse the repository at this point in the history
This commit fix very rare bug with incorrect MD5
calculation on a Telegram side for small files.

Now OpenPretender() class have '.get_expected_size()'
method, so we can correctly report file size on
upload to the underlying Telethon library.
  • Loading branch information
NotStatilko committed Jan 7, 2025
1 parent e8a77bd commit 6f58c3c
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 6 deletions.
4 changes: 2 additions & 2 deletions tgbox/api/remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -1061,7 +1061,7 @@ async def _push_file(
ifile = await upload_file(
self._tc, oe,
file_name=urlsafe_b64encode(pf.filesalt.salt).decode(),
part_size_kb=512, file_size=pf.filesize,
part_size_kb=512, file_size=oe.get_expected_size(),
progress_callback=progress_callback
)
except Exception as e:
Expand All @@ -1074,7 +1074,7 @@ async def _push_file(

ifile = await self._tc.upload_file(
oe, file_name=urlsafe_b64encode(pf.filesalt.salt).decode(),
part_size_kb=512, file_size=pf.filesize,
part_size_kb=512, file_size=oe.get_expected_size(),
progress_callback=progress_callback)
try:
if message_to_edit:
Expand Down
23 changes: 19 additions & 4 deletions tgbox/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -361,7 +361,7 @@ def __init__(
self, flo: BinaryIO,
aes_state: 'tgbox.crypto.AESwState',
hmac_state: 'hashlib.HMAC',
file_size: Optional[int] = None,
file_size: int
):
"""
Arguments:
Expand All @@ -374,9 +374,8 @@ def __init__(
hmac_state (``hmac.HMAC``):
``HMAC`` initialized with ``HMACKey``
file_size (``int``, optional):
File size of ``flo``. If not specified,
we will try to seek.
file_size (``int``):
File size of ``flo``.
"""
self._aes_state = aes_state
self._hmac_state = hmac_state
Expand Down Expand Up @@ -412,6 +411,22 @@ def concat_metadata(self, metadata: bytes) -> None:
self._buffered_bytes += metadata
self._concated_metadata_size = len(metadata)

def get_expected_size(self):
"""
Returns expected actual Telegram document size after
upload. We use it in ``push_file()``
"""
if not self._concated_metadata_size:
raise Exception('You need to concat metadata firstly')
# self._file_size already include size of Metadata, but
# we need to calculate size of encrypted File *with*
# padding, so firstly we are required to subtract
# self._concated_metadata_size from self._file_size
# for correct calculation. 32 here is HMAC blob
expected = self._file_size - self._concated_metadata_size
expected = (expected + (16 - expected % 16)) + 32
return expected + self._concated_metadata_size

async def read(self, size: int=-1) -> bytes:
"""
Returns ``size`` bytes from async Generator.
Expand Down

0 comments on commit 6f58c3c

Please sign in to comment.