datum-agent 0.10.3

use prost::Message as ProstMessage;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

use crate::dcp::{DcpError, DcpResult, proto::DcpFrame};

const FRAME_LEN_BYTES: usize = 4;
const MAX_DCP_FRAME_BYTES: usize = 16 * 1024 * 1024;

pub(crate) async fn read_frame<R>(reader: &mut R) -> DcpResult<Option<DcpFrame>>
where
    R: AsyncRead + Unpin,
{
    let mut header = [0_u8; FRAME_LEN_BYTES];
    match reader.read_exact(&mut header).await {
        Ok(_) => {}
        Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
        Err(error) => return Err(error.into()),
    }

    let len = u32::from_be_bytes(header) as usize;
    if len > MAX_DCP_FRAME_BYTES {
        return Err(DcpError::Protocol(format!(
            "DCP frame exceeds {MAX_DCP_FRAME_BYTES} byte limit: {len}"
        )));
    }
    if len == 0 {
        return Err(DcpError::Protocol("DCP frame length is zero".to_owned()));
    }

    let mut payload = vec![0_u8; len];
    reader.read_exact(&mut payload).await?;
    Ok(Some(DcpFrame::decode(payload.as_slice())?))
}

pub(crate) async fn write_frame<W>(writer: &mut W, frame: &DcpFrame) -> DcpResult<()>
where
    W: AsyncWrite + Unpin,
{
    let payload = frame.encode_to_vec();
    let len = u32::try_from(payload.len()).map_err(|_| {
        DcpError::Protocol(format!(
            "DCP frame exceeds {MAX_DCP_FRAME_BYTES} byte limit: {}",
            payload.len()
        ))
    })?;
    if payload.len() > MAX_DCP_FRAME_BYTES {
        return Err(DcpError::Protocol(format!(
            "DCP frame exceeds {MAX_DCP_FRAME_BYTES} byte limit: {}",
            payload.len()
        )));
    }

    writer.write_all(&len.to_be_bytes()).await?;
    writer.write_all(&payload).await?;
    writer.flush().await?;
    Ok(())
}