refactor: sanitize AI code

This commit is contained in:
Stanislav N Mikhailov
2026-08-03 14:52:42 +03:00
parent 0f2cb0384c
commit 7bacf7ac73
3 changed files with 139 additions and 95 deletions
+119 -81
View File
@@ -1,10 +1,25 @@
use std::net::SocketAddr; use std::{
io,
net::SocketAddr,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::{
use tokio::net::{TcpStream}; io::{AsyncReadExt, AsyncWriteExt},
use std::io; 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 { pub struct Connection {
id: u64, id: u64,
client_addr: SocketAddr, client_addr: SocketAddr,
@@ -28,99 +43,122 @@ impl Connection {
pub async fn run(mut self) -> io::Result<()> { pub async fn run(mut self) -> io::Result<()> {
println!( 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, self.id,
_remaining.len() self.client_addr
);
header
}
Err(error) => {
println!(
"Обработчик {}: ошибка парсинга заголовка: {:?}",
self.id, error
); );
return Err(io::Error::new( while let Some(message) = self.read_message().await? {
io::ErrorKind::InvalidData,
format!("Invalid header: {error:?}"),
));
}
};
//
if let Err(error) = header.validate() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
error,
));
}
self.message_count += 1; self.message_count += 1;
println!("Обработчик {}: Успешно распарсен заголовок №{}: {:?}", self.id, self.message_count, header); self.handle_message(message).await?;
}
// 4. Используем данные из распарсенного заголовка! println!(
// Теперь мы знаем точный размер тела сообщения благодаря header.body_size "Соединение {} с {} закрыто клиентом",
const MAX_BODY_SIZE: u64 = 1024 * 1024; // 1 МиБ self.id,
self.client_addr
);
if header.body_size() > MAX_BODY_SIZE { Ok(())
}
async fn read_message(
&mut self,
) -> io::Result<Option<IncomingMessage>> {
let Some(header) = self.read_header().await? else {
return Ok(None);
};
let body = self.read_body(header.body_size()).await?;
Ok(Some(IncomingMessage {
header,
body,
}))
}
async fn read_header(&mut self) -> io::Result<Option<Header>> {
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<Vec<u8>> {
if body_size > MAX_BODY_SIZE {
return Err(io::Error::new( return Err(io::Error::new(
io::ErrorKind::InvalidData, io::ErrorKind::InvalidData,
format!( format!(
"Размер тела {} превышает лимит {}", "Размер тела {body_size} превышает лимит \
header.body_size(), {MAX_BODY_SIZE}"
MAX_BODY_SIZE,
), ),
)); ));
} }
let body_len = header.body_size() as usize; let mut body = vec![0_u8; body_size as usize];
// Подгоняем размер буфера под размер тела (чтобы не читать лишнего) self.stream.read_exact(&mut body).await?;
if body_len > body_buffer.len() {
body_buffer.resize(body_len, 0); Ok(body)
} }
// Читаем из сети тело сообщения (ровно столько байт, сколько указано в заголовке) async fn handle_message(
self.stream.read_exact(&mut body_buffer[..body_len]).await?; &mut self,
message: IncomingMessage,
let received_body = String::from_utf8_lossy(&body_buffer[..body_len]); ) -> io::Result<()> {
println!("Обработчик {}: Получено тело сообщения: {:?}", self.id, received_body); println!(
"Соединение {}: сообщение №{}, заголовок: {:?}",
// 5. Отправляем ответ клиенту self.id,
let response = format!( self.message_count,
"Обработчик {}: Заголовок валиден (cmd={}). Тело принято.\n", message.header
self.id, header.command()
); );
// Это знание конкретной команды.
// Позже этот код можно вынести в отдельный 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?; self.stream.write_all(response.as_bytes()).await?;
Ok(())
} }
} }
}
+11 -11
View File
@@ -22,17 +22,17 @@
* `bytes[0..4]``magic`; * `bytes[0..4]``magic`;
* `bytes[4]``version`; * `bytes[4]``version`;
* `bytes[5]``command`. * `bytes[5]``command`.
* [ ] Проверять `magic` как четыре байта `NCLD`. * [✔] Проверять `magic` как четыре байта `NCLD`.
* [ ] Проверять версию как число `1`, а не ASCII-символ `b'1'`. * [✔] Проверять версию как число `1`, а не ASCII-символ `b'1'`.
* [ ] Пока разрешать только версию протокола `1`. * [✔] Пока разрешать только версию протокола `1`.
* [ ] Исправить команды `80 | 81` на `0x80 | 0x81`. * [✔] Исправить команды `80 | 81` на `0x80 | 0x81`.
* [ ] Записывать в `Header.version` значение из `bytes[4]`. * [✔] Записывать в `Header.version` значение из `bytes[4]`.
* [ ] Для первого теста отправлять заголовок с `body_size = 0`. * [✔] Для первого теста отправлять заголовок с `body_size = 0`.
## Проверка исправлений ## Проверка исправлений
* [ ] Отправить корректный заголовок. * [✔] Отправить корректный заголовок.
* [ ] Убедиться, что сервер вывел: * [✔] Убедиться, что сервер вывел:
```text ```text
Header { Header {
@@ -44,9 +44,9 @@
} }
``` ```
* [ ] Отправить заголовок с неправильным `magic`. * [] Отправить заголовок с неправильным `magic`.
* [ ] Убедиться, что сервер вернул `InvalidFormat`. * [] Убедиться, что сервер вернул `InvalidFormat`.
* [ ] Отправить заголовок с неподдерживаемой версией. * [] Отправить заголовок с неподдерживаемой версией.
* [ ] Убедиться, что сервер вернул `InvalidFormat`. * [ ] Убедиться, что сервер вернул `InvalidFormat`.
* [ ] Отправить заголовок с неизвестной командой. * [ ] Отправить заголовок с неизвестной командой.
* [ ] Убедиться, что сервер вернул `InvalidFormat`. * [ ] Убедиться, что сервер вернул `InvalidFormat`.
+6
View File
@@ -1,4 +1,7 @@
// Парсим заголовок на базе Nom // Парсим заголовок на базе Nom
pub const HEADER_SIZE: usize = 20;
pub const MAX_BODY_SIZE: u64 = 1024 * 1024;
use nom::{ use nom::{
bytes::complete::tag, bytes::complete::tag,
combinator::verify, combinator::verify,
@@ -82,6 +85,9 @@ impl Header {
pub fn body_size(&self) -> u64 { pub fn body_size(&self) -> u64 {
self.body_size self.body_size
} }
///Валидаторы для данных из заголовка:
///
// Валидация Header // Валидация Header
pub fn validate(&self) -> Result<(), HeaderValidationError> { pub fn validate(&self) -> Result<(), HeaderValidationError> {
if self.version != 1 { if self.version != 1 {