nardol/packet/

mod.rs

mod packet_kind;

use serde::{Deserialize, Serialize};

use itertools::Itertools;

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream, ToSocketAddrs, UdpSocket};
use std::sync::atomic::{AtomicU16, Ordering};

use crate::bytes::Bytes;
use crate::error::{Error, ErrorKind};
use crate::ron::{FromRon, ToRon};

pub use packet_kind::PacketKind;

/// For information go to [Packet::description_size].
const PACKET_DESCRIPTION_SIZE: u16 = 4;

/// For information go to [Packet::max_size].
static mut MAX_PACKET_SIZE: AtomicU16 = AtomicU16::new(1024);

/// Gives structure to data to be sent or received from `stream`.
///
/// [Packet] is the lowest abstraction above buffer in this library.
///
/// # Fields
///
/// * `size` -- size of the whole [packet](Packet) in number of bytes.
/// It is [u16] so that packet can not have size over [u16::MAX]
/// * `kind` -- [kind](PacketKind) of [packet](Packet).
/// * `content` -- data stored in the [packet](Packet).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Packet {
    size: u16,
    kind: PacketKind,
    content: Bytes,
}

impl ToRon for Packet {}
impl FromRon<'_> for Packet {}

impl Default for Packet {
    fn default() -> Self {
        Packet {
            size: Packet::description_size(),
            kind: PacketKind::Empty,
            content: Bytes::empty(),
        }
    }
}

impl TryFrom<Bytes> for Packet {
    type Error = Error;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        // Checks if value has enough bytes to parse to Packet
        // (at least 4 for packet with no content).
        if value
            .get((Packet::description_size() - 1) as usize)
            .is_none()
        {
            return Err(Error::new(
                ErrorKind::InvalidBufferSize,
                Some(
                    "impl TryFrom<Bytes> for Packet requires buffer of length of at least 4 bytes."
                        .to_string(),
                ),
            ));
        }

        let size = value.len();
        // Starts at 2 because size field takes 2 bytes in buffer.
        let kind = PacketKind::try_from(&value[2..4])?;
        let content = value[4..size].to_vec();

        Ok(Packet {
            size: size as u16,
            kind,
            content: content.into(),
        })
    }
}

impl TryFrom<&Bytes> for Packet {
    type Error = Error;

    fn try_from(value: &Bytes) -> Result<Self, Self::Error> {
        // Checks if value has enough bytes to parse to Packet
        // (at least 4 for packet with no content).
        if value
            .get((Packet::description_size() - 1) as usize)
            .is_none()
        {
            return Err(Error::new(
                ErrorKind::InvalidBufferSize,
                Some(
                    "impl TryFrom<Bytes> for Packet requires buffer of length of at least 4 bytes."
                        .to_string(),
                ),
            ));
        }

        let size = value.len();
        // Starts at 2 because size field takes 2 bytes in buffer.
        let kind = PacketKind::try_from(&value[2..4])?;
        let content = value[4..size].to_vec();

        Ok(Packet {
            size: size as u16,
            kind,
            content: content.into(),
        })
    }
}

impl TryFrom<&[u8]> for Packet {
    type Error = Error;

    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        // Checks if value has enough bytes to parse to Packet
        // (at least 4 for packet with no content).
        if value
            .get((Packet::description_size() - 1) as usize)
            .is_none()
        {
            return Err(Error::new(
                ErrorKind::InvalidBufferSize,
                Some(
                    "impl TryFrom<Bytes> for Packet requires buffer of length of at least 4 bytes."
                        .to_string(),
                ),
            ));
        }

        let size = value.len();
        // Starts at 2 because size field takes 2 bytes in buffer.
        let kind = PacketKind::try_from(&value[2..4])?;
        let content = value[4..size].to_vec();

        Ok(Packet {
            size: size as u16,
            kind,
            content: content.into(),
        })
    }
}

impl From<Packet> for Bytes {
    fn from(mut value: Packet) -> Self {
        let mut bytes = Bytes::empty();

        bytes.append(&mut value.size.into());
        bytes.append(&mut value.kind.into());
        bytes.append(&mut value.content);

        bytes
    }
}

