Xtao-Labs/MTPyroger
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Always run crypto-related functions in the dedicated thread

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pyrogram/__init__.py
View File

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

8
pyrogram/connection/transport/tcp/tcp_abridged_o.py
View File

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

14
pyrogram/session/session.py
View File

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

8
pyrogram/utils.py
View File

@ -309,11 +309,3 @@ async def parse_text_entities(
"message": text, "message": text,
"entities": entities "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)
)