btp 0.1.3

A rust library of, a lightweight protocol, Blog Transfer Protocol.
Documentation
use crate::message::{BtpHeader, BtpPackage};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
use tokio::net::ToSocketAddrs;
use tokio::spawn;
use tokio::sync::Mutex;

type Error = Box<dyn std::error::Error + Send + Sync>;

#[derive(Clone, Copy)]
pub struct BtpConfig {
    pub ver: [u8; 3],
    pub addr: SocketAddr,
}

impl BtpConfig {
    /// Takes the address and returns [`BtpConfig`]
    pub fn from_addr(addr: SocketAddr) -> BtpConfig {
        BtpConfig {
            ver: super::VERSION,
            addr,
        }
    }
}

/// Since we use this function inside the async functions, we need to send, sync it between threads,
/// so we use this [`BtpSocketInner`] wrapped in an Arc<Mutex<>> to use between threads safely without any additional problem.
/// We don't use [`BtpSocketInner`] inside the functions, it's not public because it's embedded into [`BtpSocket`]
#[derive(Clone, Copy)]
struct BtpSocketInner<S: AsyncRead + AsyncWrite + Unpin + Send + Sync> {
    pub socket: S,
}

pub struct BtpSocket<
    S: AsyncRead + AsyncWrite + Unpin + Send + Sync,
    P: ToSocketAddrs + Send + Sync,
> {
    inner: Arc<Mutex<BtpSocketInner<S>>>,
    pub btp_conf: BtpConfig,
    pub peer: P,
}

impl<S: AsyncRead + AsyncWrite + Unpin + Send + Sync> BtpSocketInner<S> {
    pub async fn write_message(&mut self, msg: BtpPackage) -> Result<(), Error> {
        let mut writer = BufWriter::new(&mut self.socket);
        writer.write_all(msg.as_vec().as_slice()).await?;
        writer.flush().await?;

        Ok(())
    }

    /// Returns a Result and a [`BtpFile`] or [`BtpMessage`], with checking the file_name_len.
    /// If it's empty, it understands that coming message is BtpMessage.
    /// If it has a file_name_len, then it is a BtpFile.
    pub async fn read(&mut self) -> Result<BtpPackage, Error> {
        let mut reader = BufReader::new(&mut self.socket);
        let mut header = [0u8; 10];
        reader.read_exact(&mut header).await?;
        let res;
        match &header[8..] {
            [0, 0] => {
                let header = BtpHeader::from_raw(header)?;
                let file_len = header.get_len();
                let mut buf = vec![0u8; file_len];
                reader.read_exact(&mut buf).await?;
                reader.flush().await?;
                res = Ok(BtpPackage::from_header(
                    header,
                    String::from_utf8(buf)?,
                    None,
                ));
            }
            _ => {
                let header = BtpHeader::from_raw(header)?;
                let file_len = header.get_len();
                let file_name_len = header.get_file_name_len();
                let mut file = vec![0u8; file_len];
                let mut file_name = vec![0u8; file_name_len];
                reader.read_exact(&mut file).await?;
                reader.flush().await?;
                reader.read_exact(&mut file_name).await?;
                reader.flush().await?;
                res = Ok(BtpPackage::from_header(
                    header,
                    String::from_utf8(file)?,
                    Some(String::from_utf8(file_name)?),
                ));
            }
        }
        res
    }
}

impl<
    S: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,
    P: ToSocketAddrs + Send + Sync + Clone + Copy + 'static,
> BtpSocket<S, P>
{
    pub(super) fn copy(&mut self) -> Self {
        let peer = self.peer.clone();
        let btp_conf = self.btp_conf.clone();
        let inner = Arc::clone(&self.inner);

        BtpSocket {
            inner,
            btp_conf,
            peer,
        }
    }

    pub async fn write(&mut self, msg: BtpPackage) -> Result<(), Error> {
        let mut inner = self.inner.lock().await;
        inner.write_message(msg).await
    }

    pub async fn read(&mut self) -> Result<BtpPackage, Error> {
        let mut inner = self.inner.lock().await;
        inner.read().await
    }

    pub fn from(socket: S, btp_conf: BtpConfig, peer: P) -> Self {
        BtpSocket {
            inner: Arc::new(Mutex::new(BtpSocketInner { socket })),
            btp_conf,
            peer,
        }
    }

    /// This function takes two function, both for msg and file. Then when a receive happens, this functions will handle it.
    pub fn attach_handler<F1>(&mut self, func: F1) -> Result<(), Error>
    where
        F1: Fn(BtpPackage, BtpSocket<S, P>) + Send + Sync + 'static,
    {
        let inner = Arc::clone(&self.inner);
        let func = Arc::new(func);

        let soc = Arc::new(Mutex::new(self.copy()));

        spawn(async move {
            loop {
                match inner.lock().await.read().await {
                    Ok(packet) => {
                        let f = Arc::clone(&func);
                        let soc = Arc::clone(&soc);
                        std::thread::spawn(move || f(packet, soc.blocking_lock().copy()));
                    }
                    Err(e) => {
                        eprintln!("Couldn't send the packet: {}", e);
                    }
                };
            }
        });
        Ok(())
    }
}