async def send(self, conn):
async with conn.lock_write():
speed: int = self.data
await conn.write(self.type_id + to_uint(speed, 4))
conn.debug(
f'[SEND {conn.address}] SET Download speed: {format_amount(speed)}'
)
async def _write_data_to_stream(self, conn, data, *_, compression):
offset = self.offset
delay = Delay(conn.download_speed)
async for chunk in data:
if offset > 0:
offset -= (i := min(offset, len(chunk)))
chunk = chunk[i:]
if not chunk:
continue
if not isinstance(chunk, Bytes):
raise MalformedDataError('Provided data chunk is not binary',
data=data,
headers=self.headers)
chunk, _ = await Compressor.compress(chunk, self.headers,
conn.allowed_compressors,
conn.default_compressor,
compression)
chunk_size = len(chunk)
if chunk_size >= 1 << 32:
raise MalformedDataError(
'Provided data chunk exceeded max chunk size',
data=data,
headers=self.headers)
await delay(chunk_size + 4)
await conn.write(to_uint(chunk_size, 4))
await conn.write(chunk)
conn.debug(
f'[SEND {conn.address}] [{int2hex(self.message_id):<4}] -> {bytes2hex(chunk[:64])}...'
)
await conn.write(b'\x00\x00\x00\x00')
async def send(self, conn):
self.conn = conn
data, compression = await self._encode_gen(conn)
header = self.type_id + self.Head(self.handler_id, self.message_id,
time_ns() // 1000_000,
self.data_type, compression).pack()
message_headers = self.headers.encode()
async with conn.lock_write():
await conn.write(header)
await conn.write(to_uint(len(message_headers), 4))
await conn.write(message_headers)
conn.debug(
f'[SEND {conn.address}] Stream '
f'H: {int2hex(self.handler_id):<4} '
f'M: {int2hex(self.message_id):<4} '
f'T: {Codec.codecs[self.data_type].type_name:<8} '
f'C: {Compressor.compressors[compression].type_name:<8} ')
conn.debug(
f'[SEND {conn.address}] [{int2hex(self.message_id):<4}] -> HEADERS {self.headers}'
)
await self._write_data_to_stream(conn,
data,
compression=compression)
async def init(self):
await self.write(to_uint(self.PROTOCOL_VERSION, 4))
result = await self.read(4)
if result != bytes(4):
raise ProtocolError(
f'Unsupported protocol version. '
f'Please upgrade your client to: {as_uint(result)}',
conn=self
)
client_stmt = ClientStatement(
api=self.api_version,
client_time=time_ns() // 1000_000,
scheme_format='JSON',
compressors=['gzip', 'zlib'],
default_compression='zlib',
)
self.set_compressors(['gzip', 'zlib'], 'zlib')
await self.write(client_stmt.pack())
stmt: ServerStatement = ServerStatement.unpack(
await self.read(as_uint(await self.read(4)))
)
self.time_delta = (stmt.server_time / 1000) - time()
if self.conf.handshake is not None:
await self.conf.handshake.send(self)
self.debug(f'{self} initialized')
async def compress_file(cls, src: Path, dst: Path,
headers: T_Headers) -> None:
compressor = zlib.compressobj(level=6)
value = 1
ln = src.stat().st_size
with src.open('rb') as rc:
with dst.open('wb') as wc:
wc.write(to_uint(ln, 4))
while line := rc.read(1 << 24):
value = zlib.adler32(line, value)
wc.write(compressor.compress(line))
wc.write(compressor.flush())
async def handle_stream(self, stream: IOStream, address: tuple[str, int]) -> None:
try:
protocol_version = as_uint(await stream.read_bytes(4))
if not self.protocols[0] <= protocol_version <= self.protocols[1]:
await stream.write(to_uint(self.protocols[1], 4))
stream.close(CatsError('Unsupported protocol version'))
return
await stream.write(bytes(4))
async with self.create_connection(stream, address, protocol_version) as conn:
conn: ServerConnection
conn.debug(f'[INIT {address}]')
await conn.init()
await conn.start()
conn.debug(f'[STOP {address}]')
except self.app.config.stream_errors:
pass
async def compress(cls, data: bytes, headers: T_Headers) -> bytes:
headers['Adler32'] = zlib.adler32(data)
return to_uint(len(data), 4) + zlib.compress(data, level=6)
async def send(self, conn):
async with conn.lock_write():
now = time_ns() // 1000_000
await conn.write(self.type_id + to_uint(now, 8))
conn.debug(f'[SEND {conn.address}] PONG {now}')
async def send(self, conn):
async with conn.lock_write():
message_id: int = self.data
await conn.write(self.type_id + to_uint(message_id, 2))
conn.debug(f'[SEND {conn.address}] CANCEL Input M: {message_id}')
async def set_download_speed(self, speed=0):
await self.write(b'\x05')
await self.write(to_uint(speed, 4))
