1 files changed, 68 insertions, 117 deletions

diff --git a/azalea-protocol/src/connect.rs b/azalea-protocol/src/connect.rs
index 67771d8e..cd51f86c 100755
--- a/azalea-protocol/src/connect.rs
+++ b/azalea-protocol/src/connect.rs
@@ -1,53 +1,56 @@
 //! parse sending and receiving packets with a server.
 
-use crate::packets::game::GamePacket;
-use crate::packets::handshake::HandshakePacket;
-use crate::packets::login::LoginPacket;
-use crate::packets::status::StatusPacket;
+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;
 
-#[derive(Debug, Clone, Copy)]
-pub enum PacketFlow {
-    ClientToServer,
-    ServerToClient,
-}
-
-pub struct HandshakeConnection {
-    pub flow: PacketFlow,
-    /// The buffered writer
-    pub stream: TcpStream,
-}
-
-pub struct GameConnection {
-    pub flow: PacketFlow,
+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 {
-    pub flow: PacketFlow,
-    /// 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 {
-    pub flow: PacketFlow,
-    /// 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;
 
@@ -60,95 +63,26 @@ impl HandshakeConnection {
             .set_nodelay(true)
             .expect("Error enabling tcp_nodelay");
 
-        Ok(HandshakeConnection {
-            flow: PacketFlow::ServerToClient,
+        Ok(Connection {
             stream,
-        })
-    }
-
-    pub fn login(self) -> LoginConnection {
-        LoginConnection {
-            flow: self.flow,
-            stream: self.stream,
             compression_threshold: None,
             enc_cipher: None,
             dec_cipher: None,
-        }
-    }
-
-    pub fn status(self) -> StatusConnection {
-        StatusConnection {
-            flow: self.flow,
-            stream: self.stream,
-        }
-    }
-
-    pub async fn read(&mut self) -> Result<HandshakePacket, String> {
-        read_packet::<HandshakePacket, _>(&self.flow, &mut self.stream, None, &mut None).await
-    }
-
-    /// Write a packet to the server
-    pub async fn write(&mut self, packet: HandshakePacket) {
-        write_packet(packet, &mut self.stream, None, &mut None).await;
-    }
-}
-
-impl GameConnection {
-    pub async fn read(&mut self) -> Result<GamePacket, String> {
-        read_packet::<GamePacket, _>(
-            &self.flow,
-            &mut self.stream,
-            self.compression_threshold,
-            &mut self.dec_cipher,
-        )
-        .await
-    }
-
-    /// Write a packet to the server
-    pub async fn write(&mut self, packet: GamePacket) {
-        write_packet(
-            packet,
-            &mut self.stream,
-            self.compression_threshold,
-            &mut self.enc_cipher,
-        )
-        .await;
+            _reading: PhantomData,
+            _writing: PhantomData,
+        })
     }
-}
 
-impl StatusConnection {
-    pub async fn read(&mut self) -> Result<StatusPacket, String> {
-        read_packet::<StatusPacket, _>(&self.flow, &mut self.stream, None, &mut None).await
+    pub fn login(self) -> Connection<ClientboundLoginPacket, ServerboundLoginPacket> {
+        Connection::from(self)
     }
 
-    /// Write a packet to the server
-    pub async fn write(&mut self, packet: StatusPacket) {
-        write_packet(packet, &mut self.stream, None, &mut None).await;
+    pub fn status(self) -> Connection<ClientboundStatusPacket, ServerboundStatusPacket> {
+        Connection::from(self)
     }
 }
 
-impl LoginConnection {
-    pub async fn read(&mut self) -> Result<LoginPacket, String> {
-        read_packet::<LoginPacket, _>(
-            &self.flow,
-            &mut self.stream,
-            self.compression_threshold,
-            &mut self.dec_cipher,
-        )
-        .await
-    }
-
-    /// Write a packet to the server
-    pub async fn write(&mut self, packet: LoginPacket) {
-        write_packet(
-            packet,
-            &mut self.stream,
-            self.compression_threshold,
-            &mut self.enc_cipher,
-        )
-        .await;
-    }
-
+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 {
@@ -165,13 +99,30 @@ impl LoginConnection {
         self.dec_cipher = Some(dec_cipher);
     }
 
-    pub fn game(self) -> GameConnection {
-        GameConnection {
-            flow: self.flow,
-            stream: self.stream,
-            compression_threshold: self.compression_threshold,
-            enc_cipher: self.enc_cipher,
-            dec_cipher: self.dec_cipher,
+    pub fn game(self) -> Connection<ClientboundGamePacket, ServerboundGamePacket> {
+        Connection::from(self)
+    }
+}
+
+// rust doesn't let us implement From because allegedly it conflicts with
+// `core`'s "impl<T> From<T> for T" so we do this instead
+impl<R1, W1> Connection<R1, W1>
+where
+    R1: ProtocolPacket + Debug,
+    W1: ProtocolPacket + Debug,
+{
+    fn from<R2, W2>(connection: Connection<R1, W1>) -> Connection<R2, W2>
+    where
+        R2: ProtocolPacket + Debug,
+        W2: ProtocolPacket + Debug,
+    {
+        Connection {
+            stream: connection.stream,
+            compression_threshold: connection.compression_threshold,
+            enc_cipher: connection.enc_cipher,
+            dec_cipher: connection.dec_cipher,
+            _reading: PhantomData,
+            _writing: PhantomData,
         }
     }
 }