qssh 0.4.3

//! Quantum-secure transport layer

use crate::compression::{CompressionAlgorithm, CompressionContext};
use crate::{crypto::SymmetricCrypto, QsshError, Result};
use bincode;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::TcpStream;
use tokio::sync::{Mutex, Notify};

pub mod channel;
pub mod protocol;
pub mod quantum_resistant;

pub use channel::*;
pub use protocol::*;
pub use quantum_resistant::*;

/// Unified transport trait for both classical and quantum-resistant transports
#[async_trait::async_trait]
pub trait QsshTransport: Send + Sync {
    /// Send a message
    async fn send_message<T: serde::Serialize + Send + Sync>(&self, message: &T) -> Result<()>;

    /// Receive a message
    async fn receive_message<T: for<'de> serde::Deserialize<'de>>(&self) -> Result<T>;

    /// Close the transport
    async fn close(&self) -> Result<()>;
}

/// Maximum raw frame size (1MB) for unauthenticated inbound data.
const MAX_RAW_MESSAGE_SIZE: usize = 1024 * 1024;

/// Maximum serialized message size (1MB).
const MAX_MESSAGE_SIZE: usize = MAX_RAW_MESSAGE_SIZE;

/// In-memory transport used by unit tests.
///
/// This mock has no network side effects and can be shared across tasks.
#[derive(Default)]
pub struct MockTransport {
    incoming: Mutex<VecDeque<u8>>,
    outgoing: Mutex<VecDeque<u8>>,
    closed: AtomicBool,
    incoming_notify: Notify,
}

impl MockTransport {
    /// Create an empty mock transport.
    pub fn new() -> Self {
        Self::default()
    }

    /// Inject a serialized frame into the incoming buffer.
    pub async fn inject_frame(&self, frame: &[u8]) {
        let mut incoming = self.incoming.lock().await;
        incoming.extend(frame.iter().copied());
        self.incoming_notify.notify_waiters();
    }

    /// Inject a typed message into the incoming buffer.
    pub async fn inject_message<T: serde::Serialize>(&self, message: &T) -> Result<()> {
        let payload = bincode::serialize(message)
            .map_err(|e| QsshError::Protocol(format!("Serialization failed: {}", e)))?;
        let frame = encode_frame(&payload)?;
        self.inject_frame(&frame).await;
        Ok(())
    }

    /// Serialize a raw payload into a framed message and queue it to outgoing.
    pub async fn send_raw_message(&self, payload: &[u8]) -> std::result::Result<Vec<u8>, QsshError> {
        if self.closed.load(Ordering::Acquire) {
            return Err(QsshError::Connection("Mock transport is closed".into()));
        }

        let frame = encode_frame(payload)?;
        let mut outgoing = self.outgoing.lock().await;
        outgoing.extend(frame.iter().copied());
        Ok(frame)
    }

    /// Wait for and return the next fully framed incoming payload.
    pub async fn receive_raw_message(&self) -> std::result::Result<Vec<u8>, QsshError> {
        loop {
            {
                let mut incoming = self.incoming.lock().await;
                if let Some(payload) = try_pop_frame(&mut incoming)? {
                    return Ok(payload);
                }
            }

            if self.closed.load(Ordering::Acquire) {
                return Err(QsshError::Connection("Mock transport is closed".into()));
            }

            self.incoming_notify.notified().await;
        }
    }

    /// Drain all raw bytes written by send_message.
    pub async fn take_outgoing_bytes(&self) -> Vec<u8> {
        let mut outgoing = self.outgoing.lock().await;
        outgoing.drain(..).collect()
    }

    /// Pop and decode one sent message from the outgoing buffer, if present.
    pub async fn take_outgoing_message<T: for<'de> serde::Deserialize<'de>>(
        &self,
    ) -> Result<Option<T>> {
        let mut outgoing = self.outgoing.lock().await;
        let Some(payload) = try_pop_frame(&mut outgoing)? else {
            return Ok(None);
        };

        let message = bincode::deserialize(&payload)
            .map_err(|e| QsshError::Protocol(format!("Deserialization failed: {}", e)))?;
        Ok(Some(message))
    }

    /// Return the current number of queued incoming bytes.
    pub async fn incoming_len(&self) -> usize {
        self.incoming.lock().await.len()
    }

    /// Return the current number of queued outgoing bytes.
    pub async fn outgoing_len(&self) -> usize {
        self.outgoing.lock().await.len()
    }
}

