add problem 03

problem03/src/main.rs

@@ -0,0 +1,175 @@
+use std::sync::Arc;
+
+use anyhow::bail;
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt, Result},
+    net::TcpStream,
+    sync::{Mutex, MutexGuard, broadcast::Sender},
+};
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let listener = tokio::net::TcpListener::bind("0.0.0.0:10000").await?;
+
+    let room = Arc::new(Mutex::new(vec![]));
+    let (tx, _) = tokio::sync::broadcast::channel(100_000);
+
+    loop {
+        let (socket, _) = listener.accept().await?;
+        let the_room = room.clone();
+        let this_tx = tx.clone();
+        tokio::spawn(async move {
+            match handle_connection(socket, the_room, this_tx).await {
+                Ok(_) => (),
+                Err(err) => {
+                    dbg!(err);
+                }
+            }
+        });
+    }
+}
+
+async fn handle_connection(
+    socket: TcpStream,
+    room: Arc<Mutex<Vec<String>>>,
+    tx: Sender<MessageType>,
+) -> anyhow::Result<()> {
+    let (mut reader, mut writer) = tokio::io::split(socket);
+
+    writer
+        .write_all(b"Welcome to budgetchat! What shall I call you?\n")
+        .await?;
+
+    let mut read_buffer = Vec::with_capacity(1024);
+
+    read_to_buffer(&mut reader, &mut read_buffer).await?;
+
+    let newline = find_next_newline(&read_buffer).unwrap();
+    let before_newline = read_buffer.drain(0..=newline).collect();
+
+    let self_name = String::from_utf8(before_newline)?.trim().to_string();
+
+    let message = generate_room_contains_message(&room).await;
+    writer.write_all(&message.into_bytes()).await?;
+
+    let mut rx = {
+        let mut locked_room = room.lock().await;
+
+        deny_invalid_user(&self_name, &locked_room)?;
+        locked_room.push(self_name.clone());
+
+        let rx = tx.subscribe();
+
+        let message = MessageType::UserJoined(self_name.clone());
+        tx.send(message)?;
+
+        rx
+    };
+
+    loop {
+        tokio::select! {
+            read_amount = read_to_buffer(&mut reader, &mut read_buffer) => {
+                match read_amount {
+                    Ok(_) => (),
+                    Err(_) => {
+                        tx.send(MessageType::UserQuit(self_name.clone()))?;
+                        {
+                            let mut the_room = room.lock().await;
+                            let idx = the_room.iter().position(|name| *name == self_name).unwrap();
+                            the_room.remove(idx);
+                        }
+                        return Ok(());
+                    }
+                };
+
+                let newline = match find_next_newline(&read_buffer) {
+                    Some(idx) => idx,
+                    None => continue,
+                };
+                let before_newline = read_buffer.drain(0..=newline).collect();
+
+                let message = String::from_utf8(before_newline)?;
+                if !message.trim().is_empty() {
+                    tx.send(MessageType::TextMessage(self_name.clone(), message.to_string()))?;
+
+                }
+            }
+            received = rx.recv() => {
+                match received? {
+                    MessageType::UserJoined(user_name) => {
+                        if user_name == self_name {
+                            continue;
+                        }
+                        let message = format!("* {user_name} has entered the room\n");
+                        writer.write_all(&message.into_bytes()).await?;
+                    }
+                    MessageType::UserQuit(user_name) => {
+                        let message = format!("* {user_name} has left the room\n");
+                        writer.write_all(&message.into_bytes()).await?;
+                    }
+                    MessageType::TextMessage(user_name, the_message) => {
+                        if user_name == self_name {
+                            continue;
+                        }
+                        let message = format!("[{user_name}] {}\n", the_message.trim());
+                        writer.write_all(&message.into_bytes()).await?;
+                    }
+                }
+            }
+
+        };
+    }
+}
+
+fn find_next_newline(read_buffer: &[u8]) -> Option<usize> {
+    read_buffer.iter().position(|c| *c == b'\n')
+}
+
+async fn read_to_buffer(
+    reader: &mut tokio::io::ReadHalf<TcpStream>,
+    read_buffer: &mut Vec<u8>,
+) -> anyhow::Result<usize> {
+    let mut read_amount = 0;
+
+    while !read_buffer.contains(&b'\n') {
+        read_amount += reader.read_buf(read_buffer).await?;
+        if read_amount == 0 {
+            bail!("user disconnected");
+        }
+    }
+    Ok(read_amount)
+}
+
+fn deny_invalid_user(name: &str, room: &MutexGuard<'_, Vec<String>>) -> anyhow::Result<()> {
+    if name.is_empty()
+        || !name.chars().all(|c| c.is_alphanumeric())
+        || room.iter().any(|n| n == name)
+    {
+        bail!("user '{name}' is invalid or already exists");
+    }
+
+    Ok(())
+}
+
+async fn generate_room_contains_message(room: &Arc<Mutex<Vec<String>>>) -> String {
+    let the_room = room.lock().await;
+    let mut message = String::from("* The room contains: ");
+    let mut iter = the_room.iter().peekable();
+
+    while let Some(user) = iter.next() {
+        message.push_str(user.as_str());
+        if iter.peek().is_some() {
+            message.push_str(", ");
+        }
+    }
+
+    message.push('\n');
+    message
+}
+
+#[derive(Clone, Debug)]
+enum MessageType {
+    UserJoined(String),
+    UserQuit(String),
+    TextMessage(String, String), // user, message
+}