impl Packet {
    /// Creates a new [Packet].
    ///
    /// Returns an [Error] if created packet would be bigger than
    /// [`MAX_PACKET_SIZE`](Packet::max_size).
    ///
    /// To avoid this check, use [Packet::new_unchecked].
    ///
    /// Size of packet is derived from `content`.
    ///
    /// # Examples
    /// Content packet is created like that.
    /// ```
    /// # use nardol::packet::{Packet, PacketKind};
    /// # use nardol::bytes::Bytes;
    /// # fn packet_new_doc_test() -> Result<(), nardol::error::Error> {
    /// #   
    ///     let packet = Packet::new(PacketKind::End, Bytes::from([1_u8, 2_u8]))?;
    /// #   Ok(())
    /// # }
    /// ```
    pub fn new(kind: PacketKind, content: Bytes) -> Result<Self, Error> {
        // Size is composed of three parts:
        // Size of size field which is always 2.
        // Size of PacketKind which is always 2.
        // Size of content.
        let size = Packet::description_size() + content.len() as u16;

        if size > Packet::max_size() {
            return Err(Error::new(
                ErrorKind::TooBigPacket,
                Some("Created Packet would be bigger than MAX_PACKET_SIZE.".to_string()),
            ));
        }

        Ok(Packet {
            size: size as u16,
            kind,
            content,
        })
    }

    /// Creates a new [Packet] without checking if resulting packet size is not bigger than
    /// [`MAX_PACKET_SIZE`](Packet::max_size).
    ///
    /// Size of packet is derived from `content`.
    pub fn new_unchecked(kind: PacketKind, content: Bytes) -> Self {
        // Size is composed of three parts:
        // Size of size field which is always 2.
        // Size of PacketKind which is always 2.
        // Size of content.
        let size = Packet::description_size() + content.len() as u16;

        Packet {
            size: size as u16,
            kind,
            content,
        }
    }

    /// Sends `self` via `stream`.
    ///
    /// This takes ownership of `self`.
    pub fn send(self, stream: &mut TcpStream) -> Result<(), Error> {
        let packet_bytes: Bytes = self.into();
        let packet_buff = packet_bytes.as_slice();

        if let Err(e) = stream.write(packet_buff) {
            let packet = Packet::try_from(packet_buff)?;
            return Err(Error::new(
                ErrorKind::WritingToStreamFailed,
                Some(format!(
                    "Failed to write packet to stream.\n{}\n{}",
                    packet.to_ron_pretty(None)?,
                    e,
                )),
            ));
        }

        Ok(())
    }

    /// Sends `self` to given `address` on [UdpSocket].
    ///
    /// This takes ownership of `self`.
    pub fn send_to<A>(self, socket: UdpSocket, address: A) -> Result<(), Error>
    where
        A: ToSocketAddrs,
    {
        let packet_bytes: Bytes = self.into();
        let packet_buff = packet_bytes.as_slice();

        if let Err(e) = socket.send_to(packet_buff, address) {
            let packet = Packet::try_from(packet_buff)?;
            return Err(Error::new(
                ErrorKind::SendingToAddressFailed,
                Some(format!(
                    "Failed to send packet to address.\n{}\n{}",
                    packet.to_ron_pretty(None)?,
                    e,
                )),
            ));
        }

        Ok(())
    }

    /// Sends `self` to connected [UdpSocket].
    ///
    /// This takes ownership of `self`.
    pub fn send_to_connected(self, socket: UdpSocket) -> Result<(), Error> {
        let packet_bytes: Bytes = self.into();
        let packet_buff = packet_bytes.as_slice();

        if let Err(e) = socket.send(packet_buff) {
            let packet = Packet::try_from(packet_buff)?;
            return Err(Error::new(
                ErrorKind::SendingToAddressFailed,
                Some(format!(
                    "Failed to send packet to connected socket.\n{}\n{}",
                    packet.to_ron_pretty(None)?,
                    e,
                )),
            ));
        }

        Ok(())
    }

