golang_ipc_rs/

message.rs

use crate::{connection::Connection, encryption, Context};
use anyhow::{anyhow, bail};
use bytes::{Bytes, BytesMut};
use std::fmt;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tracing::*;
use std::sync::Arc;

pub(crate) type Channel = crate::bichannel::Channel<Message, Message>;

#[derive(Debug, Clone)]
pub struct Message {
    msg_type: Type,
    data: Option<RawData>,
}

#[derive(Debug, Clone)]
pub enum Type {
    Control,
    StatusChange(Status),
    Error(Arc<str>),
    General(u16),
}

#[derive(Debug, Clone)]
struct RawData {
    nonce: Option<BytesMut>,
    msg_type_raw: BytesMut,
    data: BytesMut,
}

#[allow(dead_code)]
#[derive(Default, Debug, Clone, Copy, PartialEq)]
pub enum Status {
    #[default]
    NotConnected,
    Listening,
    Connecting,
    Connected,
    ReConnecting,
    Closed,
    Closing,
    Error,
    Timeout,
}

const I32_SIZE: usize = std::mem::size_of::<i32>();

impl Message {
    pub(crate) fn new(msg_type_in: i32, data_in: Bytes, encryption_padding: bool) -> Message {
        // Create a BytesMut that can hold all the data
        //  This accounts for encryption overhead and the msg_type
        let data_len = data_in.len();

        let (mut data, nonce) = if encryption_padding {
            let (data, nonce) = encryption::new_data_buffer(I32_SIZE + data_len);
            (data, Some(nonce))
        } else {
            // Reserve for just the message
            (BytesMut::with_capacity(I32_SIZE + data_len), None)
        };

        // Write the message type, and split
        data.extend(msg_type_in.to_be_bytes());
        let msg_type_raw = data.split();

        // Write the data
        data.extend(data_in);

        Message {
            msg_type: msg_type_in.into(),
            data: Some(RawData {
                nonce,
                msg_type_raw,
                data,
            }),
        }
    }

    pub fn from_status(status: Status) -> Message {
        Message {
            msg_type: Type::StatusChange(status),
            data: None,
        }
    }

    pub fn from_error(err: crate::Error) -> Message {
        Message {
            msg_type: Type::Error(err.to_string().into()),
            data: None,
        }
    }

    pub fn is_data_message(&self) -> bool {
        matches!(self.msg_type, Type::General { .. })
    }

    pub fn data_length(&self) -> usize {
        if let Some(RawData { data, .. }) = &self.data {
            data.len()
        } else {
            0
        }
    }

    pub fn data(&self) -> Option<&[u8]> {
        if let Some(RawData { data, .. }) = &self.data {
            Some(data.as_ref())
        } else {
            None
        }
    }

    pub fn msg_type(&self) -> Type {
        self.msg_type.clone()
    }

    // pub fn stauts(&self) -> Type

    #[instrument(level = "trace")]
    pub(crate) async fn from_stream(connection: &mut Connection) -> crate::Result<Option<Message>> {
        let data_len = connection.reader.read_u32().await? as usize;

        if data_len < 4 {
            bail!("Not enough data for a message");
        }

        // let msg_type = connection.reader.read_i32().await?;

        // Allocate a buffer large enough for the data.
        let mut data = BytesMut::with_capacity(data_len);

        if 0 == connection.reader.read_buf(&mut data).await? {
            // The remote closed the connection. For this to be a clean
            // shutdown, there should be no data in the read buffer. If
            // there is, this means that the peer closed the socket while
            // sending a message.
            if data.is_empty() {
                return Ok(None);
            } else {
                bail!("connection reset by peer");
            }
        }

        // Decrypt the data
        let (mut data, nonce) = encryption::maybe_decrypt(&connection.cipher, data)?;

        // Split and decode the message type
        let msg_type_raw = data.split_to(I32_SIZE);
        let msg_type = i32::from_be_bytes(msg_type_raw.as_ref().try_into().unwrap()).into();

        let message = Message {
            msg_type,
            data: Some(RawData {
                nonce,
                msg_type_raw,
                data,
            }),
        };
        info!("recv message {}", message);

        Ok(Some(message))
    }

