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))