tcp_server/
handler_base.rs

1//! The basic trait for handler.
2
3use std::net::SocketAddr;
4use async_trait::async_trait;
5use tcp_handler::common::AesCipher;
6use tcp_handler::bytes::{Buf, BytesMut};
7use tcp_handler::flate2::Compression;
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9use crate::mutable_cipher::MutableCipher;
10use crate::network::{NetworkError, recv, send};
11
12/// A wrapper of the stream with some common methods.
13pub struct IOStream<R, W> where R: AsyncReadExt + Unpin + Send, W: AsyncWriteExt + Unpin + Send {
14    pub(crate) receiver: R,
15    pub(crate) sender: W,
16    pub(crate) cipher: MutableCipher,
17    addr: SocketAddr,
18    version: String,
19}
20
21impl<R, W> IOStream<R, W> where R: AsyncReadExt + Unpin + Send, W: AsyncWriteExt + Unpin + Send {
22    pub(crate) fn new(receiver: R, sender: W, cipher: AesCipher, addr: SocketAddr, version: String) -> Self {
23        Self { receiver, sender, cipher: MutableCipher::new(cipher), addr, version }
24    }
25
26    /// Get the addr of the client.
27    pub fn get_addr(&self) -> &SocketAddr {
28        &self.addr
29    }
30
31    /// Get the version of the client.
32    pub fn get_version(&self) -> &str {
33        &self.version
34    }
35
36    /// Send a message to the client.
37    pub async fn send<B: Buf + Send>(&mut self, message: &mut B) -> Result<(), NetworkError> {
38        let (cipher, guard) = self.cipher.get().await?;
39        let cipher = send(&mut self.sender, message, cipher, Compression::default()).await?;
40        self.cipher.reset(guard, cipher);
41        Ok(())
42    }
43
44    /// Recv a message from the client.
45    pub async fn recv(&mut self) -> Result<BytesMut, NetworkError> {
46        let (cipher, guard) = self.cipher.get().await?;
47        let (response, cipher) = recv(&mut self.receiver, cipher).await?;
48        self.cipher.reset(guard, cipher);
49        Ok(response)
50    }
51
52    /// A shortcut of send and recv message.
53    pub async fn send_recv<B: Buf + Send>(&mut self, message: &mut B) -> Result<BytesMut, NetworkError> {
54        self.send(message).await?;
55        self.recv().await
56    }
57}
58
59/// The basic trait for handler.
60#[async_trait]
61pub trait FuncHandler<R, W>: Send where R: AsyncReadExt + Unpin + Send, W: AsyncWriteExt + Unpin + Send {
62    /// Define your handler here.
63    async fn handle(&self, stream: &mut IOStream<R, W>) -> anyhow::Result<()>;
64}
65
66/// Conveniently define a handler.
67///
68/// The struct is zero-sized.
69#[macro_export]
70macro_rules! func_handler {
71    ($vi: vis $name: ident, |$stream: ident| $block: expr) => {
72        $vi struct $name;
73        #[$crate::async_trait::async_trait]
74        impl<R: $crate::tokio::io::AsyncReadExt + Unpin + Send, W: $crate::tokio::io::AsyncWriteExt + Unpin + Send> $crate::handler_base::FuncHandler<R, W> for $name {
75            async fn handle(&self, $stream: &mut $crate::handler_base::IOStream<R, W>) -> $crate::anyhow::Result<()> {
76                $block
77            }
78        }
79    };
80}
tcp_server/handler_base.rs

tcp_server/
handler_base.rs
