From 7bacf7ac73079bfc11cd20f59d3f17fc47bc1010 Mon Sep 17 00:00:00 2001 From: Stanislav N Mikhailov Date: Mon, 3 Aug 2026 14:52:42 +0300 Subject: [PATCH] refactor: sanitize AI code --- src/connection.rs | 206 +++++++++++++++++++++++++++------------------- src/main.MD | 22 ++--- src/protocol.rs | 6 ++ 3 files changed, 139 insertions(+), 95 deletions(-) diff --git a/src/connection.rs b/src/connection.rs index 72ed7a2..e9b5d31 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -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, +} -use crate::protocol::parse_header; pub struct Connection { id: u64, client_addr: SocketAddr, @@ -26,101 +41,124 @@ impl Connection { } } - pub async fn run(mut self) -> io::Result<()> { - println!( - "Обработчик {} создан для {}", - self.id, self.client_addr - ); + 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); // Другие системные ошибки ввода-вывода (например, обрыв связи) + while let Some(message) = self.read_message().await? { + self.message_count += 1; + self.handle_message(message).await?; } - // 3. Вызываем ваш парсер заголовка - let header = match parse_header(&header_buffer) { - Ok((_remaining, header)) => { - println!( - "Обработчик {}: nom оставил {} байт", - self.id, - _remaining.len() - ); - header - } - - Err(error) => { - println!( - "Обработчик {}: ошибка парсинга заголовка: {:?}", - self.id, error - ); + println!( + "Соединение {} с {} закрыто клиентом", + self.id, + self.client_addr + ); - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Invalid header: {error:?}"), - )); + Ok(()) } -}; -// -if let Err(error) = header.validate() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - error, - )); -} - self.message_count += 1; - println!("Обработчик {}: Успешно распарсен заголовок №{}: {:?}", self.id, self.message_count, header); + async fn read_message( + &mut self, + ) -> io::Result> { + let Some(header) = self.read_header().await? else { + return Ok(None); + }; - // 4. Используем данные из распарсенного заголовка! - // Теперь мы знаем точный размер тела сообщения благодаря header.body_size - const MAX_BODY_SIZE: u64 = 1024 * 1024; // 1 МиБ + let body = self.read_body(header.body_size()).await?; - if header.body_size() > MAX_BODY_SIZE { + Ok(Some(IncomingMessage { + header, + body, + })) + } + + async fn read_header(&mut self) -> io::Result> { + let mut buffer = [0_u8; HEADER_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> { + 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; - - // Подгоняем размер буфера под размер тела (чтобы не читать лишнего) - if body_len > body_buffer.len() { - body_buffer.resize(body_len, 0); - } + let mut body = vec![0_u8; body_size as usize]; - // Читаем из сети тело сообщения (ровно столько байт, сколько указано в заголовке) - self.stream.read_exact(&mut body_buffer[..body_len]).await?; + self.stream.read_exact(&mut body).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() - ); - self.stream.write_all(response.as_bytes()).await?; + Ok(body) + } + + 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(()) } -} } \ No newline at end of file diff --git a/src/main.MD b/src/main.MD index 3c7e25f..92f485f 100644 --- a/src/main.MD +++ b/src/main.MD @@ -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`. diff --git a/src/protocol.rs b/src/protocol.rs index 03ff6ec..801206e 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -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 {