    #[instrument(level = "trace")]
    pub(crate) async fn to_stream(
        connection: &mut Connection,
        message: Message,
    ) -> crate::Result<()> {
        info!("send message {}", message);

        if let Some(data) = message.data {
            // Extract raw data objects
            let RawData {
                nonce,
                mut msg_type_raw,
                data,
                ..
            } = data;

            // Unsplit the buffer
            //  Note, this looks backwards but otherwise the data is in the wrong order
            msg_type_raw.unsplit(data);
            let data = msg_type_raw;

            // Encrypt the data
            let data = encryption::maybe_encrypt(&connection.cipher, data, nonce)?;

            connection.writer.write_u32(data.len() as u32).await?;
            connection.writer.write_all(data.as_ref()).await?;
        } else {
            bail!("Internal messages can not send.");
        }

        Ok(())
    }
}

pub(crate) async fn message_loop(
    mut connection: Connection,
    context: &mut Context,
) -> crate::Result<()> {
    loop {
        tokio::select! {
            result = Message::from_stream(&mut connection) => {
                match result {
                    Ok(Some(message)) => {
                        // We received a message over the socket, so send it to the consumer
                        if let Err(_err) = context.get_channel().send(message).await {
                            context.report_error_status(anyhow!("Dropped message")).await;
                        }
                    },
                    Ok(None) => {
                        context.report_error_status(anyhow!("Remote disconnect")).await;
                        context.report_status(Status::Closing).await?;
                        break;
                    },
                    Err(err) =>  {
                        let new_status = match err.downcast_ref::<std::io::Error>() {
                            Some(io_err) => {
                                if io_err.kind() == std::io::ErrorKind::UnexpectedEof {
                                    Status::ReConnecting
                                } else {
                                    Status::Closing
                                }
                            },
                            None => {
                                Status::Closing
                            },
                        };
                        context.report_error_status(err).await;
                        context.report_status(new_status).await?;
                        break;
                    },
                }

            },
            to_send = context.get_channel().recv() => {
                match to_send {
                    Some(message) => {

                        if let Err(err) = Message::to_stream(&mut connection, message).await {
                            // TODO: The message is dropped in this case.  Can we put it back into `context.get_channel()`?
                            context.report_error_status(err).await;
                        }
                    },
                    None => {
                        context.report_error_status(anyhow!("bad recv")).await;
                    },
                }
            },
        };
    }

    Ok(())
}

impl From<i32> for Type {
    fn from(value: i32) -> Type {
        match value {
            0 => Type::Control,
            -1 => Type::StatusChange(Status::Error), // not expected
            -2 => Type::Error("Unknown Error".into()), // not expected
            _ => {
                if value > 0 {
                    Type::General(value as u16)
                } else {
                    // TODO: how does Go version handle this
                    panic!("Unexpected mess_type {}", value)
                }
            }
        }
    }
}

impl From<Type> for i32 {
    fn from(value: Type) -> i32 {
        match value {
            Type::Control => 0,
            Type::StatusChange(_) => -1,
            Type::Error(_) => -2,
            Type::General(value) => value as i32,
        }
    }
}

impl fmt::Display for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match &self.msg_type {
            Type::StatusChange(status) => write!(f, "{}: {}", self.msg_type, status),
            Type::Control => write!(f, "{}", self.msg_type),
            Type::Error(err) => write!(f, "{}: {}", self.msg_type, err),
            Type::General(_) => {
                write!(
                    f,
                    "{}: [{}] {:?}",
                    self.msg_type,
                    self.data_length(),
                    self.data
                )
            }
        }
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl fmt::Display for Status {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}