Abstract out connection state with phantom protocol state markers.

src/main.rs

@@ -1,11 +1,10 @@
 #![feature(const_generics)]
+#![feature(never_type)]
 
 mod net;
 
-use crate::net::{Reader, Writer};
 use crate::net::chat::Chat;
-use crate::net::format::{PacketFormat, DefaultPacketFormat};
-use tokio::io::{BufReader, BufWriter};
+use crate::net::connection::{Connection, StateHandshake, StateStatus};
 use tokio::net::{TcpListener, TcpStream};
 use std::io;
 use std::net::IpAddr;
@@ -43,43 +42,38 @@ async fn listen(mut listener: TcpListener) {
     }
 }
 
-async fn accept_connection(mut socket: TcpStream) {
-    let (read, write) = socket.split();
-    let mut source = BufReader::new(read);
-    let mut dest = BufWriter::new(write);
+async fn accept_connection(socket: TcpStream) {
+    let con = Connection::new(socket);
 
     eprintln!("Client connected.");
-    match interact_handshake(&mut source, &mut dest).await {
+    match interact_handshake(con).await {
         Err(err) => { eprintln!("Client disconnected with error: {:?}", err); },
         Ok(_) => { eprintln!("Client disconnected without error."); }
     }
 }
 
-async fn interact_handshake<'a>(source: &mut Reader<'a>, dest: &mut Writer<'a>) -> io::Result<()> {
+async fn interact_handshake<'a>(mut con: Connection<StateHandshake>) -> io::Result<()> {
     use crate::net::packet::handshake::*;
 
-    let pkt = DefaultPacketFormat.recieve::<HandshakeServerbound>(source).await?;
-
-    match pkt {
+    match con.read().await? {
         HandshakeServerbound::Handshake(handshake) => {
-            if handshake.next_state == HandshakeNextState::Status {
-                interact_status(source, dest).await
-            } else {
-                Err(io::Error::new(io::ErrorKind::Other, "We do not support client log-in yet.".to_string()))
+            use HandshakeNextState::*;
+
+            match handshake.next_state {
+                Status => interact_status(con.into_status()).await,
+                Login => Err(io::Error::new(io::ErrorKind::Other, "We do not support client log-in yet.".to_string()))
             }
         }
     }
 }
 
-async fn interact_status<'a>(source: &mut Reader<'a>, dest: &mut Writer<'a>) -> io::Result<()> {
+async fn interact_status<'a>(mut con: Connection<StateStatus>) -> io::Result<()> {
     use crate::net::packet::status::*;
 
     loop {
-        let pkt = DefaultPacketFormat.recieve::<StatusServerbound>(source).await?;
-
-        match pkt {
+        match con.read().await? {
             StatusServerbound::Request(Request {}) => {
-                DefaultPacketFormat.send(dest, &StatusClientbound::Response(Response {
+                con.write(&StatusClientbound::Response(Response {
                     data: ResponseData {
                         version: ResponseVersion {
                             name: "1.16.1".to_string(),
@@ -96,7 +90,7 @@ async fn interact_status<'a>(source: &mut Reader<'a>, dest: &mut Writer<'a>) ->
                 })).await?;
             },
             StatusServerbound::Ping(ping) => {
-                DefaultPacketFormat.send(dest, &StatusClientbound::Pong(Pong {
+                con.write(&StatusClientbound::Pong(Pong {
                     payload: ping.payload
                 })).await?;
 

src/net/connection.rs

@@ -0,0 +1,117 @@
+use crate::net::{Reader, Writer};
+use crate::net::format::{PacketFormat, DefaultPacketFormat};
+use crate::net::state::ProtocolState;
+use crate::net::packet::handshake::{HandshakeClientbound, HandshakeServerbound};
+use crate::net::packet::status::{StatusClientbound, StatusServerbound};
+use std::io;
+use std::marker::PhantomData;
+use tokio::net::TcpStream;
+
+pub trait State {
+    type Clientbound: ProtocolState + Sync;
+    type Serverbound: ProtocolState + Sync;
+}
+
+#[allow(dead_code)]
+pub enum StateHandshake {}
+
+impl State for StateHandshake {
+    type Clientbound = HandshakeClientbound;
+    type Serverbound = HandshakeServerbound;
+}
+
+#[allow(dead_code)]
+pub enum StateStatus {}
+
+impl State for StateStatus {
+    type Clientbound = StatusClientbound;
+    type Serverbound = StatusServerbound;
+}
+
+#[allow(dead_code)]
+pub enum StateDisconnected {}
+
+impl State for StateDisconnected {
+    type Clientbound = !;
+    type Serverbound = !;
+}
+
+#[allow(dead_code)]
+pub enum StateLogin {}
+
+impl State for StateLogin {
+    type Clientbound = !;
+    type Serverbound = !;
+}
+
+#[allow(dead_code)]
+pub enum StatePlay {}
+
+impl State for StatePlay {
+    type Clientbound = !;
+    type Serverbound = !;
+}
+
+pub struct Connection<St: State> {
+    src: Reader,
+    dest: Writer,
+    st: PhantomData<St>,
+}
+
+impl<St: State> Connection<St> {
+    pub async fn write(&mut self, pkt: &St::Clientbound) -> io::Result<()> {
+        DefaultPacketFormat.send::<St::Clientbound>(&mut self.dest, pkt).await
+    }
+
+    pub async fn read(&mut self) -> io::Result<St::Serverbound> {
+        DefaultPacketFormat.recieve::<St::Serverbound>(&mut self.src).await
+    }
+
+    pub fn into_disconnected(self) -> Connection<StateDisconnected> {
+        Connection {
+            src: self.src,
+            dest: self.dest,
+            st: PhantomData,
+        }
+    }
+}
+
+impl Connection<StateHandshake> {
+    pub fn new(stream: TcpStream) -> Self {
+        use tokio::io::{BufReader, BufWriter};
+
+        let (src, dest) = stream.into_split();
+
+        Connection {
+            src: BufReader::new(src),
+            dest: BufWriter::new(dest),
+            st: PhantomData,
+        }
+    }
+
+    pub fn into_status(self) -> Connection<StateStatus> {
+        Connection {
+            src: self.src,
+            dest: self.dest,
+            st: PhantomData,
+        }
+    }
+
+    pub fn into_login(self) -> Connection<StateLogin> {
+        Connection {
+            src: self.src,
+            dest: self.dest,
+            st: PhantomData,
+        }
+    }
+}
+
+impl Connection<StateLogin> {
+    pub fn into_play(self) -> Connection<StatePlay> {
+        Connection {
+            src: self.src,
+            dest: self.dest,
+            st: PhantomData,
+        }
+    }
+}

src/net/format.rs

@@ -7,15 +7,15 @@ use tokio::io::AsyncReadExt;
 
 #[async_trait]
 pub trait PacketFormat {
-    async fn send<'wr, 'w, P: ProtocolState + Sync>(&self, dest: &'wr mut Writer<'w>, pkt: &P) -> io::Result<()>;
-    async fn recieve<'rr, 'r, P: ProtocolState>(&self, src: &'rr mut Reader<'r>) -> io::Result<P>;
+    async fn send<P: ProtocolState + Sync>(&self, dest: &mut Writer, pkt: &P) -> io::Result<()>;
+    async fn recieve<P: ProtocolState>(&self, src: &mut Reader) -> io::Result<P>;
 }
 
 pub const MAX_CLIENT_PACKET_SIZE: usize = 32767;
 
 pub struct DefaultPacketFormat;
 
-async fn read_varint<'rr, 'r>(src: &'rr mut Reader<'r>) -> io::Result<(usize, i32)> {
+async fn read_varint(src: &mut Reader) -> io::Result<(usize, i32)> {
     let mut length = 1;
     let mut acc = 0;
     while length <= 5 {
@@ -35,7 +35,7 @@ async fn read_varint<'rr, 'r>(src: &'rr mut Reader<'r>) -> io::Result<(usize, i3
 
 #[async_trait]
 impl PacketFormat for DefaultPacketFormat {
-    async fn send<'wr, 'w, P: ProtocolState + Sync>(&self, dest: &'wr mut Writer<'w>, pkt: &P) -> io::Result<()> {
+    async fn send<P: ProtocolState + Sync>(&self, dest: &mut Writer, pkt: &P) -> io::Result<()> {
         use crate::net::serialize::{PacketSerializer, VarInt};
 
         let packet_id = pkt.id();
@@ -61,7 +61,7 @@ impl PacketFormat for DefaultPacketFormat {
         Ok(())
     }
 
-    async fn recieve<'rr, 'r, P: ProtocolState>(&self, src: &'rr mut Reader<'r>) -> io::Result<P> {
+    async fn recieve<P: ProtocolState>(&self, src: &mut Reader) -> io::Result<P> {
         use crate::net::serialize::VecPacketDeserializer;
 
         let (_, length) = read_varint(src).await?;

src/net/mod.rs

@@ -1,11 +1,12 @@
 pub mod chat;
+pub mod connection;
 pub mod format;
 pub mod packet;
 pub mod serialize;
 pub mod state;
 
 use tokio::io::{BufReader, BufWriter};
-use tokio::net::tcp::{ReadHalf, WriteHalf};
+use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
 
-pub type Reader<'a> = BufReader<ReadHalf<'a>>;
-pub type Writer<'a> = BufWriter<WriteHalf<'a>>;
+pub type Reader= BufReader<OwnedReadHalf>;
+pub type Writer = BufWriter<OwnedWriteHalf>;

src/net/state.rs

@@ -9,6 +9,20 @@ pub trait ProtocolState: Sized {
     fn write(&self, ser: &mut impl PacketSerializer);
 }
 
+impl ProtocolState for ! {
+    fn id(&self) -> i32 {
+        match *self { }
+    }
+
+    fn read(_id: i32, _deser: &mut impl PacketDeserializer) -> Result<Self, String> {
+        Err("Cannot read packets; the connection state is disconnected.".to_string())
+    }
+
+    fn write(&self, _ser: &mut impl PacketSerializer) {
+        match *self { }
+    }
+}
+
 #[macro_export]
 macro_rules! define_states {
     { $( state $name:ident { $( $id:expr => $packet:ident ),* } )+ } => {