RaaS/src/robot/connection_handler.rs

use crate::robot::robot_connector::RobotConnector;
use crate::robot::ROBOT_MESSAGE_QUEUE_SIZE;
use crate::robot::{ConnectorManager, Error};
use log::{error, info};
use prost::Message;
use raas_types::raas::register::{Register, RegisterResponse};
use raas_types::raas::{recv_raas_msg, send_raas_msg};
use std::sync::Arc;
use tokio::io::duplex;
use tokio::net::TcpListener;
use tokio::sync::Mutex;

pub struct ConnectionHandler {
    tcp_listener: TcpListener,
    addr: String,
}

impl ConnectionHandler {
    pub fn new(addr: &str, tcp_listener: TcpListener) -> Self {
        Self {
            tcp_listener,
            addr: addr.to_string(),
        }
    }

    async fn handle_register(
        &mut self,
        connector_manager: Arc<Mutex<ConnectorManager>>,
    ) -> Result<(), Error> {
        info!("Listening for connection at {}", self.addr);
        let (mut socket, addr) = self.tcp_listener.accept().await?;
        info!("Got connection from {}", addr);

        let register_msg = recv_raas_msg(&mut socket).await?;

        let register = Register::decode(&*register_msg.msg)?;

        let mut connector_manager = connector_manager.lock().await;

        let id = connector_manager.next_id;
        connector_manager.next_id += 1;

        let register_resp = RegisterResponse {
            name: register.name.to_string(),
            r#type: register.bot_type,
            id,
        };

        let mut msg = Vec::new();

        register_resp.encode(&mut msg)?;

        send_raas_msg(&mut socket, msg).await?;

        let (connector_duplex_manager, connector_duplex_connector) =
            duplex(ROBOT_MESSAGE_QUEUE_SIZE);

        let mut connector = RobotConnector::new(
            id,
            vec![register.bot_type()],
            socket,
            connector_duplex_connector,
        );

        let connector_handler = crate::robot::ConnectorHandle {
            stream: connector_duplex_manager,
            tags: connector.bot_tags.clone(),
        };

        connector_manager.connectors.insert(id, connector_handler);

        tokio::spawn(async move { connector.worker().await });

        Ok(())
    }

    pub async fn worker(
        mut self,
        connector_manager: Arc<Mutex<ConnectorManager>>,
    ) -> Result<(), Error> {
        loop {
            let res = self.handle_register(connector_manager.clone()).await;

            if let Err(err) = res {
                error!("Got error handling new register: {}", err);

                if let Error::Io(_io_err) = &err {
                    error!("IO Error, exiting...");
                    return Err(err);
                }
            }
        }
    }
}