iscsi-client-rs 0.0.3

use std::any::type_name;

use anyhow::{Context, Result, bail};

use crate::{
    cfg::config::Config,
    client::pdu_connection::FromBytes,
    models::common::{BasicHeaderSegment, Builder, SendingData},
};

#[derive(Debug)]
pub struct PDUWithData<T> {
    pub header: T,
    pub aditional_heder: Vec<u8>,
    pub header_digest: Option<u32>,
    pub data: Vec<u8>,
    pub data_digest: Option<u32>,
}

impl<T> Builder for PDUWithData<T>
where T: BasicHeaderSegment + SendingData + FromBytes
{
    type Header = Vec<u8>;

    /// Appends raw bytes to the Data Segment and updates its length field.
    fn append_data(&mut self, more: Vec<u8>) {
        self.data.extend_from_slice(&more);
        let len = self.data.len() as u32;
        self.header.set_data_length_bytes(len);
    }

    /// Build finnal PDU (BHS + DataSegment)
    fn build(&mut self, cfg: &Config) -> Result<Vec<(Self::Header, Vec<u8>)>> {
        let mrdsl = cfg.login.negotiation.max_recv_data_segment_length as usize;
        if mrdsl == 0 {
            bail!("MaxRecvDataSegmentLength is zero");
        }

        let chunks = if self.data.is_empty() {
            vec![&self.data[..]]
        } else {
            self.data.chunks(mrdsl).collect::<Vec<_>>()
        };
        let padding_ahs = (4 - (self.aditional_heder.len() % 4)) % 4;

        let mut final_body = Vec::with_capacity(chunks.len());
        for (i, chunk) in chunks.iter().enumerate() {
            if i == chunks.len() - 1 {
                self.header.set_final_bit();
            } else {
                self.header.set_continue_bit();
            }
            self.header.set_data_length_bytes(chunk.len() as u32);
            let bhs = T::to_bhs_bytes(&self.header)?;
            let padding_chunk = (4 - (chunk.len() % 4)) % 4;
            let mut body = Vec::with_capacity(
                self.aditional_heder.len()
                    + padding_ahs
                    + (self.header_digest.is_some() as usize) * 4
                    + chunk.len()
                    + padding_chunk
                    + (self.data_digest.is_some() as usize) * 4,
            );

            if !self.aditional_heder.is_empty() {
                body.extend_from_slice(&self.aditional_heder);
                body.extend(std::iter::repeat_n(0u8, padding_ahs));
            }

            if let Some(hd) = self.header_digest {
                body.extend_from_slice(&hd.to_be_bytes());
            }

            body.extend_from_slice(chunk);
            body.extend(std::iter::repeat_n(0u8, padding_chunk));

            if let Some(dd) = self.data_digest {
                body.extend_from_slice(&dd.to_be_bytes());
            }

            final_body.push((bhs.to_vec(), body));
        }

        Ok(final_body)
    }
}

impl<T> PDUWithData<T>
where T: BasicHeaderSegment + FromBytes
{
    pub fn from_header(header: T) -> Self {
        Self {
            header,
            aditional_heder: vec![],
            header_digest: None,
            data: vec![],
            data_digest: None,
        }
    }

    /// Parse PDU: BHS(=48) + AHS + pad(AHS) + [HeaderDigest?] + Data +
    /// pad(Data) + [DataDigest?]
    pub fn parse(
        header: T,
        buf: &[u8],
        enable_header_digest: bool,
        enable_data_digest: bool,
    ) -> Result<Self> {
        let tn = type_name::<T>();

        let ahs_len = header.get_ahs_length_bytes();
        let data_len = header.get_data_length_bytes();
        let ahs_pad = (4 - (ahs_len % 4)) % 4;
        let data_pad = (4 - (data_len % 4)) % 4;

        let mut off = 0;

        // --- AHS ---
        let aditional_heder = if ahs_len > 0 {
            if buf.len() < off + ahs_len {
                bail!(
                    "{tn}: buffer {} too small for AHS end {}",
                    buf.len(),
                    off + ahs_len
                );
            }
            let v = buf[off..off + ahs_len].to_vec();
            off += ahs_len;

            if buf.len() < off + ahs_pad {
                bail!(
                    "{tn}: buffer {} too small for AHS padding end {}",
                    buf.len(),
                    off + ahs_pad
                );
            }
            off += ahs_pad;
            v
        } else {
            Vec::new()
        };

        // --- HeaderDigest? ---
        let header_digest = if enable_header_digest {
            if buf.len() < off + 4 {
                bail!("{tn}: no room for HeaderDigest");
            }
            let hd = u32::from_be_bytes(
                buf[off..off + 4]
                    .try_into()
                    .context("expected header_digest, but failed to build")?,
            );
            off += 4;
            Some(hd)
        } else {
            None
        };

        // --- Data ---
        if buf.len() < off + data_len {
            bail!(
                "{tn}: buffer {} too small for Data end {}",
                buf.len(),
                off + data_len
            );
        }
        let data = buf[off..off + data_len].to_vec();
        off += data_len;

        if buf.len() < off + data_pad {
            bail!(
                "{tn}: buffer {} too small for Data padding end {}",
                buf.len(),
                off + data_pad
            );
        }
        off += data_pad;

        // --- DataDigest? ---
        let data_digest = if enable_data_digest {
            if buf.len() < off + 4 {
                bail!("{tn}: no room for DataDigest");
            }
            let dd = u32::from_be_bytes(
                buf[off..off + 4]
                    .try_into()
                    .context("expected data_digest, but failed to build")?,
            );
            //off += 4;
            Some(dd)
        } else {
            None
        };

        Ok(Self {
            header,
            aditional_heder,
            header_digest,
            data,
            data_digest,
        })
    }
}