    /// Receives a [Packet] from [`stream`][TcpStream].
    pub fn receive(stream: &mut TcpStream) -> Result<Packet, Error> {
        // Reads the size of packet.
        let mut size_buff = vec![0_u8; 2];
        if let Err(e) = stream.read_exact(&mut size_buff) {
            return Err(Error::new(
                ErrorKind::ReadingFromStreamFailed,
                Some(format!("Failed to read the size of packet. \n({})", e)),
            ));
        }
        let size = u16::try_from(&Bytes::from(size_buff.clone()))?;

        // Reads rest of packet.
        // - 2 for size of packet encoded as bytes which already exist.
        let mut buff = vec![0_u8; (size - 2) as usize];
        // USE READ EXACT, WITH OTHER METHODS THIS DOES NOT WORK.
        if let Err(e) = stream.read_exact(&mut buff) {
            return Err(Error::new(
                ErrorKind::ReadingFromStreamFailed,
                Some(format!("Failed to read contents of packet. \n({})", e)),
            ));
        }

        // Connect whole buffer and change name, so it makes more sense.
        size_buff.extend(buff);
        let buff = size_buff;

        // Create and return a packet from buffer.
        Packet::try_from(buff.as_slice())
    }

    /// Receives a [Packet] and a [SocketAddr] on a [UdpSocket].
    pub fn receive_from(socket: UdpSocket) -> Result<(Packet, SocketAddr), Error> {
        let mut buff = vec![0_u8; Packet::max_size().into()];

        match socket.recv_from(&mut buff) {
            Ok((number_of_byte_read, address)) => {
                // Truncate the buffer to remove invalid data.
                buff.truncate(number_of_byte_read);
                Ok((Packet::try_from(buff.as_slice())?, address))
            }
            Err(e) => Err(Error::new(
                ErrorKind::ReceivingFromAddressFailed,
                Some(format!(
                    "Failed to receive a packet from an address. \n({e})"
                )),
            )),
        }
    }

    /// Receives a [Packet] on a connected [UdpSocket].
    pub fn receive_from_connected(socket: UdpSocket) -> Result<Packet, Error> {
        let mut buff = vec![0_u8; Packet::max_size().into()];

        match socket.recv(&mut buff) {
            Ok(number_of_byte_read) => {
                // Truncate the buffer to remove invalid data.
                buff.truncate(number_of_byte_read);
                Ok(Packet::try_from(buff.as_slice())?)
            }
            Err(e) => Err(Error::new(
                ErrorKind::ReceivingFromAddressFailed,
                Some(format!(
                    "Failed to receive a packet from the connected socket. \n({e})"
                )),
            )),
        }
    }

    /// Peeks on data on [TcpStream] if those data are valid, [Packet] is created.
    ///
    /// Since `TcpStream::peek` is used to retrieve data, calling this again returns the same data.
    pub fn peek(stream: &mut TcpStream) -> Result<Packet, Error> {
        // Peeks on the size of packet.
        let mut size_buff = vec![0_u8; 2];
        if let Err(e) = stream.peek(&mut size_buff) {
            return Err(Error::new(
                ErrorKind::ReadingFromStreamFailed,
                Some(format!("Failed to peek on the size of packet. \n({})", e)),
            ));
        }

        let size = u16::try_from(&Bytes::from(size_buff))?;

        // Peek on rest of packet.
        // Size is peeked on again as for getting it was used peek as well.
        let mut buff = vec![0_u8; (size) as usize];
        if let Err(e) = stream.peek(&mut buff) {
            return Err(Error::new(
                ErrorKind::ReadingFromStreamFailed,
                Some(format!(
                    "Failed to peek on the contents of packet. \n({})",
                    e
                )),
            ));
        }

        // Create and return a packet from buffer.
        Packet::try_from(buff.as_slice())
    }

    /// Peeks on data on `socket`, if valid data are received a [Packet] is created and
    /// returned together with a [SocketAddr].
    ///
    /// Since `UdpSocket::peek_from` is used to retrieve data,
    /// calling this again returns the same data.
    pub fn peek_from(socket: UdpSocket) -> Result<(Self, SocketAddr), Error> {
        let mut buff = vec![0_u8; Packet::max_size().into()];

        match socket.peek_from(&mut buff) {
            Ok((number_of_byte_read, address)) => {
                buff.truncate(number_of_byte_read);
                Ok((Packet::try_from(buff.as_slice())?, address))
            }
            Err(e) => Err(Error::new(
                ErrorKind::ReceivingFromAddressFailed,
                Some(format!(
                    "Failed to receive a packet from the address. \n({e})"
                )),
            )),
        }
    }