#[async_trait::async_trait]
impl QsshTransport for MockTransport {
    async fn send_message<T: serde::Serialize + Send + Sync>(&self, message: &T) -> Result<()> {
        let payload = bincode::serialize(message)
            .map_err(|e| QsshError::Protocol(format!("Serialization failed: {}", e)))?;
        let _ = self.send_raw_message(&payload).await?;
        Ok(())
    }

    async fn receive_message<T: for<'de> serde::Deserialize<'de>>(&self) -> Result<T> {
        let payload = self.receive_raw_message().await?;
        let message = bincode::deserialize(&payload)
            .map_err(|e| QsshError::Protocol(format!("Deserialization failed: {}", e)))?;
        Ok(message)
    }

    async fn close(&self) -> Result<()> {
        self.closed.store(true, Ordering::Release);
        self.incoming_notify.notify_waiters();
        Ok(())
    }
}

fn encode_frame(payload: &[u8]) -> Result<Vec<u8>> {
    if payload.len() > MAX_RAW_MESSAGE_SIZE {
        return Err(QsshError::Protocol("Message too large".into()));
    }

    let mut frame = Vec::with_capacity(4 + payload.len());
    frame.extend_from_slice(&(payload.len() as u32).to_be_bytes());
    frame.extend_from_slice(payload);
    Ok(frame)
}

fn try_pop_frame(buffer: &mut VecDeque<u8>) -> Result<Option<Vec<u8>>> {
    if buffer.len() < 4 {
        return Ok(None);
    }

    let len_bytes = [
        *buffer.front()
            .ok_or_else(|| QsshError::Protocol("Truncated frame: missing length byte 0".into()))?,
        *buffer
            .get(1)
            .ok_or_else(|| QsshError::Protocol("Truncated frame: missing length byte 1".into()))?,
        *buffer
            .get(2)
            .ok_or_else(|| QsshError::Protocol("Truncated frame: missing length byte 2".into()))?,
        *buffer
            .get(3)
            .ok_or_else(|| QsshError::Protocol("Truncated frame: missing length byte 3".into()))?,
    ];

    let frame_len = u32::from_be_bytes(len_bytes) as usize;
    if frame_len > MAX_RAW_MESSAGE_SIZE {
        return Err(QsshError::Protocol("Frame too large".into()));
    }

    if buffer.len() < 4 + frame_len {
        return Ok(None);
    }

    let mut payload = Vec::with_capacity(frame_len);
    for i in 0..frame_len {
        let byte = *buffer
            .get(4 + i)
            .ok_or_else(|| QsshError::Protocol("Truncated frame: missing payload byte".into()))?;
        payload.push(byte);
    }

    for _ in 0..(4 + frame_len) {
        let _ = buffer
            .pop_front()
            .ok_or_else(|| QsshError::Protocol("Truncated frame: buffer underflow".into()))?;
    }

    Ok(Some(payload))
}

/// Transport layer for encrypted communication
#[derive(Clone)]
pub struct Transport {
    reader: Arc<Mutex<OwnedReadHalf>>,
    writer: Arc<Mutex<OwnedWriteHalf>>,
    send_crypto: Arc<RwLock<SymmetricCrypto>>,
    recv_crypto: Arc<RwLock<SymmetricCrypto>>,
    send_sequence: Arc<Mutex<u64>>,
    recv_sequence: Arc<Mutex<u64>>,
    /// Number of completed rekey operations
    rekey_count: Arc<std::sync::atomic::AtomicU32>,
    /// Compression context for outgoing data
    send_compression: Arc<Mutex<CompressionContext>>,
    /// Compression context for incoming data
    recv_compression: Arc<Mutex<CompressionContext>>,
}

impl Transport {
    /// Create new transport from established TCP connection (unidirectional - for backwards compatibility)
    pub fn new(stream: TcpStream, crypto: SymmetricCrypto) -> Result<Self> {
        let (reader, writer) = stream.into_split();
        let default_recv_crypto = SymmetricCrypto::from_shared_secret(&[0u8; 32])?;
        Ok(Self {
            reader: Arc::new(Mutex::new(reader)),
            writer: Arc::new(Mutex::new(writer)),
            send_crypto: Arc::new(RwLock::new(crypto)),
            recv_crypto: Arc::new(RwLock::new(default_recv_crypto)),
            send_sequence: Arc::new(Mutex::new(0)),
            recv_sequence: Arc::new(Mutex::new(0)),
            rekey_count: Arc::new(std::sync::atomic::AtomicU32::new(0)),
            send_compression: Arc::new(Mutex::new(CompressionContext::new(
                CompressionAlgorithm::None,
                6,
            ))),
            recv_compression: Arc::new(Mutex::new(CompressionContext::new(
                CompressionAlgorithm::None,
                6,
            ))),
        })
    }

