rustchat/src/main.rs

use clap::Parser;
use std::fs::File;
use std::io::Write;
use std::io::{self, BufRead};
use std::net::{TcpListener, TcpStream};
use std::sync::{Arc, Mutex};
use std::thread;

mod chat_error;
use chat_error::ChatError;

mod chat_user;

mod chat_server;
use chat_server::ChatServer;

mod config;
use config::ConfigArgs;

fn handle_client(
    mut stream: TcpStream,
    server: Arc<Mutex<ChatServer>>,
    client_id: &mut Option<u64>,
) -> Result<(), ChatError> {
    let mut writer = stream.try_clone()?;
    let reader = io::BufReader::new(&mut stream);
    let mut lines = reader.lines();

    // greet and recv nick
    write!(writer, "Hello and welcome to RustChat!\nNick: ")?;
    let mut nick;
    loop {
        nick = lines
            .next()
            .ok_or(ChatError::Protocol(String::from("Could not recv nickname")))??;
        // check for nick name uniqueness
        let mut srv = server.lock()?;
        if !srv.is_nick_in_use(&nick) {
            // print joined msg and register at server
            srv.send_to_all(&format!("* {} has joined the chat", nick))?;
            let curr_users = srv.get_user_count();
            let users_in_room_msg = if curr_users == 1 {
                String::from("is 1 user")
            } else {
                format!("are {} users", curr_users)
            };
            write!(
                writer,
                " *** {}\n *** You are known as {}. There {} in the room.\n",
                srv.config.greeting_msg, nick, users_in_room_msg
            )?;
            *client_id = Some(srv.register(nick.clone(), writer));

            break;
        }
        write!(
            writer,
            "This nickname is already in use, please choose another.\nNick: "
        )?;
    }

    // read lines
    for line in lines {
        let msg_raw = line?;
        if msg_raw.is_empty() {
            continue;
        }
        let msg = format!("<{}> {}", nick, msg_raw);
        server.lock()?.send_to_all(&msg)?;
    }
    Ok(())
}

fn run_thread(stream: TcpStream, server: Arc<Mutex<ChatServer>>) {
    let mut client_id = None;
    let quit_reason = match handle_client(stream, Arc::clone(&server), &mut client_id) {
        Ok(()) => String::from("Connection closed by user's choice"),
        Err(err) => format!("Client foobar't their thread: {:?}", err),
    };

    if let Some(client_id) = client_id {
        let mut srv = server.lock().unwrap();
        let name = srv
            .deregister(client_id)
            .map(|client| client.name)
            .unwrap_or("<unknown user>".into());
        let quit_msg = format!("* {} has left the chat ({})", name, quit_reason);
        println!("{}", quit_msg);
        srv.send_to_all(&quit_msg).unwrap();
    } else {
        println!("Client thread terminated before identifying themselves");
    }
}

/// Chatserver in RUST!
#[derive(Parser, Debug)]
#[clap(about, version, author)]
struct Args {
    /// Path to config file
    #[clap(short, long, default_value = "config.yaml")]
    config: String,

    /// IP address to bind to
    #[clap(long)]
    host: Option<String>,

    /// Port to run the server on
    #[clap(short, long)]
    port: Option<u16>,
}

fn main() -> std::io::Result<()> {
    // argument parsing
    let args = Args::parse();

    // read config file
    let config_reader = File::open(args.config).expect("Could not open config file");
    let mut config: ConfigArgs =
        serde_yaml::from_reader(&config_reader).expect("Could not parse config file");

    if let Some(host) = args.host {
        config.host = host;
    }
    if let Some(port) = args.port {
        config.port = port;
    }
    let binding_host = format!("{}:{}", config.host, config.port);
    println!("Binding to {}", binding_host);

    let listener = TcpListener::bind(binding_host)?;
    let server = Arc::new(Mutex::new(ChatServer::new(config)));

    // accept connections and process them serially
    for stream in listener.incoming() {
        if let Ok(stream) = stream {
            let server_clone = Arc::clone(&server);
            thread::spawn(move || {
                run_thread(stream, server_clone);
            });
        }
    }
    Ok(())
}