make Connection a single struct with generics

This isn't that good of a solution but I couldn't come up with a better one and this was pretty simple to implement

1 files changed, 63 insertions, 87 deletions

diff --git a/azalea-protocol/src/connect.rs b/azalea-protocol/src/connect.rs
index bf730fc4..a438ccfb 100755
--- a/azalea-protocol/src/connect.rs
+++ b/azalea-protocol/src/connect.rs
@@ -4,40 +4,58 @@ use crate::packets::game::{ClientboundGamePacket, ServerboundGamePacket};
 use crate::packets::handshake::{ClientboundHandshakePacket, ServerboundHandshakePacket};
 use crate::packets::login::{ClientboundLoginPacket, ServerboundLoginPacket};
 use crate::packets::status::{ClientboundStatusPacket, ServerboundStatusPacket};
+use crate::packets::ProtocolPacket;
 use crate::read::read_packet;
 use crate::write::write_packet;
 use crate::ServerIpAddress;
 use azalea_crypto::{Aes128CfbDec, Aes128CfbEnc};
+use std::fmt::Debug;
+use std::marker::PhantomData;
 use tokio::net::TcpStream;
 
-pub struct HandshakeConnection {
-    /// The buffered writer
-    pub stream: TcpStream,
-}
+pub struct Handshake;
+pub struct Game;
+pub struct Status;
+pub struct Login;
 
-pub struct GameConnection {
+pub struct Connection<R: ProtocolPacket, W: ProtocolPacket> {
     /// The buffered writer
     pub stream: TcpStream,
     pub compression_threshold: Option<u32>,
     pub enc_cipher: Option<Aes128CfbEnc>,
     pub dec_cipher: Option<Aes128CfbDec>,
+    _reading: PhantomData<R>,
+    _writing: PhantomData<W>,
 }
 
-pub struct StatusConnection {
-    /// The buffered writer
-    pub stream: TcpStream,
-}
+impl<R, W> Connection<R, W>
+where
+    R: ProtocolPacket + Debug,
+    W: ProtocolPacket + Debug,
+{
+    pub async fn read(&mut self) -> Result<R, String> {
+        read_packet::<R, _>(
+            &mut self.stream,
+            self.compression_threshold,
+            &mut self.dec_cipher,
+        )
+        .await
+    }
 
-pub struct LoginConnection {
-    /// The buffered writer
-    pub stream: TcpStream,
-    pub compression_threshold: Option<u32>,
-    pub enc_cipher: Option<Aes128CfbEnc>,
-    pub dec_cipher: Option<Aes128CfbDec>,
+    /// Write a packet to the server
+    pub async fn write(&mut self, packet: W) {
+        write_packet(
+            packet,
+            &mut self.stream,
+            self.compression_threshold,
+            &mut self.enc_cipher,
+        )
+        .await;
+    }
 }
 
-impl HandshakeConnection {
-    pub async fn new(address: &ServerIpAddress) -> Result<HandshakeConnection, String> {
+impl Connection<ClientboundHandshakePacket, ServerboundHandshakePacket> {
+    pub async fn new(address: &ServerIpAddress) -> Result<Self, String> {
         let ip = address.ip;
         let port = address.port;
 
@@ -50,88 +68,44 @@ impl HandshakeConnection {
             .set_nodelay(true)
             .expect("Error enabling tcp_nodelay");
 
-        Ok(HandshakeConnection { stream })
-    }
-
-    pub fn login(self) -> LoginConnection {
-        LoginConnection {
-            stream: self.stream,
+        Ok(Connection {
+            stream,
             compression_threshold: None,
             enc_cipher: None,
             dec_cipher: None,
-        }
+            _reading: PhantomData,
+            _writing: PhantomData,
+        })
     }
 
-    pub fn status(self) -> StatusConnection {
-        StatusConnection {
+    pub fn login(self) -> Connection<ClientboundLoginPacket, ServerboundLoginPacket> {
+        Connection {
             stream: self.stream,
+            compression_threshold: self.compression_threshold,
+            enc_cipher: self.enc_cipher,
+            dec_cipher: self.dec_cipher,
+            _reading: PhantomData,
+            _writing: PhantomData,
         }
     }
 
-    pub async fn read(&mut self) -> Result<ClientboundHandshakePacket, String> {
-        read_packet::<ClientboundHandshakePacket, _>(&mut self.stream, None, &mut None).await
-    }
-
-    /// Write a packet to the server
-    pub async fn write(&mut self, packet: ServerboundHandshakePacket) {
-        write_packet(packet, &mut self.stream, None, &mut None).await;
-    }
-}
-
-impl GameConnection {
-    pub async fn read(&mut self) -> Result<ClientboundGamePacket, String> {
-        read_packet::<ClientboundGamePacket, _>(
-            &mut self.stream,
-            self.compression_threshold,
-            &mut self.dec_cipher,
-        )
-        .await
-    }
-
-    /// Write a packet to the server
-    pub async fn write(&mut self, packet: ServerboundGamePacket) {
-        write_packet(
-            packet,
-            &mut self.stream,
-            self.compression_threshold,
-            &mut self.enc_cipher,
-        )
-        .await;
-    }
-}
-
-impl StatusConnection {
-    pub async fn read(&mut self) -> Result<ClientboundStatusPacket, String> {
-        read_packet::<ClientboundStatusPacket, _>(&mut self.stream, None, &mut None).await
-    }
-
-    /// Write a packet to the server
-    pub async fn write(&mut self, packet: ServerboundStatusPacket) {
-        write_packet(packet, &mut self.stream, None, &mut None).await;
+    pub fn status(self) -> Connection<ClientboundStatusPacket, ServerboundStatusPacket> {
+        Connection {
+            stream: self.stream,
+            compression_threshold: self.compression_threshold,
+            enc_cipher: self.enc_cipher,
+            dec_cipher: self.dec_cipher,
+            _reading: PhantomData,
+            _writing: PhantomData,
+        }
     }
 }
 
-impl LoginConnection {
-    pub async fn read(&mut self) -> Result<ClientboundLoginPacket, String> {
-        read_packet::<ClientboundLoginPacket, _>(
-            &mut self.stream,
-            self.compression_threshold,
-            &mut self.dec_cipher,
-        )
-        .await
-    }
+impl Connection<ClientboundGamePacket, ServerboundGamePacket> {}
 
-    /// Write a packet to the server
-    pub async fn write(&mut self, packet: ServerboundLoginPacket) {
-        write_packet(
-            packet,
-            &mut self.stream,
-            self.compression_threshold,
-            &mut self.enc_cipher,
-        )
-        .await;
-    }
+impl Connection<ClientboundStatusPacket, ServerboundStatusPacket> {}
 
+impl Connection<ClientboundLoginPacket, ServerboundLoginPacket> {
     pub fn set_compression_threshold(&mut self, threshold: i32) {
         // if you pass a threshold of 0 or less, compression is disabled
         if threshold > 0 {
@@ -148,12 +122,14 @@ impl LoginConnection {
         self.dec_cipher = Some(dec_cipher);
     }
 
-    pub fn game(self) -> GameConnection {
-        GameConnection {
+    pub fn game(self) -> Connection<ClientboundGamePacket, ServerboundGamePacket> {
+        Connection {
             stream: self.stream,
             compression_threshold: self.compression_threshold,
             enc_cipher: self.enc_cipher,
             dec_cipher: self.dec_cipher,
+            _reading: PhantomData,
+            _writing: PhantomData,
         }
     }
 }