    /// Create new transport with separate send and receive crypto
    pub fn new_bidirectional(
        stream: TcpStream,
        send_crypto: SymmetricCrypto,
        recv_crypto: SymmetricCrypto,
    ) -> Self {
        let (reader, writer) = stream.into_split();
        Self {
            reader: Arc::new(Mutex::new(reader)),
            writer: Arc::new(Mutex::new(writer)),
            send_crypto: Arc::new(RwLock::new(send_crypto)),
            recv_crypto: Arc::new(RwLock::new(recv_crypto)),
            send_sequence: Arc::new(Mutex::new(0)),
            recv_sequence: Arc::new(Mutex::new(0)),
            rekey_count: Arc::new(std::sync::atomic::AtomicU32::new(0)),
            send_compression: Arc::new(Mutex::new(CompressionContext::new(
                CompressionAlgorithm::None,
                6,
            ))),
            recv_compression: Arc::new(Mutex::new(CompressionContext::new(
                CompressionAlgorithm::None,
                6,
            ))),
        }
    }

    /// Send an encrypted message
    pub async fn send_message<T: Serialize>(&self, message: &T) -> Result<()> {
        log::trace!("Transport: sending message");
        // Serialize message
        let plaintext = bincode::serialize(message)
            .map_err(|e| QsshError::Protocol(format!("Serialization failed: {}", e)))?;

        if plaintext.len() > MAX_MESSAGE_SIZE {
            return Err(QsshError::Protocol("Message too large".into()));
        }

        // Compress serialized payload before encryption
        let plaintext = {
            let mut comp = self.send_compression.lock().await;
            comp.compress(&plaintext)?
        };

        // Get sequence number
        let seq = {
            let mut seq_lock = self.send_sequence.lock().await;
            let current = *seq_lock;
            *seq_lock += 1;
            current
        };

        // Add sequence number to prevent replay attacks
        let mut authenticated_data = Vec::new();
        authenticated_data.extend_from_slice(&seq.to_be_bytes());
        authenticated_data.extend_from_slice(&plaintext);

        // Encrypt (read lock: non-blocking for concurrent recv)
        let (ciphertext, nonce) = {
            let crypto = self
                .send_crypto
                .read()
                .map_err(|_| QsshError::Crypto("Send crypto lock poisoned".into()))?;
            crypto.encrypt(&authenticated_data)?
        };

        // Frame format: [4 bytes length][12 bytes nonce][ciphertext]
        let frame_length = (nonce.len() + ciphertext.len()) as u32;
        let mut frame = Vec::new();
        frame.extend_from_slice(&frame_length.to_be_bytes());
        frame.extend_from_slice(&nonce);
        frame.extend_from_slice(&ciphertext);

        // Send
        let mut writer = self.writer.lock().await;
        log::trace!("Transport: writing {} bytes to stream", frame.len());
        writer.write_all(&frame).await.map_err(|e| {
            log::error!("Transport: failed to write to stream: {}", e);
            QsshError::Io(e)
        })?;
        writer.flush().await.map_err(|e| {
            log::error!("Transport: failed to flush stream: {}", e);
            QsshError::Io(e)
        })?;
        log::trace!("Transport: message sent successfully");

        Ok(())
    }

