diff --git a/src/oneme_tcp/proto.py b/src/oneme_tcp/proto.py index 7deb2c1..0c085ba 100644 --- a/src/oneme_tcp/proto.py +++ b/src/oneme_tcp/proto.py @@ -4,8 +4,19 @@ class Proto: def __init__(self) -> None: self.logger = logging.getLogger(__name__) + # TODO узнать какие должны быть лимиты и поменять, + # сейчас это больше заглушка + MAX_PAYLOAD_SIZE = 1048576 # 1 MB + MAX_DECOMPRESSED_SIZE = 1048576 # 1 MB + HEADER_SIZE = 10 # 1+2+1+2+4 + ### Работа с протоколом def unpack_packet(self, data: bytes) -> dict | None: + # Проверяем минимальный размер пакета + if len(data) < self.HEADER_SIZE: + self.logger.warning(f"Пакет слишком маленький: {len(data)} байт") + return None + # Распаковываем заголовок ver = int.from_bytes(data[0:1], "big") cmd = int.from_bytes(data[1:3], "big") @@ -18,6 +29,17 @@ class Proto: # Парсим данные пакета payload_length = packed_len & 0xFFFFFF + + # Проверяем размер payload + if payload_length > self.MAX_PAYLOAD_SIZE: + self.logger.warning(f"Payload слишком большой: {payload_length} B (лимит {self.MAX_PAYLOAD_SIZE})") + return None + + # Проверяем длину пакета + if len(data) < self.HEADER_SIZE + payload_length: + self.logger.warning(f"Пакет неполный: требуется {self.HEADER_SIZE + payload_length} B, получено {len(data)}") + return None + payload_bytes = data[10 : 10 + payload_length] payload = None @@ -27,14 +49,14 @@ class Proto: if comp_flag != 0: compressed_data = payload_bytes try: - payload_bytes = lz4.block.decompress( compressed_data, - uncompressed_size=99999, + uncompressed_size=self.MAX_DECOMPRESSED_SIZE, ) except lz4.block.LZ4BlockError: + self.logger.warning("Ошибка декомпрессии LZ4") return None - + # Распаковываем msgpack payload = msgpack.unpackb(payload_bytes, raw=False, strict_map_key=False) diff --git a/src/oneme_tcp/server.py b/src/oneme_tcp/server.py index 49d4d91..75023f2 100644 --- a/src/oneme_tcp/server.py +++ b/src/oneme_tcp/server.py @@ -16,9 +16,12 @@ class OnemeMobileServer: self.proto = Proto() self.processors = Processors(db_pool=db_pool, clients=clients, send_event=send_event, telegram_bot=telegram_bot) - # rate limiter anti ddos brute force protection here + # rate limiter anti ddos brute force protection self.auth_rate_limiter = RateLimiter(max_attempts=5, window_seconds=60) + self.read_timeout = 300 # Таймаут чтения из сокета (секунды) + self.max_read_size = 65536 # Максимальный размер данных из сокета + async def handle_client(self, reader, writer): """Функция для обработки подключений""" # IP-адрес клиента @@ -34,16 +37,33 @@ class OnemeMobileServer: try: while True: - # Читаем новые данные из сокета - data = await reader.read(4098) + # Читаем новые данные из сокета с таймаутом + try: + data = await asyncio.wait_for( + reader.read(self.max_read_size), + timeout=self.read_timeout + ) + except asyncio.TimeoutError: + self.logger.info(f"Таймаут соединения для {address[0]}:{address[1]}") + break # Если сокет закрыт - выходим из цикла if not data: break + + if len(data) > self.max_read_size: + self.logger.warning(f"Пакет от {address[0]}:{address[1]} превышает лимит ({len(data)} байт)") + break + # Распаковываем данные packet = self.proto.unpack_packet(data) + # Скип если пакет невалидный + if packet is None: + self.logger.warning(f"Невалидный пакет от {address[0]}:{address[1]}") + continue + opcode = packet.get("opcode") seq = packet.get("seq") payload = packet.get("payload") diff --git a/src/tamtam_tcp/proto.py b/src/tamtam_tcp/proto.py index 26ae97a..e055224 100644 --- a/src/tamtam_tcp/proto.py +++ b/src/tamtam_tcp/proto.py @@ -4,8 +4,19 @@ class Proto: def __init__(self) -> None: self.logger = logging.getLogger(__name__) + # TODO узнать какие должны быть лимиты и поменять, + # сейчас это больше заглушка + MAX_PAYLOAD_SIZE = 1048576 # 1 MB + MAX_DECOMPRESSED_SIZE = 1048576 # 1 MB + HEADER_SIZE = 10 # 1+2+1+2+4 + ### Работа с протоколом def unpack_packet(self, data: bytes) -> dict | None: + # Проверяем минимальный размер пакета + if len(data) < self.HEADER_SIZE: + self.logger.warning(f"Пакет слишком маленький: {len(data)} байт") + return None + # Распаковываем заголовок ver = int.from_bytes(data[0:1], "big") cmd = int.from_bytes(data[1:3], "big") @@ -18,6 +29,17 @@ class Proto: # Парсим данные пакета payload_length = packed_len & 0xFFFFFF + + # Проверяем размер payload + if payload_length > self.MAX_PAYLOAD_SIZE: + self.logger.warning(f"Payload слишком большой: {payload_length} B (лимит {self.MAX_PAYLOAD_SIZE})") + return None + + # Проверяем длину пакета + if len(data) < self.HEADER_SIZE + payload_length: + self.logger.warning(f"Пакет неполный: требуется {self.HEADER_SIZE + payload_length} B, получено {len(data)}") + return None + payload_bytes = data[10 : 10 + payload_length] payload = None @@ -27,14 +49,14 @@ class Proto: if comp_flag != 0: compressed_data = payload_bytes try: - payload_bytes = lz4.block.decompress( compressed_data, - uncompressed_size=99999, + uncompressed_size=self.MAX_DECOMPRESSED_SIZE, ) except lz4.block.LZ4BlockError: + self.logger.warning("Ошибка декомпрессии LZ4") return None - + # Распаковываем msgpack payload = msgpack.unpackb(payload_bytes, raw=False, strict_map_key=False) diff --git a/src/tamtam_tcp/server.py b/src/tamtam_tcp/server.py index 7cc5968..1fd1814 100644 --- a/src/tamtam_tcp/server.py +++ b/src/tamtam_tcp/server.py @@ -19,6 +19,9 @@ class TTMobileServer: # rate limiter self.auth_rate_limiter = RateLimiter(max_attempts=5, window_seconds=60) + self.read_timeout = 300 # Таймаут чтения из сокета (секунды) + self.max_read_size = 65536 # Максимальный размер данных из сокета + async def handle_client(self, reader, writer): """Функция для обработки подключений""" # IP-адрес клиента @@ -34,16 +37,33 @@ class TTMobileServer: try: while True: - # Читаем новые данные из сокета - data = await reader.read(4098) + # Читаем новые данные из сокета (с таймаутом!) + try: + data = await asyncio.wait_for( + reader.read(self.max_read_size), + timeout=self.read_timeout + ) + except asyncio.TimeoutError: + self.logger.info(f"Таймаут соединения для {address[0]}:{address[1]}") + break # Если сокет закрыт - выходим из цикла if not data: break + # Проверяем размер данных + if len(data) > self.max_read_size: + self.logger.warning(f"Пакет от {address[0]}:{address[1]} превышает лимит ({len(data)} байт)") + break + # Распаковываем данные packet = self.proto.unpack_packet(data) + # Если пакет невалидный — пропускаем + if packet is None: + self.logger.warning(f"Невалидный пакет от {address[0]}:{address[1]}") + continue + opcode = packet.get("opcode") seq = packet.get("seq") payload = packet.get("payload") diff --git a/src/tamtam_ws/proto.py b/src/tamtam_ws/proto.py index 6bf093f..bd7417f 100644 --- a/src/tamtam_ws/proto.py +++ b/src/tamtam_ws/proto.py @@ -12,14 +12,18 @@ class Proto: "payload": payload }) + MAX_PACKET_SIZE = 65536 # 64 KB, заглушка, нужно узнать реальные лимиты и поменять, хотя кто будет это делать... + def unpack_packet(self, packet): - # нужно try catch сделать - # чтобы не сыпалось всё при неверных пакетах + # try catch чтобы не сыпалось всё при неверных пакетах + if isinstance(packet, (str, bytes)) and len(packet) > self.MAX_PACKET_SIZE: + return {} + try: parsed_packet = json.loads(packet) - except: + except (json.JSONDecodeError, TypeError, ValueError): return {} - + return parsed_packet # мне кажется долго вручную всё писать # а как еще diff --git a/src/tamtam_ws/server.py b/src/tamtam_ws/server.py index 6fd0dce..340dbc5 100644 --- a/src/tamtam_ws/server.py +++ b/src/tamtam_ws/server.py @@ -21,12 +21,17 @@ class TTWSServer: # Распаковываем пакет packet = self.proto.unpack_packet(message) + if not packet: + self.logger.warning("Невалидный пакет от ws клиента") + continue + # Валидируем структуру пакета try: MessageModel.model_validate(packet) except ValidationError as e: - self.logger.error(e) - + self.logger.warning(f"Ошибка валидации пакета: {e}") + continue + # Извлекаем данные из пакета seq = packet['seq'] opcode = packet['opcode'] @@ -57,5 +62,10 @@ class TTWSServer: async def start(self): self.logger.info(f"Вебсокет запущен на порту {self.port}") - async with serve(self.handle_client, self.host, self.port): + async with serve( + self.handle_client, self.host, self.port, + max_size=65536, + open_timeout=10, + close_timeout=10, + ): await asyncio.Future() \ No newline at end of file