markbase/rust-iscsi-initiator/src/connection/mod.rs

use crate::Result;
use crate::pdu::IscsiPdu;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

/// iSCSI Connection
pub struct IscsiConnection {
    /// TCP stream
    stream: TcpStream,

    /// Session ID
    session_id: u64,

    /// Command Sequence Number
    cmd_sn: u32,

    /// Status Sequence Number
    stat_sn: u32,
}

impl IscsiConnection {
    /// Create new connection to iSCSI target
    pub async fn connect(addr: &str) -> Result<Self> {
        let stream = TcpStream::connect(addr).await?;

        Ok(Self {
            stream,
            session_id: 0,
            cmd_sn: 0,
            stat_sn: 0,
        })
    }

    /// Send PDU to target
    pub async fn send_pdu(&mut self, pdu: &IscsiPdu) -> Result<()> {
        let data = pdu.encode();
        self.stream.write_all(&data).await?;
        Ok(())
    }

    /// Receive PDU from target
    pub async fn recv_pdu(&mut self) -> Result<IscsiPdu> {
        // Read header (48 bytes)
        let mut header = [0u8; 48];
        self.stream.read_exact(&mut header).await?;

        // Parse data segment length
        let data_len =
            ((header[5] as usize) << 16) | ((header[6] as usize) << 8) | (header[7] as usize);

        // Read data segment if present
        if data_len > 0 {
            let mut full_pdu = Vec::with_capacity(48 + data_len);
            full_pdu.extend_from_slice(&header);

            let mut data = vec![0u8; data_len];
            self.stream.read_exact(&mut data).await?;
            full_pdu.extend(data);

            IscsiPdu::decode(&full_pdu)
        } else {
            IscsiPdu::decode(&header)
        }
        .map_err(Into::into)
    }

    /// Perform iSCSI login
    pub async fn login(&mut self, initiator_name: &str, target_name: &str) -> Result<()> {
        let login_pdu = IscsiPdu::login_request(initiator_name, target_name);
        self.send_pdu(&login_pdu).await?;

        let response = self.recv_pdu().await?;

        // Check login response
        if response.opcode != 0x23 {
            return Err(crate::Error::Protocol("Invalid login response".into()));
        }

        // Parse login parameters
        if !response.data.is_empty() {
            let params = String::from_utf8_lossy(&response.data);
            log::info!("Login response: {}", params);
        }

        Ok(())
    }

    /// Send NOP-Out (keepalive)
    pub async fn nop_out(&mut self) -> Result<()> {
        let pdu = IscsiPdu::nop_out();
        self.send_pdu(&pdu).await?;

        let response = self.recv_pdu().await?;

        if response.opcode != 0x20 {
            return Err(crate::Error::Protocol("Invalid NOP-In response".into()));
        }

        Ok(())
    }

    /// Close connection
    pub async fn close(&mut self) -> Result<()> {
        self.stream.shutdown().await?;
        Ok(())
    }

    /// Get next command sequence number
    pub fn next_cmd_sn(&mut self) -> u32 {
        let sn = self.cmd_sn;
        self.cmd_sn += 1;
        sn
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_connection_mock() {
        // Mock test - would need real iSCSI target for full test
        let result = IscsiConnection::connect("127.0.0.1:3260").await;
        // Should fail if no target running
        assert!(result.is_ok() || result.is_err());
    }
}