fix: CDN client can't connect

This commit is contained in:
HBcao233 2024-08-08 01:45:51 +08:00
parent 317ee12ab8
commit 50f6816fd7
2 changed files with 30 additions and 30 deletions

View File

@ -39,7 +39,7 @@ class _DirectDownloadIter(RequestIter):
self.request = functions.upload.GetFileRequest( self.request = functions.upload.GetFileRequest(
file, offset=offset, limit=request_size) file, offset=offset, limit=request_size)
self._client = self.client self._client = self.client
self.cdn_redirect = cdn_redirect self._cdn_redirect = cdn_redirect
if cdn_redirect is not None: if cdn_redirect is not None:
self.request = functions.upload.GetCdnFileRequest(cdn_redirect.file_token, offset=offset, limit=request_size) self.request = functions.upload.GetCdnFileRequest(cdn_redirect.file_token, offset=offset, limit=request_size)
self._client = await self.client._get_cdn_client(cdn_redirect) self._client = await self.client._get_cdn_client(cdn_redirect)
@ -54,22 +54,22 @@ class _DirectDownloadIter(RequestIter):
self._exported = dc_id and self._client.session.dc_id != dc_id self._exported = dc_id and self._client.session.dc_id != dc_id
if not self._exported: if not self._exported:
# The used sender will also change if ``FileMigrateError`` occurs # The used sender will also change if ``FileMigrateError`` occurs
self._sender = self._client._sender self._sender = self.client._sender
else: else:
try: try:
self._sender = await self._client._borrow_exported_sender(dc_id) self._sender = await self.client._borrow_exported_sender(dc_id)
except errors.DcIdInvalidError: except errors.DcIdInvalidError:
# Can't export a sender for the ID we are currently in # Can't export a sender for the ID we are currently in
config = await self._client(functions.help.GetConfigRequest()) config = await self.client(functions.help.GetConfigRequest())
for option in config.dc_options: for option in config.dc_options:
if option.ip_address == self._client.session.server_address: if option.ip_address == self.client.session.server_address:
self._client.session.set_dc( self.client.session.set_dc(
option.id, option.ip_address, option.port) option.id, option.ip_address, option.port)
self._client.session.save() self.client.session.save()
break break
# TODO Figure out why the session may have the wrong DC ID # TODO Figure out why the session may have the wrong DC ID
self._sender = self._client._sender self._sender = self.client._sender
self._exported = False self._exported = False
async def _load_next_chunk(self): async def _load_next_chunk(self):
@ -84,12 +84,14 @@ class _DirectDownloadIter(RequestIter):
async def _request(self): async def _request(self):
try: try:
result = await self._client._call(self._sender, self.request) result = await self._client._call(self._sender, self.request)
self.client._log[__name__].debug('result: %s', result)
self._timed_out = False self._timed_out = False
if isinstance(result, types.upload.FileCdnRedirect): if isinstance(result, types.upload.FileCdnRedirect):
# raise NotImplementedError # TODO Implement if self.client._mb_entity_cache.self_bot:
raise ValueError('FileCdnRedirect but the GetCdnFileRequest API access for bot users is restricted. Try to change api_id to avoid FileCdnRedirect')
raise _CdnRedirect(result) raise _CdnRedirect(result)
if isinstance(result, types.upload.CdnFileReuploadNeeded): if isinstance(result, types.upload.CdnFileReuploadNeeded):
result = await self._client._call(self._sender, functions.upload.reuploadCdnFile(file_token=self.cdn_redirect.file_token, request_token=result.request_token)) await self.client._call(self.client._sender, functions.upload.ReuploadCdnFileRequest(file_token=self._cdn_redirect.file_token, request_token=result.request_token))
result = await self._client._call(self._sender, self.request) result = await self._client._call(self._sender, self.request)
return result.bytes return result.bytes
else: else:
@ -579,7 +581,7 @@ class DownloadMethods:
if in_memory: if in_memory:
return f.getvalue() return f.getvalue()
except _CdnRedirect as e: except _CdnRedirect as e:
self._log[__name__].info('FileCdnRedirect') self._log[__name__].info('FileCdnRedirect to CDN data center %s', e.cdn_redirect.dc_id)
return await self._download_file( return await self._download_file(
input_location=input_location, input_location=input_location,
file=file, file=file,

View File

@ -800,10 +800,13 @@ class TelegramBaseClient(abc.ABC):
'Failed to get DC %s (cdn = %s) with use_ipv6 = %s; retrying ignoring IPv6 check', 'Failed to get DC %s (cdn = %s) with use_ipv6 = %s; retrying ignoring IPv6 check',
dc_id, cdn, self._use_ipv6 dc_id, cdn, self._use_ipv6
) )
try:
return next( return next(
dc for dc in cls._config.dc_options dc for dc in cls._config.dc_options
if dc.id == dc_id and bool(dc.cdn) == cdn if dc.id == dc_id and bool(dc.cdn) == cdn
) )
except StopIteration:
raise ValueError(f'Failed to get DC {dc_id} (cdn = {cdn})')
async def _create_exported_sender(self: 'TelegramClient', dc_id): async def _create_exported_sender(self: 'TelegramClient', dc_id):
""" """
@ -891,9 +894,6 @@ class TelegramBaseClient(abc.ABC):
async def _get_cdn_client(self: 'TelegramClient', cdn_redirect): async def _get_cdn_client(self: 'TelegramClient', cdn_redirect):
"""Similar to ._borrow_exported_client, but for CDNs""" """Similar to ._borrow_exported_client, but for CDNs"""
# TODO Implement
# raise NotImplementedError
from .telegramclient import TelegramClient
session = self._exported_sessions.get(cdn_redirect.dc_id) session = self._exported_sessions.get(cdn_redirect.dc_id)
if not session: if not session:
dc = await self._get_dc(cdn_redirect.dc_id, cdn=True) dc = await self._get_dc(cdn_redirect.dc_id, cdn=True)
@ -902,24 +902,22 @@ class TelegramBaseClient(abc.ABC):
self._exported_sessions[cdn_redirect.dc_id] = session self._exported_sessions[cdn_redirect.dc_id] = session
self._log[__name__].info('Creating new CDN client') self._log[__name__].info('Creating new CDN client')
client = TelegramClient( client = self.__class__(
session, self.api_id, self.api_hash, session, self.api_id, self.api_hash,
proxy=self._proxy, proxy=self._proxy,
timeout=self._timeout, timeout=self._timeout,
loop=self.loop loop=self.loop
) )
# This will make use of the new RSA keys for this specific CDN. session.auth_key = self._sender.auth_key
# await client._sender.connect(self._connection(
# We won't be calling GetConfigRequest because it's only called session.server_address,
# when needed by ._get_dc, and also it's static so it's likely session.port,
# set already. Avoid invoking non-CDN methods by not syncing updates. session.dc_id,
loggers=self._log,
self_id = self._mb_entity_cache.self_id proxy=self._proxy,
self_user = self.session.get_input_entity(self_id) local_addr=self._local_addr
client._mb_entity_cache.set_self_user(self_id, True, self_user.access_hash) ))
await client.start()
return client return client
# endregion # endregion