literustlib 0.4.0

use std::collections::HashMap;

use crate::packet::Packet;

#[derive(Debug, Clone)]
pub struct PacketPool<D: crate::packet::PacketData> {
    fragments: HashMap<u16, HashMap<u16, Packet>>, // (fragment id, (fragment part, packet))
    _d: std::marker::PhantomData<fn() -> D>,
}

impl <D: crate::packet::PacketData> PacketPool<D> {
    pub fn new() -> Self {
        Self {
            fragments: Default::default(),
            _d: Default::default(),
        }
    }

    pub fn handle_packet(&mut self, packet: Packet) -> std::io::Result<Option<D>> {
        if let Some(fragment) = packet.header.fragment {
            if !fragment.is_valid() {
                return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "Fragment data failed validation"));
            }
            let key = fragment.id;
            let fragment_count = fragment.total as usize;
            let header = packet.header;
            let frags_available = if let Some(frags) = self.fragments.get_mut(&key) {
                frags.insert(fragment.part, packet);
                frags.len()
            } else {
                let mut map = HashMap::with_capacity(fragment.total as _);
                map.insert(fragment.part, packet);
                self.fragments.insert(key, map);
                1
            };
            if frags_available == fragment_count {
                let mut all_fragments = self.fragments.remove(&key).unwrap();
                let total_size = all_fragments.values().map(|p| p.data.len()).sum();
                let mut buf = bytes::BytesMut::with_capacity(total_size);
                for i in 0..fragment_count as u16 {
                    buf.extend_from_slice(&all_fragments.remove(&i).unwrap().data);
                }
                Ok(Some(D::parse(buf.freeze(), &header)?))
            } else {
                Ok(None)
            }
        } else {
            Ok(Some(D::parse(packet.data, &packet.header)?))
        }
    }
}