refactor: sanitize AI code

2026-08-03 14:52:42 +03:00
parent 0f2cb0384c
commit 7bacf7ac73
3 changed files with 139 additions and 95 deletions
@@ -1,10 +1,25 @@
-use std::net::SocketAddr;
+use std::{
+    io,
+    net::SocketAddr,
+};

-use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::{TcpStream};
-use std::io;
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::TcpStream,
+};
+
+use crate::protocol::{
+    parse_header,
+    Header,
+    HEADER_SIZE,
+    MAX_BODY_SIZE,
+};
+
+struct IncomingMessage {
+    header: Header,
+    body: Vec<u8>,
+}

-use crate::protocol::parse_header; 
 pub struct Connection {
    id: u64,
    client_addr: SocketAddr,
@@ -28,99 +43,122 @@ impl Connection {

    pub async fn run(mut self) -> io::Result<()> {
        println!(
-        "Обработчик {} создан для {}",
-        self.id, self.client_addr
-    );
-
-    // Буфер для данных тела сообщения (теперь он не хранит заголовок)
-    let mut body_buffer = vec![0_u8; 1024];
-
-    loop {
-        // 1. Создаем буфер СТРОГО под размер заголовка
-        let mut header_buffer = [0_u8; 20];
-
-        // 2. Читаем из сети РОВНО 20 байт
-        // Если клиент закроет соединение до того, как пришлет 20 байт,
-        // read_exact вернет ошибку UnexpectedEof.
-        if let Err(e) = self.stream.read_exact(&mut header_buffer).await {
-            if e.kind() == io::ErrorKind::UnexpectedEof {
-                println!(
-                    "Обработчик {}: клиент {} отключился (или прислал неполный заголовок)",
-                    self.id, self.client_addr
-                );
-                return Ok(());
-            }
-            return Err(e); // Другие системные ошибки ввода-вывода (например, обрыв связи)
-        }
-
-        // 3. Вызываем ваш парсер заголовка
-        let header = match parse_header(&header_buffer) {
-            Ok((_remaining, header)) => {
-                println!(
-                    "Обработчик {}: nom оставил {} байт",
+            "Соединение {} установлено с {}",
            self.id,
-                    _remaining.len()
+            self.client_addr
        );
-                header
+
+        while let Some(message) = self.read_message().await? {
+            self.message_count += 1;
+            self.handle_message(message).await?;
        }

-            Err(error) => {
        println!(
-                    "Обработчик {}: ошибка парсинга заголовка: {:?}",
-                    self.id, error
+            "Соединение {} с {} закрыто клиентом",
+            self.id,
+            self.client_addr
        );

-        return Err(io::Error::new(
-            io::ErrorKind::InvalidData,
-            format!("Invalid header: {error:?}"),
-        ));
+        Ok(())
    }
+
+    async fn read_message(
+        &mut self,
+    ) -> io::Result<Option<IncomingMessage>> {
+        let Some(header) = self.read_header().await? else {
+            return Ok(None);
        };

-// 
-if let Err(error) = header.validate() {
-    return Err(io::Error::new(
-        io::ErrorKind::InvalidData,
-        error,
-    ));
+        let body = self.read_body(header.body_size()).await?;
+
+        Ok(Some(IncomingMessage {
+            header,
+            body,
+        }))
    }
-        self.message_count += 1;
-        println!("Обработчик {}: Успешно распарсен заголовок №{}: {:?}", self.id, self.message_count, header);

-        // 4. Используем данные из распарсенного заголовка!
-        // Теперь мы знаем точный размер тела сообщения благодаря header.body_size
-        const MAX_BODY_SIZE: u64 = 1024 * 1024; // 1 МиБ
+    async fn read_header(&mut self) -> io::Result<Option<Header>> {
+        let mut buffer = [0_u8; HEADER_SIZE];

-        if header.body_size() > MAX_BODY_SIZE {
+        // Первый байт читаем отдельно, чтобы отличить нормальное
+        // закрытие соединения от оборванного заголовка.
+        let bytes_read = self.stream.read(&mut buffer[..1]).await?;
+
+        if bytes_read == 0 {
+            return Ok(None);
+        }
+
+        self.stream.read_exact(&mut buffer[1..]).await?;
+
+        let (remaining, header) =
+            parse_header(&buffer).map_err(|error| {
+                io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    format!("Ошибка разбора заголовка: {error:?}"),
+                )
+            })?;
+
+        // Мы дали парсеру ровно HEADER_SIZE байт.
+        debug_assert!(remaining.is_empty());
+
+        header.validate().map_err(|error| {
+            io::Error::new(io::ErrorKind::InvalidData, error)
+        })?;
+
+        Ok(Some(header))
+    }
+
+    async fn read_body(
+        &mut self,
+        body_size: u64,
+    ) -> io::Result<Vec<u8>> {
+        if body_size > MAX_BODY_SIZE {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!(
-                    "Размер тела {} превышает лимит {}",
-                    header.body_size(),
-                    MAX_BODY_SIZE,
+                    "Размер тела {body_size} превышает лимит \
+                     {MAX_BODY_SIZE}"
                ),
            ));
        }

-        let body_len = header.body_size() as usize;
+        let mut body = vec![0_u8; body_size as usize];

-        // Подгоняем размер буфера под размер тела (чтобы не читать лишнего)
-        if body_len > body_buffer.len() {
-            body_buffer.resize(body_len, 0);
+        self.stream.read_exact(&mut body).await?;
+
+        Ok(body)
    }

-        // Читаем из сети тело сообщения (ровно столько байт, сколько указано в заголовке)
-        self.stream.read_exact(&mut body_buffer[..body_len]).await?;
-
-        let received_body = String::from_utf8_lossy(&body_buffer[..body_len]);
-        println!("Обработчик {}: Получено тело сообщения: {:?}", self.id, received_body);
-
-        // 5. Отправляем ответ клиенту
-        let response = format!(
-            "Обработчик {}: Заголовок валиден (cmd={}). Тело принято.\n",
-            self.id, header.command()
+    async fn handle_message(
+        &mut self,
+        message: IncomingMessage,
+    ) -> io::Result<()> {
+        println!(
+            "Соединение {}: сообщение №{}, заголовок: {:?}",
+            self.id,
+            self.message_count,
+            message.header
        );
+
+        // Это знание конкретной команды.
+        // Позже этот код можно вынести в отдельный handler.rs.
+        let body_text = String::from_utf8_lossy(&message.body);
+
+        println!(
+            "Соединение {}: тело: {:?}",
+            self.id,
+            body_text
+        );
+
+        let response = format!(
+            "Соединение {}: запрос {} принят, cmd={}\n",
+            self.id,
+            message.header.request_id(),
+            message.header.command(),
+        );
+
        self.stream.write_all(response.as_bytes()).await?;
-    }
+
+        Ok(())
    }
 }
@@ -22,17 +22,17 @@
    * `bytes[0..4]` — `magic`;
    * `bytes[4]` — `version`;
    * `bytes[5]` — `command`.
-    * [ ] Проверять `magic` как четыре байта `NCLD`.
-    * [ ] Проверять версию как число `1`, а не ASCII-символ `b'1'`.
-    * [ ] Пока разрешать только версию протокола `1`.
-    * [ ] Исправить команды `80 | 81` на `0x80 | 0x81`.
-    * [ ] Записывать в `Header.version` значение из `bytes[4]`.
-    * [ ] Для первого теста отправлять заголовок с `body_size = 0`.
+    * [✔] Проверять `magic` как четыре байта `NCLD`.
+    * [✔] Проверять версию как число `1`, а не ASCII-символ `b'1'`.
+    * [✔] Пока разрешать только версию протокола `1`.
+    * [✔] Исправить команды `80 | 81` на `0x80 | 0x81`.
+    * [✔] Записывать в `Header.version` значение из `bytes[4]`.
+    * [✔] Для первого теста отправлять заголовок с `body_size = 0`.

    ## Проверка исправлений

-    * [ ] Отправить корректный заголовок.
-    * [ ] Убедиться, что сервер вывел:
+    * [✔] Отправить корректный заголовок.
+    * [✔] Убедиться, что сервер вывел:

    ```text
    Header {
@@ -44,9 +44,9 @@
    }
    ```

-    * [ ] Отправить заголовок с неправильным `magic`.
-    * [ ] Убедиться, что сервер вернул `InvalidFormat`.
-    * [ ] Отправить заголовок с неподдерживаемой версией.
+    * [✔] Отправить заголовок с неправильным `magic`.
+    * [✔] Убедиться, что сервер вернул `InvalidFormat`.
+    * [✔] Отправить заголовок с неподдерживаемой версией.
    * [ ] Убедиться, что сервер вернул `InvalidFormat`.
    * [ ] Отправить заголовок с неизвестной командой.
    * [ ] Убедиться, что сервер вернул `InvalidFormat`.
@@ -1,4 +1,7 @@
 // Парсим заголовок на базе Nom
+pub const HEADER_SIZE: usize = 20;
+pub const MAX_BODY_SIZE: u64 = 1024 * 1024;
+
 use nom::{
    bytes::complete::tag,
    combinator::verify,
@@ -82,6 +85,9 @@ impl Header {
    pub fn body_size(&self) -> u64 {
        self.body_size
    }
+
+    ///Валидаторы для данных из заголовка:
+    /// 
    // Валидация Header
    pub fn validate(&self) -> Result<(), HeaderValidationError> {
    if self.version != 1 {