    /// Receive and decrypt a message
    pub async fn receive_message<T: for<'de> Deserialize<'de>>(&self) -> Result<T> {
        let mut reader = self.reader.lock().await;

        // Read frame length
        let mut length_bytes = [0u8; 4];
        log::trace!("Transport: attempting to read 4 bytes for frame length");
        reader.read_exact(&mut length_bytes).await.map_err(|e| {
            log::error!("Transport: failed to read frame length: {}", e);
            QsshError::Io(e)
        })?;
        let frame_length = u32::from_be_bytes(length_bytes) as usize;

        if frame_length > MAX_MESSAGE_SIZE {
            return Err(QsshError::Protocol("Frame too large".into()));
        }

        // Read nonce and ciphertext
        let mut frame = vec![0u8; frame_length];
        reader.read_exact(&mut frame).await.map_err(QsshError::Io)?;

        let (nonce, ciphertext) = frame.split_at(12);

        // Decrypt (read lock: non-blocking for concurrent send)
        let authenticated_data = {
            let crypto = self
                .recv_crypto
                .read()
                .map_err(|_| QsshError::Crypto("Recv crypto lock poisoned".into()))?;
            crypto.decrypt(ciphertext, nonce)?
        };

        // Verify sequence number
        if authenticated_data.len() < 8 {
            return Err(QsshError::Protocol("Invalid message format".into()));
        }

        let (seq_bytes, plaintext) = authenticated_data.split_at(8);
        let received_seq = u64::from_be_bytes(
            seq_bytes
                .try_into()
                .map_err(|_| QsshError::Protocol("Invalid sequence number format".into()))?,
        );

        // Check and update sequence number
        let expected_seq = {
            let mut seq_lock = self.recv_sequence.lock().await;
            let current = *seq_lock;
            *seq_lock += 1;
            current
        };

        if received_seq != expected_seq {
            return Err(QsshError::Protocol(format!(
                "Invalid sequence number: expected {}, got {}",
                expected_seq, received_seq
            )));
        }

        // Decompress before deserialization
        let plaintext = {
            let mut comp = self.recv_compression.lock().await;
            comp.decompress(plaintext)?
        };

        // Deserialize
        let message = bincode::deserialize(&plaintext)
            .map_err(|e| QsshError::Protocol(format!("Deserialization failed: {}", e)))?;

        Ok(message)
    }

    /// Enable compression on this transport (called after handshake negotiation)
    pub async fn set_compression(&self, algorithm: CompressionAlgorithm, level: u32) {
        let mut send = self.send_compression.lock().await;
        *send = CompressionContext::new(algorithm, level);
        let mut recv = self.recv_compression.lock().await;
        *recv = CompressionContext::new(algorithm, level);
        if algorithm.is_enabled() {
            log::info!(
                "Transport compression enabled: {:?} (level {})",
                algorithm,
                level
            );
        }
    }

    /// Get compression statistics for sent data
    pub async fn compression_stats(
        &self,
    ) -> (
        crate::compression::CompressionStats,
        crate::compression::CompressionStats,
    ) {
        let send = self.send_compression.lock().await;
        let recv = self.recv_compression.lock().await;
        (send.stats().clone(), recv.stats().clone())
    }

    /// Update encryption keys (rekey). Acquires write locks on both crypto instances
    /// and resets sequence numbers to prevent nonce reuse with new keys.
    pub async fn update_keys(
        &self,
        new_send: SymmetricCrypto,
        new_recv: SymmetricCrypto,
    ) -> Result<()> {
        // Acquire writer mutex to ensure no send is in flight
        let _writer_guard = self.writer.lock().await;
        // Acquire reader mutex to ensure no recv is in flight
        let _reader_guard = self.reader.lock().await;

        // Swap crypto keys
        {
            let mut send = self
                .send_crypto
                .write()
                .map_err(|_| QsshError::Crypto("Send crypto lock poisoned during rekey".into()))?;
            *send = new_send;
        }
        {
            let mut recv = self
                .recv_crypto
                .write()
                .map_err(|_| QsshError::Crypto("Recv crypto lock poisoned during rekey".into()))?;
            *recv = new_recv;
        }

        // Reset sequence numbers to 0 for the new key epoch
        {
            let mut seq = self.send_sequence.lock().await;
            *seq = 0;
        }
        {
            let mut seq = self.recv_sequence.lock().await;
            *seq = 0;
        }

        let count = self
            .rekey_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
            + 1;
        log::info!("Rekey #{} complete — new session keys active", count);

        Ok(())
    }

    /// Get the number of completed rekey operations
    pub fn rekey_count(&self) -> u32 {
        self.rekey_count.load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Close the transport
    pub async fn close(&self) -> Result<()> {
        let mut writer = self.writer.lock().await;
        writer.shutdown().await.map_err(QsshError::Io)?;
        Ok(())
    }
}

/// Implement unified transport trait for classical transport
#[async_trait::async_trait]
impl QsshTransport for Transport {
    async fn send_message<T: serde::Serialize + Send + Sync>(&self, message: &T) -> Result<()> {
        self.send_message(message).await
    }

    async fn receive_message<T: for<'de> serde::Deserialize<'de>>(&self) -> Result<T> {
        self.receive_message().await
    }

