merge

src/main.rs

@@ -0,0 +1,99 @@
+use tokio::net::{TcpListener, TcpStream};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use std::env;
+use std::error::Error;
+use std::collections::HashMap;
+use std::net::SocketAddr;
+use std::sync::{Arc, Mutex};
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn Error>> {
+    let args: Vec<String> = env::args().collect();
+    if args.len() != 3 {
+        println!("Usage: {} [server|client] [addr:port]", args[0]);
+        return Ok(());
+    }
+
+    let command = &args[1];
+    let addr = &args[2];
+
+    match command.as_str() {
+        "server" => start_server(addr).await?,
+        "client" => start_client(addr).await?,
+        _ => println!("Unknown command {}", command),
+    }
+
+    Ok(())
+}
+
+async fn start_server(addr: &str) -> Result<(), Box<dyn Error>> {
+    let listener = TcpListener::bind(addr).await?;
+
+    println!("Server listening on {}", addr);
+
+    let clients: Arc<Mutex<HashMap<SocketAddr, String>>> = Arc::new(Mutex::new(HashMap::new()));
+
+    loop {
+        let (mut socket, client_addr) = listener.accept().await?;
+
+        let mut buf = vec![0; 1024];
+        socket.read(&mut buf).await?;
+    
+        let username = String::from_utf8(buf.clone()).unwrap();
+
+        println!("Accepted connection from {}", client_addr);
+
+        let clients_shared = Arc::clone(&clients);
+
+        clients_shared.lock().unwrap().insert(client_addr, username);
+
+        tokio::spawn(async move {
+            let mut buf = vec![0; 1024];
+
+            loop {
+                let n = match socket.read(&mut buf).await {
+                    Ok(n) if n == 0 => {
+                        let username = clients_shared.lock().unwrap().remove(&client_addr).unwrap_or("Unknown".to_string());
+                        println!("Client {} ({}) disconnected!", username, client_addr);
+                        return;
+                    },
+                    Ok(n) => n,
+                    Err(e) => {
+                        eprintln!("failed to read from socket; err = {:?}", e);
+                        return;
+                    }
+                };
+
+                let username = clients_shared.lock().unwrap().get(&client_addr).unwrap_or(&"Unknown".to_string()).clone();
+
+                println!("Received message from {} ({}): {}",
+                    username,
+                    client_addr,
+                    String::from_utf8_lossy(&buf[..n]));
+
+                if let Err(e) = socket.write_all(&buf[0..n]).await {
+                    eprintln!("failed to write to socket; err = {:?}", e);
+                    return;
+                }
+            }
+        });
+    }
+}
+
+async fn start_client(addr: &str) -> Result<(), Box<dyn Error>> {
+    let mut stream = TcpStream::connect(addr).await?;
+    
+    println!("Enter a username: ");
+    let mut username = String::new(); 
+    std::io::stdin().read_line(&mut username)?;
+    
+    stream.write_all(username.as_bytes()).await?;
+
+    println!("Connected to server at {}", addr);
+    let mut input = String::new();
+    loop {
+        std::io::stdin().read_line(&mut input)?;
+        stream.write_all(input.as_bytes()).await?;
+        input.clear();
+    }
+}