Skip to content

Commit 2104731

Browse files
HitaloMHitaloM
committed
fix(client): better logic to get and reuse sessions in get_file 
1 parent f1f817f commit 2104731

2 files changed

Lines changed: 110 additions & 97 deletions

File tree

‎hydrogram/client.py‎

Lines changed: 110 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
from hydrogram import __license__, __version__, enums, raw, utils
4545
from hydrogram.crypto import aes
4646
from hydrogram.errors import (
47+
AuthBytesInvalid,
4748
BadRequest,
4849
CDNFileHashMismatch,
4950
ChannelPrivate,
@@ -52,7 +53,6 @@
5253
)
5354
from hydrogram.handlers.handler import Handler
5455
from hydrogram.methods import Methods
55-
from hydrogram.methods.messages.inline_session import get_session
5656
from hydrogram.session import Auth, Session
5757
from hydrogram.storage import BaseStorage, SQLiteStorage
5858
from hydrogram.types import ListenerTypes, TermsOfService, User
@@ -900,10 +900,9 @@ async def get_file(
900900
offset: int = 0,
901901
progress: Callable | None = None,
902902
progress_args: tuple = (),
903-
) -> AsyncGenerator[bytes, None] | None:
903+
) -> AsyncGenerator[bytes, None]:
904904
async with self.get_file_semaphore:
905905
file_type = file_id.file_type
906-
907906
if file_type == FileType.CHAT_PHOTO:
908907
if file_id.chat_id > 0:
909908
peer = raw.types.InputPeerUser(
@@ -916,7 +915,6 @@ async def get_file(
916915
channel_id=utils.get_channel_id(file_id.chat_id),
917916
access_hash=file_id.chat_access_hash,
918917
)
919-
920918
location = raw.types.InputPeerPhotoFileLocation(
921919
peer=peer,
922920
photo_id=file_id.media_id,
@@ -941,22 +939,48 @@ async def get_file(
941939
total = abs(limit) or (1 << 31) - 1
942940
chunk_size = 1024 * 1024
943941
offset_bytes = abs(offset) * chunk_size
944-
945942
dc_id = file_id.dc_id
946943

947944
try:
948-
session = await get_session(self, dc_id)
949-
r = await session.invoke(
950-
raw.functions.upload.GetFile(
951-
location=location, offset=offset_bytes, limit=chunk_size
952-
),
953-
sleep_threshold=30,
954-
)
945+
session = self.media_sessions.get(dc_id)
946+
if not session:
947+
auth_key = (
948+
await Auth(self, dc_id, await self.storage.test_mode()).create()
949+
if dc_id != await self.storage.dc_id()
950+
else await self.storage.auth_key()
951+
)
952+
session = self.media_sessions[dc_id] = Session(
953+
self, dc_id, auth_key, await self.storage.test_mode(), is_media=True
954+
)
955+
await session.start()
955956

956-
if isinstance(r, raw.types.upload.File):
957-
while True:
958-
chunk = r.bytes
957+
if dc_id != await self.storage.dc_id():
958+
for _ in range(3):
959+
exported_auth = await self.invoke(
960+
raw.functions.auth.ExportAuthorization(dc_id=dc_id)
961+
)
962+
try:
963+
await session.invoke(
964+
raw.functions.auth.ImportAuthorization(
965+
id=exported_auth.id, bytes=exported_auth.bytes
966+
)
967+
)
968+
break
969+
except AuthBytesInvalid:
970+
continue
971+
else:
972+
raise AuthBytesInvalid
973+
974+
while True:
975+
r = await session.invoke(
976+
raw.functions.upload.GetFile(
977+
location=location, offset=offset_bytes, limit=chunk_size
978+
),
979+
sleep_threshold=30,
980+
)
959981

982+
if isinstance(r, raw.types.upload.File):
983+
chunk = r.bytes
960984
yield chunk
961985

962986
current += 1
@@ -978,107 +1002,97 @@ async def get_file(
9781002
if len(chunk) < chunk_size or current >= total:
9791003
break
9801004

981-
r = await session.invoke(
982-
raw.functions.upload.GetFile(
983-
location=location, offset=offset_bytes, limit=chunk_size
984-
),
985-
sleep_threshold=30,
1005+
elif isinstance(r, raw.types.upload.FileCdnRedirect):
1006+
cdn_session = Session(
1007+
self,
1008+
r.dc_id,
1009+
await Auth(self, r.dc_id, await self.storage.test_mode()).create(),
1010+
await self.storage.test_mode(),
1011+
is_media=True,
1012+
is_cdn=True,
9861013
)
9871014

988-
elif isinstance(r, raw.types.upload.FileCdnRedirect):
989-
cdn_session = Session(
990-
self,
991-
r.dc_id,
992-
await Auth(self, r.dc_id, await self.storage.test_mode()).create(),
993-
await self.storage.test_mode(),
994-
is_media=True,
995-
is_cdn=True,
996-
)
1015+
try:
1016+
await cdn_session.start()
9971017

998-
try:
999-
await cdn_session.start()
1000-
1001-
while True:
1002-
r2 = await cdn_session.invoke(
1003-
raw.functions.upload.GetCdnFile(
1004-
file_token=r.file_token,
1005-
offset=offset_bytes,
1006-
limit=chunk_size,
1018+
while True:
1019+
r2 = await cdn_session.invoke(
1020+
raw.functions.upload.GetCdnFile(
1021+
file_token=r.file_token,
1022+
offset=offset_bytes,
1023+
limit=chunk_size,
1024+
)
10071025
)
1008-
)
10091026

1010-
if isinstance(r2, raw.types.upload.CdnFileReuploadNeeded):
1011-
try:
1012-
await session.invoke(
1013-
raw.functions.upload.ReuploadCdnFile(
1014-
file_token=r.file_token,
1015-
request_token=r2.request_token,
1027+
if isinstance(r2, raw.types.upload.CdnFileReuploadNeeded):
1028+
try:
1029+
await session.invoke(
1030+
raw.functions.upload.ReuploadCdnFile(
1031+
file_token=r.file_token,
1032+
request_token=r2.request_token,
1033+
)
10161034
)
1017-
)
1018-
except VolumeLocNotFound:
1019-
break
1020-
else:
1021-
continue
1022-
1023-
chunk = r2.bytes
1024-
1025-
# https://core.telegram.org/cdn#decrypting-files
1026-
decrypted_chunk = aes.ctr256_decrypt(
1027-
chunk,
1028-
r.encryption_key,
1029-
bytearray(
1030-
r.encryption_iv[:-4] + (offset_bytes // 16).to_bytes(4, "big")
1031-
),
1032-
)
1033-
1034-
hashes = await session.invoke(
1035-
raw.functions.upload.GetCdnFileHashes(
1036-
file_token=r.file_token, offset=offset_bytes
1035+
except VolumeLocNotFound:
1036+
break
1037+
else:
1038+
continue
1039+
1040+
chunk = r2.bytes
1041+
1042+
# https://core.telegram.org/cdn#decrypting-files
1043+
decrypted_chunk = aes.ctr256_decrypt(
1044+
chunk,
1045+
r.encryption_key,
1046+
bytearray(
1047+
r.encryption_iv[:-4]
1048+
+ (offset_bytes // 16).to_bytes(4, "big")
1049+
),
10371050
)
1038-
)
10391051

1040-
# https://core.telegram.org/cdn#verifying-files
1041-
for i, h in enumerate(hashes):
1042-
cdn_chunk = decrypted_chunk[h.limit * i : h.limit * (i + 1)]
1043-
CDNFileHashMismatch.check(
1044-
h.hash == sha256(cdn_chunk).digest(),
1045-
"h.hash == sha256(cdn_chunk).digest()",
1052+
hashes = await session.invoke(
1053+
raw.functions.upload.GetCdnFileHashes(
1054+
file_token=r.file_token, offset=offset_bytes
1055+
)
10461056
)
10471057

1048-
yield decrypted_chunk
1058+
# https://core.telegram.org/cdn#verifying-files
1059+
for i, h in enumerate(hashes):
1060+
cdn_chunk = decrypted_chunk[h.limit * i : h.limit * (i + 1)]
1061+
CDNFileHashMismatch.check(
1062+
h.hash == sha256(cdn_chunk).digest(),
1063+
"h.hash == sha256(cdn_chunk).digest()",
1064+
)
10491065

1050-
current += 1
1051-
offset_bytes += chunk_size
1066+
yield decrypted_chunk
10521067

1053-
if progress:
1054-
func = functools.partial(
1055-
progress,
1056-
min(offset_bytes, file_size)
1057-
if file_size != 0
1058-
else offset_bytes,
1059-
file_size,
1060-
*progress_args,
1061-
)
1068+
current += 1
1069+
offset_bytes += chunk_size
10621070

1063-
if inspect.iscoroutinefunction(progress):
1064-
await func()
1065-
else:
1066-
await self.loop.run_in_executor(self.executor, func)
1071+
if progress:
1072+
func = functools.partial(
1073+
progress,
1074+
min(offset_bytes, file_size)
1075+
if file_size != 0
1076+
else offset_bytes,
1077+
file_size,
1078+
*progress_args,
1079+
)
10671080

1068-
if len(chunk) < chunk_size or current >= total:
1069-
break
1070-
except Exception as e:
1071-
raise e
1072-
finally:
1073-
await cdn_session.stop()
1081+
if inspect.iscoroutinefunction(progress):
1082+
await func()
1083+
else:
1084+
await self.loop.run_in_executor(self.executor, func)
1085+
1086+
if len(chunk) < chunk_size or current >= total:
1087+
break
1088+
finally:
1089+
await cdn_session.stop()
10741090
except hydrogram.StopTransmission:
10751091
raise
10761092
except hydrogram.errors.FloodWait:
10771093
raise
10781094
except Exception as e:
10791095
log.exception(e)
1080-
finally:
1081-
await session.stop()
10821096

10831097
def guess_mime_type(self, filename: str) -> str | None:
10841098
return self.mimetypes.guess_type(filename)[0]

‎news/15.bugfix.rst‎

Lines changed: 0 additions & 1 deletion
This file was deleted.

0 commit comments

Comments
 (0)