rust_mc_proto/src/lib.rs

#[cfg(test)]
mod tests;

pub mod data;
pub mod packet;

pub mod prelude {
    pub use crate::{DataReader, DataWriter};
}

pub use crate::{
    data::{DataReader, DataWriter},
    packet::Packet,
};

use flate2::{read::ZlibDecoder, write::ZlibEncoder, Compression};
use std::{
    error::Error, fmt, io::{Read, Write}, net::{TcpStream, ToSocketAddrs}, usize
};

/// Minecraft protocol error
#[derive(Debug)]
pub enum ProtocolError {
    AddressParseError,
    DataRanOutError,
    StringParseError,
    StreamConnectError,
    VarIntError,
    ReadError,
    WriteError,
    ZlibError,
    CloneError,
    ConnectionClosedError
}

impl fmt::Display for ProtocolError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "An protocol error occured")
    }
}

impl Error for ProtocolError {}

/// Minecraft connection, wrapper for stream with compression
pub struct MinecraftConnection<T: Read + Write> {
    stream: T,
    compression: Option<usize>,
    compression_type: u32
}

impl<T: Read + Write> DataReader for MinecraftConnection<T> {
    fn read_bytes(&mut self, size: usize) -> Result<Vec<u8>, ProtocolError> {
        let mut buf = vec![0; size];
        match self.stream.read_exact(&mut buf) {
            Ok(_) => Ok(buf),
            Err(_) => Err(ProtocolError::ReadError),
        }
    }
}

impl<T: Read + Write> DataWriter for MinecraftConnection<T> {
    fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), ProtocolError> {
        self.stream.write_all(bytes).map_err(|_| ProtocolError::WriteError)
    }
}

impl<T: Read + Write> MinecraftConnection<T> {
    /// Create new MinecraftConnection from stream
    pub fn new(stream: T) -> MinecraftConnection<T> {
        MinecraftConnection {
            stream,
            compression: None,
            compression_type: 1,
        }
    }

    /// Set compression threshold
    pub fn set_compression(&mut self, threshold: Option<usize>) {
        self.compression = threshold;
    }

    /// Get compression threshold
    pub fn compression(&self) -> Option<usize> {
        self.compression
    }

    /// Set compression type
    ///
    /// `compression_type` is integer from 0 (none) to 9 (longest)
    /// 1 is fast compression
    /// 6 is normal compression
    pub fn set_compression_type(&mut self, compression_type: u32) {
        self.compression_type = compression_type;
    }

    /// Get compression type
    ///
    /// `compression_type` is integer from 0 (none) to 9 (longest)
    /// 1 is fast compression
    /// 6 is normal compression
    pub fn compression_type(&self) -> u32 {
        self.compression_type
    }

    /// Get mutable reference of stream
    pub fn get_mut(&mut self) -> &mut T {
        &mut self.stream
    }

    /// Get immutable reference of stream
    pub fn get_ref(&self) -> &T {
        &self.stream
    }

    /// Read [`Packet`](Packet) from connection
    pub fn read_packet(&mut self) -> Result<Packet, ProtocolError> {
        read_packet(&mut self.stream, self.compression)
    }

    /// Write [`Packet`](Packet) to connection
    pub fn write_packet(&mut self, packet: &Packet) -> Result<(), ProtocolError> {
        write_packet(&mut self.stream, self.compression, self.compression_type, packet)
    }
}

impl<T: Read + Write + Clone> MinecraftConnection<T> {
    /// Clone MinecraftConnection with compression and stream
    pub fn clone(&mut self) -> MinecraftConnection<T> {
        MinecraftConnection {
            stream: self.stream.clone(),
            compression: self.compression,
            compression_type: self.compression_type,
        }
    }
}