    /// Peeks on data on a connected `socket`, if valid data are received a [Packet] is created.
    ///
    /// Since `UdpSocket::peek` is used to retrieve data, calling this again returns the same data.
    pub fn peek_from_connected(socket: UdpSocket) -> Result<Self, Error> {
        let mut buff = vec![0_u8; Packet::max_size().into()];

        match socket.peek(&mut buff) {
            Ok(number_of_byte_read) => {
                buff.truncate(number_of_byte_read);
                Ok(Packet::try_from(buff.as_slice())?)
            }
            Err(e) => Err(Error::new(
                ErrorKind::ReceivingFromAddressFailed,
                Some(format!(
                    "Failed to receive a packet from the connected socket. \n({e})"
                )),
            )),
        }
    }

    // MISSING EXAMPLE

    /// Returns number of packets needed to send data with given byte `length`.
    pub fn number_of_packets(length: usize) -> u32 {
        // Get number of packets by dividing by MAX_PACKET_CONTENT_SIZE.
        let mut number_of_packets = length / (Packet::max_content_size() as usize);
        // Add one packet if there is any remainder after the division.
        if length % (Packet::max_content_size() as usize) != 0 {
            number_of_packets += 1;
        }
        number_of_packets as u32
    }

    // MISSING EXAMPLE

    /// Takes [Bytes] and returns [Vec] of [Bytes] each with maximum length of
    /// [`MAX_CONTENT_SIZE`](Packet::max_content_size).
    pub fn split_to_max_packet_size(buffer: Bytes) -> Vec<Bytes> {
        // This splits given buffer to multiple owned chunks with chunks method
        // from itertools crate, then it will split every chunk to iterator as well
        // which are then collected to vectors of bytes, that are collected to single vector.
        // From: https://stackoverflow.com/a/67009164.
        let vectored_content: Vec<Bytes> = buffer
            .into_iter()
            .chunks(Packet::max_content_size() as usize)
            .into_iter()
            .map(|chunk| chunk.collect::<Vec<u8>>().into())
            .collect::<Vec<Bytes>>();

        vectored_content
    }

    /// Returns `size`.
    pub fn size(&self) -> u16 {
        self.size
    }

    /// Returns `kind`.
    pub fn kind(&self) -> PacketKind {
        self.kind.clone()
    }

    /// Returns `content`.
    ///
    /// Content is cloned.
    pub fn content(&self) -> Bytes {
        self.content.clone()
    }

    /// Returns a reference to `content`.
    pub fn content_ref(&self) -> &Bytes {
        &self.content
    }

    /// Returns a mutable reference to `content`.
    pub fn content_mut(&mut self) -> &mut Bytes {
        &mut self.content
    }

    /// Takes ownership of `self` and returns `content`.
    pub fn content_move(self) -> Bytes {
        self.content
    }

    /// Returns maximum number of bytes that can one [Packet] hold as [u16].
    /// This is global value for every [Packet] that can be changed with [Packet::set_max_size].
    ///
    /// Minimum value is 5 bytes, 2 for packet `size`([u16]), 2 for [`packet kind`](PacketKind),
    /// and at least 1 for [`content`](Bytes).
    ///
    /// This accesses a `mutable` static with [AtomicU16] to provide global access to this data
    /// and an ability to change its value for specific needs, but since [AtomicU16] is used
    /// it is capped by [u16::MAX].
    pub fn max_size() -> u16 {
        unsafe { MAX_PACKET_SIZE.load(Ordering::SeqCst) }
    }

    /// Sets [`MAX_PACKET_SIZE`](Packet::max_size) to `size`.
    ///
    /// Minimum value is 5 bytes, 2 for packet `size`([u16]), 2 for [`packet kind`](PacketKind),
    /// and at least 1 for [`content`](Bytes).
    ///
    /// For more information look at [Packet::max_size].
    pub fn set_max_size(size: u16) {
        unsafe {
            MAX_PACKET_SIZE.store(size, Ordering::SeqCst);
        }
    }

    /// Returns minimal number of bytes that is every packet guaranteed to have,
    /// 2 bytes are for its size and 2 for its kind as [u16].
    pub const fn description_size() -> u16 {
        PACKET_DESCRIPTION_SIZE
    }

    /// Maximum amount of bytes that a [Packet] can use for its content, its lower than
    /// [`MAX_PACKET_SIZE`](Packet::max_size) by
    /// [`PACKET_DESCRIPTION_SIZE`](Packet::description_size).
    pub fn max_content_size() -> u16 {
        Packet::max_size() - Packet::description_size()
    }
}