1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
//! The basic trait for client.

use std::io::{Error, ErrorKind};
use async_trait::async_trait;
use tcp_handler::bytes::{Buf, BufMut, BytesMut};
use tcp_handler::flate2::Compression;
use tcp_handler::variable_len_reader::{VariableReader, VariableWriter};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use crate::mutable_cipher::MutableCipher;
use crate::network::{NetworkError, recv, send};

/// The basic trait for client.
/// See [`ClientFactory`][crate::ClientFactory] for example.
#[async_trait]
pub trait ClientBase<R, W> where R: AsyncReadExt + Unpin + Send, W: AsyncWriteExt + Unpin + Send {
    /// Get the receiver and mutable cipher.
    /// # Note:
    /// This should be a const expr.
    fn get_receiver<'a>(&'a mut self) -> (&'a mut R, &MutableCipher);

    /// Get the sender and mutable cipher.
    /// # Note:
    /// This should be a const expr.
    fn get_sender<'a>(&'a mut self) -> (&'a mut W, &MutableCipher);

    /// Send a message to the server.
    async fn send<B: Buf + Send>(&mut self, message: &mut B) -> Result<(), NetworkError> {
        let (sender, mutable_cipher) = self.get_sender();
        let (cipher, guard) = mutable_cipher.get().await?;
        match send(sender, message, cipher, Compression::default()).await {
            Ok(cipher) => { mutable_cipher.reset(guard, Some(cipher)); Ok(()) }
            Err(e) => { mutable_cipher.reset(guard, None); Err(e) }
        }
    }

    /// Recv a message from the server.
    async fn recv(&mut self) -> Result<BytesMut, NetworkError> {
        let (receiver, mutable_cipher) = self.get_receiver();
        let (cipher, guard) = mutable_cipher.get().await?;
        match recv(receiver, cipher).await {
            Ok((response, cipher)) => { mutable_cipher.reset(guard, Some(cipher)); Ok(response) }
            Err(e) => { mutable_cipher.reset(guard, None); Err(e) }
        }
    }

    /// A shortcut of send and recv message.
    async fn send_recv<B: Buf + Send>(&mut self, message: &mut B) -> Result<BytesMut, NetworkError> {
        self.send(message).await?;
        self.recv().await
    }

    /// Check if the function is supported by the server.
    ///
    /// This is corresponding to `tcp-server` crate.
    async fn check_func(&mut self, func: &str) -> Result<(), NetworkError> {
        let mut writer = BytesMut::new().writer();
        writer.write_string(func)?;
        let mut reader = self.send_recv(&mut writer.into_inner()).await?.reader();
        if reader.read_bool()? {
            Ok(())
        } else {
            Err(NetworkError::BufError(Error::new(ErrorKind::Other, format!("func is not available: {}", func))))
        }
    }
}