    async fn close(&self) -> Result<()> {
        self.close().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::time::{timeout, Duration};

    #[tokio::test]
    async fn mock_transport_inject_and_receive_message() {
        let transport = MockTransport::new();
        let expected = Message::Ping(42);

        transport.inject_message(&expected).await.unwrap();
        let received: Message = transport.receive_message().await.unwrap();

        match received {
            Message::Ping(v) => assert_eq!(v, 42),
            _ => panic!("unexpected message variant"),
        }
    }

    #[tokio::test]
    async fn mock_transport_send_and_drain_outgoing_message() {
        let transport = MockTransport::new();
        let sent = Message::Pong(7);

        transport.send_message(&sent).await.unwrap();
        let drained: Option<Message> = transport.take_outgoing_message().await.unwrap();

        match drained {
            Some(Message::Pong(v)) => assert_eq!(v, 7),
            _ => panic!("unexpected outgoing message"),
        }
    }

    #[tokio::test]
    async fn mock_transport_close_unblocks_receiver() {
        let transport = Arc::new(MockTransport::new());
        let receiver = transport.clone();

        let wait_task = tokio::spawn(async move { receiver.receive_message::<Message>().await });

        tokio::task::yield_now().await;
        transport.close().await.unwrap();

        let result = timeout(Duration::from_secs(1), wait_task).await;
        assert!(result.is_ok(), "receiver task did not complete in time");

        let join_result = result.unwrap().unwrap();
        assert!(join_result.is_err(), "receiver should fail after close");
    }
}

/// Kani bounded model checking harnesses for classical transport.
///
/// Verifies integer cast safety, bounds checking, and panic-freedom
/// for the encrypted transport layer.
///
/// Run with: `cargo kani --harness <harness_name>`
#[cfg(kani)]
mod kani_proofs {
    use super::*;

    // ── Step 6: Integer Cast Safety ────────────────────────────────────────

    /// Proves the `(nonce.len() + ciphertext.len()) as u32` cast at line 103
    /// never truncates, given MAX_MESSAGE_SIZE bounds.
    /// AES-256-GCM: nonce=12 bytes, tag=16 bytes, so max ciphertext
    /// = MAX_MESSAGE_SIZE + 8 (sequence) + 16 (tag) = 1048600 bytes.
    /// Total frame_length = 12 + 1048600 = 1048612, well within u32::MAX.
    #[kani::proof]
    fn proof_transport_frame_length_cast() {
        let nonce_len: usize = 12; // AES-GCM nonce is always 12 bytes
        let ciphertext_len: usize = kani::any();
        // Max ciphertext = plaintext + sequence(8) + AES-GCM tag(16)
        kani::assume(ciphertext_len <= MAX_MESSAGE_SIZE + 8 + 16);

        let total = nonce_len + ciphertext_len;
        let cast_result = total as u32;

        // Prove no truncation
        assert_eq!(cast_result as usize, total);
        assert!(total <= u32::MAX as usize);
    }

    // ── Step 7: Message Size Bounds ────────────────────────────────────────

    /// Proves that Transport::receive_message checks frame_length against
    /// MAX_MESSAGE_SIZE before allocating. This prevents OOM DoS attacks.
    #[kani::proof]
    fn proof_transport_receive_bounded() {
        let length_bytes: [u8; 4] = kani::any();
        let frame_length = u32::from_be_bytes(length_bytes) as usize;

        // Simulate the bounds check from line 141
        if frame_length > MAX_MESSAGE_SIZE {
            // Would return Err — no allocation occurs
            assert!(frame_length > MAX_MESSAGE_SIZE);
        } else {
            // Allocation is bounded to at most MAX_MESSAGE_SIZE (1 MB)
            assert!(frame_length <= MAX_MESSAGE_SIZE);
            assert!(frame_length <= 1024 * 1024);
        }
    }

    /// Proves the send_raw `data.len() as u32` cast in handshake
    /// never truncates for messages within MAX_MESSAGE_SIZE.
    #[kani::proof]
    fn proof_handshake_send_raw_no_truncation() {
        let data_len: usize = kani::any();
        kani::assume(data_len <= MAX_MESSAGE_SIZE); // 1 MB

        let cast_result = data_len as u32;

        // MAX_MESSAGE_SIZE = 1048576, which fits in u32 (max 4294967295)
        assert_eq!(cast_result as usize, data_len);
    }
}