m32chat/src/main.rs

use chrono::{DateTime, Duration, Local};
use log::debug;
use std::{collections::HashMap, net::UdpSocket, thread};

const CLIENT_TIMEOUT: u32 = 600;
const KEEPALIVE: u64 = 10;
const MAX_CLIENTS: usize = 20;
const MY_SPEED: u8 = 30;

fn strip_header(msg: &[u8]) -> Vec<u8> {
    let mut stripped = Vec::new();
    stripped.push(msg[1] & 3);
    stripped.append(msg[2..].to_vec().as_mut());
    stripped.to_owned()
}

fn mopp(speed: u8, data: &[u8]) -> Vec<u8> {
    let serial = 1u8; // Currently not really used

    let morse = HashMap::from([
        (b'0', "-----"),
        (b'1', ".----"),
        (b'2', "..---"),
        (b'3', "...--"),
        (b'4', "....-"),
        (b'5', "....."),
        (b'6', "-...."),
        (b'7', "--..."),
        (b'8', "---.."),
        (b'9', "----."),
        (b'a', ".-"),
        (b'b', "-..."),
        (b'c', "-.-."),
        (b'd', "-.."),
        (b'e', "."),
        (b'f', "..-."),
        (b'g', "--."),
        (b'h', "...."),
        (b'i', ".."),
        (b'j', ".---"),
        (b'k', "-.-"),
        (b'l', ".-.."),
        (b'm', "--"),
        (b'n', "-."),
        (b'o', "---"),
        (b'p', ".--."),
        (b'q', "--.-"),
        (b'r', ".-."),
        (b's', "..."),
        (b't', "-"),
        (b'u', "..-"),
        (b'v', "...-"),
        (b'w', ".--"),
        (b'x', "-..-"),
        (b'y', "-.--"),
        (b'z', "--.."),
        (b'=', "-...-"),
        (b'/', "-..-."),
        (b'+', ".-.-."),
        (b'-', "-....-"),
        (b'.', ".-.-.-"),
        (b',', "--..--"),
        (b'?', "..--.."),
        (b':', "---..."),
        (b'!', "-.-.--"),
        (b'\'', ".----."),
    ]);

    let mut m = "01".to_owned(); // protocol version 1
    m.push_str(&format!("{:06b}", serial));
    m.push_str(&format!("{:06b}", speed));

    for &char in data {
        if char == b' ' {
            continue;
        }

        for b in morse.get(&char).unwrap().as_bytes() {
            if *b == b'.' {
                m.push_str("01");
            } else {
                m.push_str("10");
            }
        }

        m.push_str("00"); // EOC
    }

    m.replace_range(m.len() - 2.., "11"); // EOW

    // Extend to full bytes
    m = format!(
        "{1:0<0$}",
        (8 as f32 * (m.len() as f32 / 8 as f32).ceil()) as usize,
        m
    );

    let mut res = Vec::<u8>::new();

    // Convert string representation to "real" bits
    for i in (0..m.len()).step_by(8) {
        let value = u8::from_str_radix(&m[i..i + 8], 2).unwrap();
        res.push(value);
    }

    res.to_owned()
}

fn broadcast(
    socket: &UdpSocket,
    receivers: &HashMap<String, DateTime<Local>>,
    client: &str,
    data: &[u8],
) {
    for rec in receivers {
        // Do not broadcast to origin
        if rec.0 == client {
            continue;
        }

        debug!("Broadcasting to client {}", rec.0);
        
        socket.send_to(&data, &rec.0).unwrap();
    }
}

fn remove_old_clients(socket: &UdpSocket, subscribers: &mut HashMap<String, DateTime<Local>>) {
    let timestamp = Local::now();

    //debug!("Number of clients: {}", subscribers.len());

    for cl in &*subscribers {
        if *cl.1 + Duration::seconds(CLIENT_TIMEOUT as i64) < timestamp {
            debug!("Removing old client {}", cl.0);
            socket
                .send_to(mopp(MY_SPEED, b":bye").as_slice(), &cl.0)
                .unwrap();
        }
    }

    subscribers.retain(|_, val| *val + Duration::seconds(CLIENT_TIMEOUT as i64) >= timestamp);
}

fn main() -> std::io::Result<()> {
    env_logger::init();

    let socket = UdpSocket::bind("0.0.0.0:7373")?;
    socket
        .set_read_timeout(Some(core::time::Duration::from_secs(KEEPALIVE)))
        .expect("Could not set read timeout");

    debug!("Morserino chat server started.");

    let mut subscribers: HashMap<String, DateTime<Local>> = HashMap::new();

    loop {
        thread::sleep(core::time::Duration::from_millis(200)); // Anti flood (?)

        let mut buf = [0; 64];

        // Waiting for incoming packets. Otherwise, after timeout, send keepalive packet
        let result = socket.recv_from(&mut buf);
        let (number_of_bytes, src_addr) = match result {
            Ok((num, s)) => (num, s),
            Err(_) => {
                remove_old_clients(&socket, &mut subscribers);

                for rec in &subscribers {
                    debug!("Sending keepalive packet to {}", rec.0);
                    socket.send_to(b"", rec.0).unwrap();
                }
                continue;
            }
        };

        // Just do the very least of plausibility checks
        if number_of_bytes < 2 {
            // Abort if not at least 2 bytes have been received
            debug!("Dropping unknown packet (too small)");
            continue;
        } else if (buf[0] >> 6) != 1u8 {
            // Abort if protocol version is not "1"
            debug!("Dropping unknown packet (version != 1)");
            continue;
        }

        let client_addr = src_addr.to_string();
        let speed = buf[1] >> 2;
        let data = &buf[0..number_of_bytes];

        if subscribers.contains_key(&client_addr) {
            debug!("Client known: {client_addr}");

            if strip_header(data) == strip_header(mopp(speed, b":bye").as_slice()) {
                debug!("Removing client {client_addr} as requested.");
                socket
                    .send_to(mopp(speed, b":bye").as_slice(), &client_addr)
                    .unwrap();
                subscribers.remove(&client_addr);
            } else {
                broadcast(&socket, &subscribers, &client_addr, data);
                subscribers.insert(client_addr.to_owned(), Local::now());
            }
        } else if strip_header(data) == strip_header(mopp(speed, b"hi").as_slice()) {
            debug!("Adding new client {client_addr}");
            if subscribers.len() < MAX_CLIENTS {
                subscribers.insert(client_addr.to_owned(), Local::now());
                socket
                    .send_to(
                        mopp(speed, format!("{}{}", ":hi ", subscribers.len()).as_bytes())
                            .as_slice(),
                        &client_addr,
                    )
                    .unwrap();
            } else {
                debug!("There are already too many clients connected.");
                socket
                    .send_to(mopp(speed, b":qrl").as_slice(), &client_addr)
                    .unwrap();
            }
        } else {
            debug!("Ignoring unknown client {client_addr}");
        }

        remove_old_clients(&socket, &mut subscribers);
    }
}