Handle disconnected users

src/client.rs

@@ -1,17 +1,22 @@
 use std::sync::Arc;
 
-use tokio::io::{AsyncWrite, AsyncRead};
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::sync::mpsc;
 use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
+use tokio::task::JoinHandle;
 
 use crate::connection::Connection;
 use crate::db::{Db, User};
-use crate::protocol::{MumblePacket, VoicePacket, MumblePacketWriter};
-use crate::proto::mumble::{UserState, Ping};
-use tokio::task::JoinHandle;
+use crate::proto::mumble::{Ping, UserRemove, UserState};
+use crate::protocol::{MumblePacket, MumblePacketWriter, VoicePacket};
 
 pub enum Message {
-    NewUser(u32)
+    UserConnected(u32),
+    UserDisconnected(u32),
+}
+
+pub enum ResponseMessage {
+    Disconnected
 }
 
 pub struct Client {
@@ -21,42 +26,45 @@ pub struct Client {
 }
 
 pub enum Error {
-    UserNotFound,
     StreamError(crate::protocol::Error),
 }
 
 struct Handler<W> {
     db: Arc<Db>,
     writer: MumblePacketWriter<W>,
-    response_sender: UnboundedSender<Message>,
+    session_id: u32,
+    response_sender: UnboundedSender<ResponseMessage>,
 }
 
 enum InnerMessage {
     Message(Message),
     Packet(MumblePacket),
+    Disconnected,
 }
 
-type ResponseReceiver = UnboundedReceiver<Message>;
+type ResponseReceiver = UnboundedReceiver<ResponseMessage>;
 
 impl Client {
     pub async fn new<S>(connection: Connection<S>, db: Arc<Db>) -> (Client, ResponseReceiver)
-    where
-        S: 'static + AsyncRead + AsyncWrite + Unpin + Send,
+        where
+            S: 'static + AsyncRead + AsyncWrite + Unpin + Send,
     {
         let (sender, mut receiver) = mpsc::unbounded_channel();
         let (response_sender, response_receiver) = mpsc::unbounded_channel();
 
         let writer = connection.writer;
+        let session_id = connection.session_id;
         let handler_task = tokio::spawn(async move {
             let mut handler = Handler {
                 db,
                 writer,
+                session_id,
                 response_sender,
             };
             loop {
                 let message = match receiver.recv().await {
-                    None => return,
                     Some(msg) => msg,
+                    None => return,
                 };
 
                 match message {
@@ -72,6 +80,10 @@ impl Client {
                             return;
                         }
                     }
+                    InnerMessage::Disconnected => {
+                        handler.self_disconnected().await;
+                        return;
+                    }
                 }
             }
         });
@@ -80,12 +92,13 @@ impl Client {
         let mut reader = connection.reader;
         let packet_task = tokio::spawn(async move {
             loop {
-                let packet = match reader.read().await{
-                    Ok(packet) => packet,
-                    Err(_) => return, //TODO
+                match reader.read().await {
+                    Ok(packet) => sender.send(InnerMessage::Packet(packet)),
+                    Err(_) => {
+                        sender.send(InnerMessage::Disconnected);
+                        return;
+                    }
                 };
-
-                sender.send(InnerMessage::Packet(packet));
             }
         });
 
@@ -124,7 +137,7 @@ impl<W> Handler<W>
             MumblePacket::UdpTunnel(voice) => {
                 match voice {
                     VoicePacket::Ping(_) => {
-                        self.writer.write(MumblePacket::UdpTunnel(voice));
+                        self.writer.write(MumblePacket::UdpTunnel(voice)).await;
                         println!("VoicePing");
                     }
                     VoicePacket::AudioData(_) => { println!("AudioData"); }
@@ -137,7 +150,8 @@ impl<W> Handler<W>
 
     async fn handle_message(&mut self, message: Message) -> Result<(), Error> {
         match message {
-            Message::NewUser(session_id) => self.new_user_connected(session_id).await?,
+            Message::UserConnected(session_id) => self.new_user_connected(session_id).await?,
+            Message::UserDisconnected(session_id) => self.user_disconnected(session_id).await?,
         }
 
         Ok(())
@@ -149,6 +163,17 @@ impl<W> Handler<W>
         }
         Ok(())
     }
+
+    async fn user_disconnected(&mut self, session_id: u32) -> Result<(), Error> {
+        let mut user_remove = UserRemove::new();
+        user_remove.set_session(session_id);
+        Ok(self.writer.write(MumblePacket::UserRemove(user_remove)).await?)
+    }
+
+    async fn self_disconnected(&mut self) {
+        self.db.remove_connected_user(self.session_id).await;
+        self.response_sender.send(ResponseMessage::Disconnected);
+    }
 }
 
 impl From<User> for UserState {

src/connection.rs

@@ -72,7 +72,7 @@ impl<S> Connection<S>
         let mut permission_query = PermissionQuery::new();
         permission_query.set_permissions(134743822);
         permission_query.set_channel_id(0);
-        writer.write(MumblePacket::PermissionQuery(permission_query));
+        writer.write(MumblePacket::PermissionQuery(permission_query)).await?;
 
         //User states
         let connected_users = db.get_connected_users().await;

src/db.rs

@@ -1,8 +1,7 @@
 use std::collections::HashMap;
 
-use tokio::sync::RwLock;
-
 use serde::{Deserialize, Serialize};
+use tokio::sync::RwLock;
 
 const ROOT_CHANNEL_ID: u32 = 0;
 const USER_TREE_NAME: &[u8] = b"users";
@@ -81,16 +80,21 @@ impl Db {
 
     pub async fn get_connected_users(&self) -> Vec<User> {
         let users = self.connected_users.read().await;
-        users.values().map(|el| el.clone()).collect()
+        users.values().cloned().collect()
     }
 
     pub async fn get_user_by_session_id(&self, session_id: u32) -> Option<User> {
         let connected_users = self.connected_users.read().await;
         if let Some(user) = connected_users.get(&session_id) {
-            return Some(user.clone())
+            return Some(user.clone());
         }
         None
     }
+
+    pub async fn remove_connected_user(&self, session_id: u32) {
+        let mut connected_users = self.connected_users.write().await;
+        connected_users.remove(&session_id);
+    }
 }
 
 impl Clone for User {

src/protocol.rs

@@ -73,7 +73,7 @@ pub enum VoicePacket {
 
 pub enum Error {
     UnknownPacketType,
-    ConnectionError,
+    ConnectionError(std::io::Error),
     ParsingError,
 }
 
@@ -442,17 +442,14 @@ fn serialize_voice_packet(packet: VoicePacket) -> Vec<u8> {
 }
 
 impl From<std::io::Error> for Error {
-    fn from(_: std::io::Error) -> Self {
-        Error::ConnectionError
+    fn from(err: std::io::Error) -> Self {
+        Error::ConnectionError(err)
     }
 }
 
 impl From<ProtobufError> for Error {
-    fn from(error: ProtobufError) -> Self {
-        match error {
-            ProtobufError::IoError(_) | ProtobufError::WireError(_) => Error::ConnectionError,
-            ProtobufError::Utf8(_) | ProtobufError::MessageNotInitialized { .. } => Error::ParsingError
-        }
+    fn from(_: ProtobufError) -> Self {
+        Error::ParsingError
     }
 }
 

src/server.rs

@@ -1,19 +1,15 @@
-use std::collections::{HashMap, HashSet};
-use std::future::Future;
+use std::collections::HashMap;
 use std::net::{IpAddr, SocketAddr};
 use std::sync::Arc;
-use std::task::{Context, Poll};
 
 use tokio::net::{TcpListener, TcpStream};
-use tokio::sync::{mpsc, RwLock};
-use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
+use tokio::sync::RwLock;
 use tokio_rustls::{server::TlsStream, TlsAcceptor};
 use tokio_rustls::rustls::{Certificate, NoClientAuth, PrivateKey, ServerConfig};
 
+use crate::client::{Client, Message, ResponseMessage};
 use crate::connection::{Connection, ConnectionConfig};
 use crate::db::Db;
-use crate::protocol::MumblePacket;
-use crate::client::{Client, Message};
 
 pub struct Config {
     pub ip_address: IpAddr,
@@ -58,7 +54,7 @@ async fn process(db: Arc<Db>, stream: TlsStream<TcpStream>, clients: Clients) {
         max_bandwidth: 128000,
         welcome_text: "Welcome!".to_string(),
     };
-    let mut connection = match Connection::setup_connection(db.clone(), stream, connection_config).await {
+    let connection = match Connection::setup_connection(db.clone(), stream, connection_config).await {
         Ok(connection) => connection,
         Err(_) => {
             eprintln!("Error establishing a connection");
@@ -71,7 +67,7 @@ async fn process(db: Arc<Db>, stream: TlsStream<TcpStream>, clients: Clients) {
     {
         let mut clients = clients.write().await;
         for client in clients.values() {
-            client.post_message(Message::NewUser(session_id))
+            client.post_message(Message::UserConnected(session_id))
         }
         clients.insert(session_id, client);
     }
@@ -83,7 +79,14 @@ async fn process(db: Arc<Db>, stream: TlsStream<TcpStream>, clients: Clients) {
         };
 
         match message {
-            _ => {}
+            ResponseMessage::Disconnected => {
+                let mut clients = clients.write().await;
+                clients.remove(&session_id);
+                for client in clients.values() {
+                    client.post_message(Message::UserDisconnected(session_id))
+                }
+                return;
+            }
         }
     }
 }