Compression and encryption

Cargo.toml

@@ -12,4 +12,7 @@ serde = { version = "1.0", features = ["derive", "std"] }
 serde_json = "1.0"
 bytes = "1.1"
 uuid = "0.8"
+cfb8 = "0.7"
+aes = "0.7"
+miniz_oxide = "0.5"
 lixcraft-derive = { path = "lixcraft-derive" }

src/frames.rs

@@ -1,10 +1,55 @@
-use std::io::ErrorKind;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use aes::Aes128;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
-use tokio::io::Error;
+use cfb8::cipher::AsyncStreamCipher;
+use cfb8::Cfb8;
+use miniz_oxide::deflate::{compress_to_vec_zlib, CompressionLevel};
+use miniz_oxide::inflate::decompress_to_vec_zlib_with_limit;
+use tokio::io::{AsyncRead, AsyncWrite, Error, ReadBuf};
 use tokio_util::codec::{Decoder, Encoder};
-use crate::VarInt;
+use crate::{JavaSerializable, VarInt};
+use crate::serde::JavaDeserError;
 
-pub struct LengthPrefixedFrame {
+
+pub struct JavaEncryptor<W: AsyncWrite> {
+    inner: W,
+    cipher: Cfb8<Aes128>
+}
+impl<W: AsyncWrite + Unpin> AsyncWrite for JavaEncryptor<W> {
+    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Error>> {
+        let mut out =  buf.to_vec();
+        self.cipher.encrypt(&mut out[..]);
+        Pin::new(&mut (self.inner)).poll_write(cx, &out[..])
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        Pin::new(&mut self.inner).poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        Pin::new(&mut self.inner).poll_shutdown(cx)
+    }
+}
+
+pub struct JavaDecryptor<R: AsyncRead> {
+    inner: R,
+    cipher: Cfb8<Aes128>
+}
+impl<R: AsyncRead + Unpin> AsyncRead for JavaDecryptor<R> {
+    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
+        let cur = buf.filled().len();
+        match Pin::new(&mut self.inner).poll_read(cx, buf) {
+            Poll::Ready(Ok(())) => {
+                self.cipher.decrypt(&mut buf.filled_mut()[cur..]);
+                Poll::Ready(Ok(()))
+            },
+            v => v
+        }
+    }
+}
+
+pub struct LengthPrefixedCodec {
     pub max_length: usize
 }
 
@@ -13,7 +58,7 @@ enum VarIntReadError {
     InvalidData
 }
 
