heimdall/crates/hd-bus/src/connection.rs

use std::{
    io::{Read, Write},
    net::TcpStream,
    sync::{Arc, mpsc::Receiver},
    time::Duration,
};

use crate::{
    bus::{Event, EventBus},
    error::HError,
    thread_pool::ThreadPool,
};

#[derive(Copy, Clone, PartialEq)]
enum EventBusClientConnectionState {
    Connect,
    ReadPrelude,
    ReadCursor,
    SendEvent,
    ReceiveAck,
    Disconnect,
}

struct EventBusClientConnection {
    stream: TcpStream,
    pub addr: String,
    state: EventBusClientConnectionState,
    subscription: Option<Receiver<Arc<Event>>>,
    bus: EventBus,
}

impl EventBusClientConnection {
    fn new(bus: EventBus, stream: TcpStream) -> Result<Self, HError> {
        let addr = stream
            .peer_addr()
            .map(|a| a.to_string())
            .unwrap_or("unknown".to_string());

        stream
            .set_read_timeout(Some(Duration::from_secs(5)))
            .map_err(HError::TcpPeerError)?;

        let state = EventBusClientConnectionState::Connect;

        Ok(Self {
            stream,
            addr,
            state,
            bus,
            subscription: None,
        })
    }
}

pub fn handle_event_bus_client(thread_pool: &ThreadPool, bus: EventBus, stream: TcpStream) -> Result<(), HError> {
    thread_pool.execute(move || {
        let mut client = EventBusClientConnection::new(bus, stream)?;

        loop {
            if let EventBusClientConnectionState::Disconnect = client.state {
                eprintln!("[{}] Closing connection due to invalid input", client.addr);
                return Ok(());
            }

            let prev_state = client.state;
            handle_event_bus_client_state(&mut client)
                .inspect_err(|e| eprintln!("[{}] Closing connection due to error {}", client.addr, e))?;
            if prev_state != client.state {
                println!("[{}] State change {} -> {}", client.addr, prev_state, client.state);
            }
        }
    })
}

fn handle_event_bus_client_state(client: &mut EventBusClientConnection) -> Result<(), HError> {
    let next_state = match client.state {
        EventBusClientConnectionState::Connect => handle_state_connect(client),
        EventBusClientConnectionState::ReadPrelude => todo!(),
        EventBusClientConnectionState::ReadCursor => todo!(),
        EventBusClientConnectionState::SendEvent => handle_state_send_event(client),
        EventBusClientConnectionState::ReceiveAck => todo!(),
        _ => Ok(EventBusClientConnectionState::Disconnect),
    }?;

    client.state = next_state;
    Ok(())
}

fn handle_state_connect(client: &mut EventBusClientConnection) -> Result<EventBusClientConnectionState, HError> {
    let mut buf = [0u8; 1];
    client.stream.read_exact(&mut buf).map_err(HError::TcpPeerError)?;
    if buf[0] == 1 {
        let subscription = client.bus.subscribe()?;
        client.subscription = Some(subscription);
        return Ok(EventBusClientConnectionState::SendEvent);
    }

    Ok(EventBusClientConnectionState::Disconnect)
}

fn handle_state_send_event(client: &mut EventBusClientConnection) -> Result<EventBusClientConnectionState, HError> {
    if let Some(subscription) = &mut client.subscription {
        loop {
            let Ok(evt) = subscription.recv() else {
                return Ok(EventBusClientConnectionState::Disconnect);
            };

            let data_len = evt.data.data_utf8.len();
            // [ ID  ] [TYPE] [ LEN ] [ DATA ]
            // 0  to 7   8    9 to 15 . . . .
            let buf_length = data_len + 17;
            let mut buf = vec![0u8; buf_length];
            write_le_bytes(evt.id, &mut buf[0..8]);
            write_le_bytes(data_len as u64, &mut buf[9..17]);
            buf[8] = evt.data.event_type;

            for i in 17..buf_length {
                buf[i] = evt.data.data_utf8[i - 17];
            }

            print!("Sending message : ");
            for i in &buf {
                print!("{} ", i);
            }
            println!("End");

            let stream = &mut client.stream;
            stream.write(&mut buf).map_err(HError::TcpPeerError)?;
            stream.flush().map_err(HError::TcpPeerError)?;

            let mut ack = vec![0];
            stream.read(&mut ack).map_err(HError::TcpPeerError)?;
        }
    }

    Ok(EventBusClientConnectionState::Disconnect)
}

fn write_le_bytes(v64: u64, buf: &mut [u8]) {
    let le_bytes = v64.to_le_bytes();
    buf[0] = le_bytes[0];
    buf[1] = le_bytes[1];
    buf[2] = le_bytes[2];
    buf[3] = le_bytes[3];
    buf[4] = le_bytes[4];
    buf[5] = le_bytes[5];
    buf[6] = le_bytes[6];
    buf[7] = le_bytes[7];
}

impl std::fmt::Display for EventBusClientConnectionState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let to_write = match self {
            EventBusClientConnectionState::Connect => "CONNECT",
            EventBusClientConnectionState::ReadPrelude => "READ_PRELUDE",
            EventBusClientConnectionState::ReadCursor => "READ_CURSOR",
            EventBusClientConnectionState::SendEvent => "SEND_EVENT",
            EventBusClientConnectionState::ReceiveAck => "RECEIVE_ACK",
            EventBusClientConnectionState::Disconnect => "DISCONNECT",
        };

        write!(f, "{}", to_write)
    }
}