Stanislav N Mikhailov
159 lines
4.0 KiB
Rust
159 lines
4.0 KiB
Rust
use std::{
|
||
io,
|
||
net::SocketAddr,
|
||
};
|
||
|
||
use tokio::{
|
||
io::{AsyncReadExt, AsyncWriteExt},
|
||
net::TcpStream,
|
||
};
|
||
|
||
use crate::protocol::{
|
||
decode_header,
|
||
Header,
|
||
HEADER_SIZE,
|
||
MAX_BODY_SIZE,
|
||
};
|
||
|
||
struct IncomingMessage {
|
||
header: Header,
|
||
body: Vec<u8>,
|
||
}
|
||
|
||
pub struct Connection {
|
||
id: u64,
|
||
client_addr: SocketAddr,
|
||
stream: TcpStream,
|
||
message_count: u64,
|
||
}
|
||
|
||
impl Connection {
|
||
pub fn new(
|
||
id: u64,
|
||
client_addr: SocketAddr,
|
||
stream: TcpStream,
|
||
) -> Self {
|
||
Self {
|
||
id,
|
||
client_addr,
|
||
stream,
|
||
message_count: 0,
|
||
}
|
||
}
|
||
|
||
pub async fn run(mut self) -> io::Result<()> {
|
||
println!(
|
||
"Соединение {} установлено с {}",
|
||
self.id,
|
||
self.client_addr
|
||
);
|
||
|
||
while let Some(message) = self.read_message().await? {
|
||
self.message_count += 1;
|
||
self.handle_message(message).await?;
|
||
}
|
||
|
||
println!(
|
||
"Соединение {} с {} закрыто клиентом",
|
||
self.id,
|
||
self.client_addr
|
||
);
|
||
|
||
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);
|
||
}
|
||
|
||
// Читаем остальную часть буфера от 1 до конца
|
||
self.stream.read_exact(&mut buffer[1..]).await?;
|
||
|
||
//Парсим скачанный заголовок через nocloud_core::protocol parse_header
|
||
let header =
|
||
decode_header(&buffer).map_err(|error| {
|
||
io::Error::new(
|
||
io::ErrorKind::InvalidData,
|
||
format!("Ошибка разбора заголовка: {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!(
|
||
"Размер тела {body_size} превышает лимит \
|
||
{MAX_BODY_SIZE}"
|
||
),
|
||
));
|
||
}
|
||
|
||
let mut body = vec![0_u8; body_size as usize];
|
||
|
||
self.stream.read_exact(&mut body).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(())
|
||
}
|
||
} |