113 lines
3.6 KiB
Rust
113 lines
3.6 KiB
Rust
use std::env;
|
|
use std::error::Error;
|
|
use std::net::SocketAddr;
|
|
use std::sync::Arc;
|
|
use std::collections::HashMap;
|
|
use tokio::io::{AsyncReadExt, AsyncBufReadExt, AsyncWriteExt, BufReader};
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
use tokio::sync::Mutex;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn Error>> {
|
|
// Parse command-line arguments
|
|
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>> {
|
|
// Start the TCP listener
|
|
let listener = TcpListener::bind(addr).await?;
|
|
println!("Server listening on {}", addr);
|
|
|
|
// Shared data for client connections
|
|
let clients: Arc<Mutex<HashMap<SocketAddr, String>>> = Arc::new(Mutex::new(HashMap::new()));
|
|
|
|
loop {
|
|
// Accept a new client connection
|
|
let (mut socket, client_addr) = listener.accept().await?;
|
|
|
|
// Read username from the client
|
|
let mut buf = vec![0; 1024];
|
|
let n = socket.read(&mut buf).await?;
|
|
let username = String::from_utf8_lossy(&buf[..n]).trim().to_string();
|
|
|
|
println!("Accepted connection from {}", client_addr);
|
|
|
|
let clients_shared = Arc::clone(&clients);
|
|
|
|
// Insert client into the shared HashMap
|
|
clients_shared.lock().await.insert(client_addr, username.clone());
|
|
|
|
// Clone the HashMap and release the lock
|
|
let clients_map = clients_shared.lock().await.clone();
|
|
|
|
// Spawn a new task for each client connection
|
|
tokio::spawn(async move {
|
|
let mut buf = vec![0; 1024];
|
|
|
|
loop {
|
|
let n = match socket.read(&mut buf).await {
|
|
Ok(n) if n == 0 => {
|
|
// Client disconnected
|
|
let username = clients_shared.lock().await.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_map.get(&client_addr).unwrap_or(&"Unknown".to_string()).clone();
|
|
|
|
let message = String::from_utf8_lossy(&buf[..n]).to_string();
|
|
println!(
|
|
"Received message from ({})\n{}: {}",
|
|
client_addr,
|
|
username,
|
|
message
|
|
);
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
async fn start_client(addr: &str) -> Result<(), Box<dyn Error>> {
|
|
// Connect to the server
|
|
let mut stream = TcpStream::connect(addr).await?;
|
|
|
|
println!("Enter a username: ");
|
|
let mut username = String::new();
|
|
let mut reader = BufReader::new(tokio::io::stdin());
|
|
reader.read_line(&mut username).await?;
|
|
|
|
// Send username to the server
|
|
stream.write_all(username.trim().as_bytes()).await?;
|
|
|
|
println!("Connected to server at {}", addr);
|
|
loop {
|
|
let mut input = String::new();
|
|
reader.read_line(&mut input).await?;
|
|
stream.write_all(input.trim().as_bytes()).await?;
|
|
|
|
// Print sent message
|
|
println!("Sent message from {}: {}", username.trim(), input.trim());
|
|
}
|
|
}
|