impl MinecraftConnection<TcpStream> {
    /// Connect to Minecraft Server with TcpStream
    pub fn connect(addr: impl ToSocketAddrs) -> Result<MinecraftConnection<TcpStream>, ProtocolError> {
        let addr = match addr.to_socket_addrs() {
            Ok(mut i) => match i.next() {
                Some(i) => i,
                None => return Err(ProtocolError::AddressParseError),
            },
            Err(_) => return Err(ProtocolError::AddressParseError),
        };

        let stream: TcpStream = match TcpStream::connect(&addr) {
            Ok(i) => i,
            Err(_) => return Err(ProtocolError::StreamConnectError),
        };

        Ok(MinecraftConnection {
            stream,
            compression: None,
            compression_type: 1,
        })
    }

    /// Close TcpStream
    pub fn close(&mut self) {
        let _ = self.stream.shutdown(std::net::Shutdown::Both);
    }

    /// Try clone MinecraftConnection with compression and stream
    pub fn try_clone(&self) -> Result<MinecraftConnection<TcpStream>, ProtocolError> {
        match self.stream.try_clone() {
            Ok(stream) => Ok(MinecraftConnection {
                stream,
                compression: self.compression,
                compression_type: self.compression_type,
            }),
            _ => Err(ProtocolError::CloneError),
        }
    }
}

fn compress_zlib(bytes: &[u8], compression: u32) -> Result<Vec<u8>, ProtocolError> {
    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::new(compression));
    encoder.write_all(bytes).or(Err(ProtocolError::ZlibError))?;
    encoder.finish().or(Err(ProtocolError::ZlibError))
}

fn decompress_zlib(bytes: &[u8]) -> Result<Vec<u8>, ProtocolError> {
    let mut decoder = ZlibDecoder::new(bytes);
    let mut output = Vec::new();
    decoder
        .read_to_end(&mut output)
        .or(Err(ProtocolError::ZlibError))?;
    Ok(output)
}

/// MinecraftConnection shorter alias
pub type MCConn<T> = MinecraftConnection<T>;

/// MinecraftConnection\<TcpStream\> shorter alias
pub type MCConnTcp = MinecraftConnection<TcpStream>;


/// Read [`Packet`](Packet) from stream
///
/// `compression` here is atomic usize
/// usize::MAX means that compression is disabled
///
/// `ordering` is order how to load atomic
pub fn read_packet<T: Read>(
    stream: &mut T,
    compression: Option<usize>
) -> Result<Packet, ProtocolError> {
    let mut data: Vec<u8>;

    let packet_length = stream.read_usize_varint_size()?;

    if compression.is_some() {
        let data_length = stream.read_usize_varint_size()?;

        data = stream.read_bytes(packet_length.0 - data_length.1)?;

        if data_length.0 != 0 {
            data = decompress_zlib(&data)?;
        }
    } else {
        data = stream.read_bytes(packet_length.0 as usize)?;
    }

    Ok(Packet::from_data(&data)?)
}

/// Write [`Packet`](Packet) to stream
///
/// `compression` here is usize
/// usize::MAX means that compression is disabled
///
/// `ordering` is order how to load atomic
///
/// `compression_type` is integer from 0 (none) to 9 (longest)
/// 1 is fast compression
/// 6 is normal compression
pub fn write_packet<T: Write>(
    stream: &mut T,
    compression: Option<usize>,
    compression_type: u32,
    packet: &Packet,
) -> Result<(), ProtocolError> {
    let mut buf = Vec::new();

    let mut data_buf = Vec::new();
    data_buf.write_varint(packet.id() as i32)?;
    data_buf.write_bytes(packet.get_bytes())?;

    if let Some(compression) = compression {
        let mut packet_buf = Vec::new();

        if data_buf.len() >= compression {
            let compressed_data = compress_zlib(&data_buf, compression_type)?;
            packet_buf.write_varint(data_buf.len() as i32)?;
            packet_buf
                .write_all(&compressed_data)
                .or(Err(ProtocolError::WriteError))?;
        } else {
            packet_buf.write_varint(0)?;
            packet_buf.write_bytes(&data_buf)?;
        }

        buf.write_varint(packet_buf.len() as i32)?;
        buf.write_bytes(&packet_buf)?;
    } else {
        buf.write_varint(data_buf.len() as i32)?;
        buf.write_bytes(&data_buf)?;
    }

    stream.write_bytes(&buf)?;

    Ok(())
}