Always run crypto-related functions in the dedicated thread

This commit is contained in:
Dan 2021-01-02 18:45:43 +01:00
parent 7dda167c09
commit a48d27f501
4 changed files with 13 additions and 19 deletions

View File

@ -46,4 +46,4 @@ main_event_loop = asyncio.get_event_loop()
CRYPTO_EXECUTOR_SIZE_THRESHOLD = 512
crypto_executor = ThreadPoolExecutor(2, thread_name_prefix="CryptoWorker")
crypto_executor = ThreadPoolExecutor(1, thread_name_prefix="CryptoWorker")

View File

@ -18,11 +18,9 @@
import logging
import os
from typing import Optional
from pyrogram import utils
import pyrogram
from pyrogram.crypto import aes
from .tcp import TCP
@ -60,7 +58,7 @@ class TCPAbridgedO(TCP):
async def send(self, data: bytes, *args):
length = len(data) // 4
data = (bytes([length]) if length <= 126 else b"\x7f" + length.to_bytes(3, "little")) + data
payload = await utils.maybe_run_in_executor(aes.ctr256_encrypt, data, len(data), self.loop, *self.encrypt)
payload = await self.loop.run_in_executor(pyrogram.crypto_executor, aes.ctr256_encrypt, data, *self.encrypt)
await super().send(payload)
@ -85,4 +83,4 @@ class TCPAbridgedO(TCP):
if data is None:
return None
return await utils.maybe_run_in_executor(aes.ctr256_decrypt, data, len(data), self.loop, *self.decrypt)
return await self.loop.run_in_executor(pyrogram.crypto_executor, aes.ctr256_decrypt, data, *self.decrypt)

View File

@ -25,7 +25,7 @@ from hashlib import sha1
from io import BytesIO
import pyrogram
from pyrogram import __copyright__, __license__, __version__, utils
from pyrogram import __copyright__, __license__, __version__
from pyrogram import raw
from pyrogram.connection import Connection
from pyrogram.crypto import mtproto
@ -217,8 +217,10 @@ class Session:
await self.start()
async def handle_packet(self, packet):
data = await utils.maybe_run_in_executor(
mtproto.unpack, BytesIO(packet), len(packet), self.loop,
data = await self.loop.run_in_executor(
pyrogram.crypto_executor,
mtproto.unpack,
BytesIO(packet),
self.session_id,
self.auth_key,
self.auth_key_id
@ -360,8 +362,10 @@ class Session:
log.debug(f"Sent:")
log.debug(message)
payload = await utils.maybe_run_in_executor(
mtproto.pack, message, len(message), self.loop,
payload = await self.loop.run_in_executor(
pyrogram.crypto_executor,
mtproto.pack,
message,
self.current_salt.salt,
self.session_id,
self.auth_key,

View File

@ -309,11 +309,3 @@ async def parse_text_entities(
"message": text,
"entities": entities
}
async def maybe_run_in_executor(func, data, length, loop, *args):
return (
func(data, *args)
if length <= pyrogram.CRYPTO_EXECUTOR_SIZE_THRESHOLD
else await loop.run_in_executor(pyrogram.crypto_executor, func, data, *args)
)