-impl LengthPrefixedFrame {
+impl LengthPrefixedCodec {
     fn try_read_var_int(&self, src: &mut BytesMut) -> Result<(VarInt, usize), VarIntReadError> {
         let src = src.chunk();
         let mut num_read = 0;
@@ -42,9 +87,9 @@ impl LengthPrefixedFrame {
     }
 }
 
-impl Decoder for LengthPrefixedFrame {
-    type Item = BytesMut;
-    type Error = tokio::io::Error;
+impl Decoder for LengthPrefixedCodec {
+    type Item = Bytes;
+    type Error = JavaDeserError;
 
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
         println!("!");
@@ -54,34 +99,72 @@ impl Decoder for LengthPrefixedFrame {
 
                 if src.len() >= total_length {
                     src.advance(prefix_length);
-                    Ok(Some(src.split_to(usize::from(length))))
+                    Ok(Some(src.split_to(usize::from(length)).freeze()))
                 } else {
                     Ok(None)
                 }
             },
             Err(VarIntReadError::BytesMissing) => Ok(None),
-            Err(VarIntReadError::InvalidData) => Err(Error::from(ErrorKind::InvalidData))
+            Err(VarIntReadError::InvalidData) => Err(JavaDeserError::InvalidValue)
         }
     }
 }
-impl Encoder<Bytes> for LengthPrefixedFrame {
-    type Error = tokio::io::Error;
+impl Encoder<Bytes> for LengthPrefixedCodec {
+    type Error = JavaDeserError;
 
     fn encode(&mut self, data: Bytes, dst: &mut BytesMut) -> Result<(), Self::Error> {
-        dst.reserve(5 + data.len());
-
-        let mut v = data.len();
-        loop {
-            if v & 0x80 == 0 {
-                dst.put_u8(v as u8);
-                break;
-            }
-
-            dst.put_u8((v as u8 & 0x7f) | 0x80);
-            v >>= 7;
-        }
+        let len = VarInt::from(data.len());
 
+        dst.reserve(len.size_in_bytes() + data.len());
+        len.serialize(dst);
         dst.put_slice(&data);
         Ok(())
     }
+}
+
+pub struct CompressedCodec {
+    pub length_prefixed: LengthPrefixedCodec,
+    pub threshold: usize
+}
+
+impl Encoder<Bytes> for CompressedCodec {
+    type Error = JavaDeserError;
+
+    fn encode(&mut self, data: Bytes, dst: &mut BytesMut) -> Result<(), Self::Error> {
+        let (inner_len, data) = if data.len() < self.threshold {
+            (VarInt(0), data)
+        } else {
+            (VarInt::from(data.len()), Bytes::from(compress_to_vec_zlib(&data[..], CompressionLevel::BestSpeed as u8)))
+        };
+
+        let total_len = VarInt::from(inner_len.size_in_bytes() + data.len());
+        dst.reserve(total_len.size_in_bytes() + inner_len.size_in_bytes() + data.len());
+        total_len.serialize(dst);
+        inner_len.serialize(dst);
+        dst.put_slice(&data[..]);
+        Ok(())
+    }
+}
+impl Decoder for CompressedCodec {
+    type Item = Bytes;
+    type Error = JavaDeserError;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        match self.length_prefixed.decode(src) {
+            Ok(Some(mut buf)) => {
+                let inner_len = VarInt::deserialize(&mut buf)?;
+                if *inner_len == 0 {
+                    Ok(Some(buf))
+                } else {
+                    let decompressed = decompress_to_vec_zlib_with_limit(&buf[..], *inner_len as usize).map_err(|_| JavaDeserError::InvalidValue)?;
+                    if *inner_len as usize != decompressed.len() {
+                        Err(JavaDeserError::InvalidValue)
+                    } else {
+                        Ok(Some(Bytes::from(decompressed)))
+                    }
+                }
+            },
+            v => v
+        }
+    }
 }

src/main.rs

@@ -3,16 +3,15 @@ extern crate core;
 pub mod serde;
 pub mod packets;
 pub mod frames;
+pub mod protocol;
 
-use std::collections::HashMap;
 use std::error::Error;
 use std::net::SocketAddr;
 use std::sync::Arc;
-use bytes::BytesMut;
 use tokio::net::{TcpListener, TcpStream};
 use tokio_stream::StreamExt;
 use tokio_util::codec::Framed;
-use crate::frames::LengthPrefixedFrame;
+use crate::frames::LengthPrefixedCodec;
 use crate::packets::handshake::HandshakePacket;
 use crate::serde::{JavaSerializable, VarInt};
 
@@ -59,29 +58,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
-enum Reason {
-    Disconnected,
-    ConnectionDrop,
-    ProtocolError
-}
-
-struct Protocol {
-    packet_handlers: HashMap<VarInt, Box<dyn FnMut(BytesMut) -> Reason>>
-}
-
-fn handle_handshake(buf: BytesMut) -> Result<Framed<TcpStream, LengthPrefixedFrame>, Reason> {
-    Err(Reason::Disconnected)
-}
-
 async fn process_initial(stream: TcpStream, _addr: SocketAddr) {
-    /*let handshake_protocol = Protocol {
-        packet_handlers: HashMap::from([(VarInt(0), Box::new(handle_handshake) as Box<dyn FnMut(BytesMut) -> Reason>)]),
-    };*/
-
-    let mut framed = Framed::new(stream, LengthPrefixedFrame {max_length: 263});
+    let mut framed = Framed::new(stream, LengthPrefixedCodec {max_length: 263});
     match framed.next().await {
-        Some(Ok(buf)) => {
-            let mut buf = buf.freeze();
+        Some(Ok(mut buf)) => {
             match VarInt::deserialize(&mut buf) {
                 Ok(VarInt(0)) => {
                     match HandshakePacket::deserialize(&mut buf) {

src/protocol.rs

@@ -0,0 +1,74 @@
+use std::collections::HashMap;
+use bytes::Bytes;
+use crate::{HandshakePacket, VarInt};
+
+pub const CURRENT_VERSION: VarInt = VarInt(758);
+
+pub fn version_to_string(version: VarInt) -> Option<&'static str> {
+    Some(match *version {
+        758 => "1.18.2",
+        757 => "1.18.1",
+        756 => "1.17.1",
+        755 => "1.17",
+        754 => "1.16.5",
+        753 => "1.16.3",
+        751 => "1.16.2",
+        736 => "1.16.1",
+        735 => "1.16",
+        578 => "1.15.2",
+        575 => "1.15.1",
+        573 => "1.15",
+        498 => "1.14.4",
+        490 => "1.14.3",
+        485 => "1.14.2",
+        480 => "1.14.1",
+        477 => "1.14",
+        404 => "1.13.2",
+        401 => "1.13.1",
+        393 => "1.13",
+        340 => "1.12.2",
+        338 => "1.12.1",
+        335 => "1.12",
+        316 => "1.11.2",
+        315 => "1.11",
+        210 => "1.10.2",
+        110 => "1.9.4",
+        109 => "1.9.2",
+        108 => "1.9.1",
+        107 => "1.9",
+        47 => "1.8.9",
+        5 => "1.7.10",
+        4 => "1.7.5",
+        _ => return None
+    })
+}
+
+pub enum Todo {
+    Continue,
+    Disconnect
+}
+
+pub trait Protocol {
+    fn handle_buf(&self, buf: Bytes) -> Option<Box<dyn Protocol>>;
+}
+pub struct HandshakeProtocol {}
+impl Protocol for HandshakeProtocol {
+    fn handle_buf(&self, buf: Bytes) -> Option<Box<dyn Protocol>> {
+        todo!()
+    }
+}
+
+pub struct GenericProtocol {
+    packet_handlers: HashMap<VarInt, Box<dyn FnMut(Bytes) -> Todo>>
+}
+
+fn handle_handshake(handshake: HandshakePacket) -> Todo {
+    todo!()
+}
+/*
+pub fn handshake_protocol() -> Protocol {
+    Protocol {
+        packet_handlers: HashMap::from([(VarInt(0), Box::new(handle_handshake) as Box<dyn FnMut(Bytes) -> Todo>)]),
+    }
+}
+*/

src/serde.rs

@@ -1,6 +1,7 @@
+use std::io::Error;
 use std::ops::{DerefMut, Deref};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
-use crate::serde::JavaDeserError::{InvalidValue, MissingBytes};
+use crate::serde::JavaDeserError::{InvalidValue, MissingBytes, OtherError};
 
 pub use lixcraft_derive::JavaSerializable;
 use uuid::Uuid;
@@ -8,7 +9,14 @@ use uuid::Uuid;
 #[derive(Debug)]
 pub enum JavaDeserError {
     InvalidValue,
-    MissingBytes
+    MissingBytes,
+    OtherError
+}
+
+impl From<std::io::Error> for JavaDeserError {
+    fn from(_: Error) -> Self {
+        OtherError
+    }
 }
 
 pub trait JavaSerializable: Sized {