*SNEEZE*

src/main.rs

@@ -1,4 +1,5 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
+use std::time::Instant;
 use std::{io::Cursor, net::SocketAddr, str::FromStr, sync::Arc, time::Duration};
 
 use irc_rust::Message;
@@ -23,7 +24,8 @@ struct Client {
     username: Option<String>,
     realname: Option<String>,
 
-    modes: HashSet<u8>, // [byte per mode](https://www.unrealircd.org/docs/User_modes)
+    // modes: HashSet<u8>, // [byte per mode](https://www.unrealircd.org/docs/User_modes)
+    last_pong: Instant,
 
     c2s_tx: Option<Sender<Message>>,
     s2c_rx: Option<Receiver<Message>>,
@@ -42,7 +44,8 @@ impl Client {
             username: None,
             realname: None,
 
-            modes: HashSet::new(),
+            // modes: HashSet::new(),
+            last_pong: Instant::now(),
 
             c2s_tx: Some(c2s_tx),
             s2c_rx: Some(s2c_rx),
@@ -58,7 +61,7 @@ impl Client {
 }
 
 #[derive(Default, Debug)]
-struct AppState {
+pub struct AppState {
     clients: RwLock<HashMap<SocketAddr, Client>>,
 }
 
@@ -80,11 +83,18 @@ impl AppState {
         }
     }
 
+    pub async fn nick(&self, peer_addr: SocketAddr) -> Option<String> {
+        if let Some(client) = self.clients.read().await.get(&peer_addr) {
+            return client.nick.clone();
+        }
+        None
+    }
+
     pub async fn user(
         &self,
         peer_addr: SocketAddr,
         username: &str,
-        mode: u16,
+        // mode: u16,
         realname: Option<&str>,
     ) {
         if let Some(client) = self.clients.write().await.get_mut(&peer_addr) {
@@ -101,10 +111,17 @@ impl AppState {
         false
     }
 
+    async fn pong(&self, peer_addr: SocketAddr) {
+        if let Some(client) = self.clients.write().await.get_mut(&peer_addr) {
+            client.last_pong = Instant::now();
+        }
+    }
+
     async fn quit(&self, peer_addr: SocketAddr, reason: Option<&str>) {
         // broadcast user leaving unless invisible flag was set
 
         self.clients.write().await.remove(&peer_addr);
+        println!("Bye {}: {}", peer_addr, reason.unwrap_or_default());
     }
 
     // clients still needs mod/op status stuff
@@ -142,15 +159,14 @@ impl AppState {
         Ok(())
     }
 
-    pub async fn broadcast(&self, msg: Message) -> Result<()> {
+    // pub async fn broadcast(&self, msg: Message) -> Result<()> {
-        let clients = self.clients.read().await;
+    //     let clients = self.clients.read().await;
-        let addresses = clients.keys().cloned();
+    //     let addresses = clients.keys().cloned();
-        for peer_addr in addresses {
+    //     for peer_addr in addresses {
-            self.send(peer_addr, msg.clone()).await?;
+    //         self.send(peer_addr, msg.clone()).await?;
-        }
+    //     }
-
+    //     Ok(())
-        Ok(())
+    // }
-    }
 }
 
 #[tokio::main]
@@ -173,7 +189,7 @@ async fn main() -> Result<()> {
     let acceptor = TlsAcceptor::from(Arc::new(config));
     let listener = TcpListener::bind("0.0.0.0:6697").await?;
 
-    // which means, this goes into yet another thread
+    // Connection loop
     let state: Arc<AppState> = app_state.clone();
     tokio::spawn(async move {
         loop {
@@ -199,12 +215,17 @@ async fn main() -> Result<()> {
         }
     });
 
-    // todo: the "processing" loop goes here...
+    // Processing loop
     let state: Arc<AppState> = app_state.clone();
     tokio::spawn(async move {
         loop {
             let messages = state.get_messages().await;
-            // println!("Got {} messages", messages.len());
+            if messages.is_empty() {
+                tokio::time::sleep(Duration::from_millis(50)).await;
+                continue;
+            }
+
+            println!("Got {} messages", messages.len());
 
             for (peer_addr, msg) in messages {
                 println!("{}: {}", peer_addr, msg);
@@ -212,7 +233,6 @@ async fn main() -> Result<()> {
                     eprintln!("Error handling command: {:?}", e);
                 }
             }
-            tokio::time::sleep(Duration::from_millis(1)).await;
         }
     });
 
@@ -221,21 +241,33 @@ async fn main() -> Result<()> {
     tokio::spawn(async move {
         loop {
             let clients = state.clients.read().await;
+            let mut mark_for_delete = vec![];
             for (peer_addr, client) in clients.iter() {
+                if Instant::now().duration_since(client.last_pong).as_secs() >= 110 {
+                    mark_for_delete.push(*peer_addr);
+                }
                 println!("PING :{}", client.username.clone().unwrap_or_default());
-                state
+                let _ = state
                     .send(
                         *peer_addr,
                         Message::from_str(&format!(
-                            ":localhost PING :{}",
+                            "PING :{}",
                             client.username.clone().unwrap_or_default()
                         ))?,
                     )
-                    .await?;
+                    .await;
             }
-
             drop(clients);
 
+            // Drop stale clients, if there are any
+            if !mark_for_delete.is_empty() {
+                state
+                    .clients
+                    .write()
+                    .await
+                    .retain(|p, _| !mark_for_delete.contains(p));
+            }
+
             tokio::time::sleep(Duration::from_secs(30)).await;
         }
 
@@ -255,7 +287,13 @@ async fn handle(
 ) -> Result<Option<Message>> {
     match msg.command()? {
         "CAP" => {
-            state.send(peer_addr, Message::from_str("NONE")?).await?;
+            dbg!(&msg);
+            state
+                .send(
+                    peer_addr,
+                    "CAP * LS :multi-prefix sasl=PLAIN,EXTERNAL".into(),
+                )
+                .await?;
         }
         "NICK" => {
             let nick = msg.params()?.next().unwrap_or_default();
@@ -264,20 +302,28 @@ async fn handle(
         "USER" => {
             let mut params = msg.params()?;
             let username = params.next().unwrap_or_default(); // aviinl
-            let mode = params.next().unwrap_or_default().parse::<u16>()?; // mode
+            // let mode = params.next().unwrap_or_default().parse::<u16>()?; // mode
-            params.next().unwrap_or_default(); // unused
+            // params.next().unwrap_or_default(); // unused
 
             let realname = msg.trailing()?; // realname
+            state.user(peer_addr, username, /*mode,*/ realname).await;
 
-            dbg!(username, realname);
+            let nick = state.nick(peer_addr).await.unwrap_or(username.to_string());
-
+            state
-            state.user(peer_addr, username, mode, realname).await;
+                .send(
+                    peer_addr,
+                    format!(":localhost 001 {} :Welcome to the IRC server.", nick).into(),
+                )
+                .await?;
         }
         "PASS" => {
             let mut params = msg.params()?;
             let password = params.next().unwrap_or_default(); // aviinl
             state.pass(peer_addr, password).await;
         }
+        "PONG" => {
+            state.pong(peer_addr).await;
+        }
         "QUIT" => {
             let reason = msg.trailing()?; // realname
             state.quit(peer_addr, reason).await;
@@ -332,7 +378,10 @@ async fn connect(
                         break 'outer;
                     }
                 }
+                continue;
             }
+            let _ = c2s_tx.send("QUIT :Broken Pipe".into()).await;
+            break;
         }
         println!("{:?} closed", peer_addr);
     });
@@ -340,7 +389,8 @@ async fn connect(
     // fire off a send loop
     tokio::spawn(async move {
         while let Some(i) = s2c_rx.recv().await {
-            if let Err(e) = tx.write(i.to_string().as_bytes()).await {
+            let data = i.to_string() + "\r\n";
+            if let Err(e) = tx.write(data.as_bytes()).await {
                 eprintln!("Error sending message: {:?}", e);
                 break;
             }