appsec/src/bin/connection_pool.rs

use std::{
    sync::{Arc, Mutex},
    collections::HashMap,
    net::{TcpStream},
    io::{Write},
};

pub struct ConnectionPool{
    conns: Arc<Mutex<HashMap<u32, TcpStream>>>,
}

impl ConnectionPool{

    pub fn new() -> Self {
        ConnectionPool{
            conns: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub fn add_connection(&self, iid: u32, conn: TcpStream) {
        self.conns.lock().unwrap().insert(iid, conn);
    }

    pub fn del_connection(&self, iid: u32) {
        self.conns.lock().unwrap().remove(&iid);
    }

    pub fn unicast(&self, recv_iid: &u32, msg: &String) {
        let mut del_conn = false;

        match self.conns.lock().unwrap().get(recv_iid) {
            Some(mut conn) => {
                if conn.write_all(msg.as_bytes()).is_err() {
                    del_conn = true;
                }
            },
            None => {
                println!("Unknown recv_iid: {:#?}", recv_iid);
                return;
            }
        }
        if del_conn{
            self.del_connection(*recv_iid);
        }
    }

    pub fn broadcast(&self, sender_iid: &u32, msg: &String) {
        let mut del_conns: Vec<u32> = Vec::new();
        for (iid, mut conn) in self.conns.lock().unwrap().iter() {
            if sender_iid.ne(iid){
                if conn.write_all(msg.as_bytes()).is_err() {
                    del_conns.push(*iid);
                    continue;
                }
                else if conn.write_all(b"\n").is_err() {
                    del_conns.push(*iid);
                }
            }
        }
        for iid in del_conns {
            self.del_connection(iid);
        }
    }
}