quicburn-shared 0.1.0

// crates/quicburn-shared/src/frame.rs

use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};
use std::fmt;

// ============================================================
// Status Codes
// ============================================================

/// Response status codes for frame messages
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Status {
    /// Success - operation completed successfully
    Ok = 0,
    /// Generic error - payload contains error message
    Error = 1,
    /// Resource not found - payload may contain identifier
    NotFound = 2,
    /// Authentication required - client must authenticate
    Unauthorized = 3,
    /// Forbidden - authenticated but not authorized
    Forbidden = 4,
    /// Invalid request - malformed or invalid parameters
    BadRequest = 5,
    /// Internal server error
    InternalError = 6,
    /// Request timeout
    Timeout = 7,
    /// Rate limited - too many requests
    RateLimited = 8,
    /// Feature not implemented
    NotImplemented = 9,
    /// Service unavailable
    ServiceUnavailable = 10,
    /// Conflict - resource already exists or state conflict
    Conflict = 11,
    /// Too large - request exceeds maximum size
    TooLarge = 12,
}

impl Status {
    /// Parse a status from a raw byte
    pub fn from_byte(b: u8) -> Result<Self> {
        match b {
            0 => Ok(Self::Ok),
            1 => Ok(Self::Error),
            2 => Ok(Self::NotFound),
            3 => Ok(Self::Unauthorized),
            4 => Ok(Self::Forbidden),
            5 => Ok(Self::BadRequest),
            6 => Ok(Self::InternalError),
            7 => Ok(Self::Timeout),
            8 => Ok(Self::RateLimited),
            9 => Ok(Self::NotImplemented),
            10 => Ok(Self::ServiceUnavailable),
            11 => Ok(Self::Conflict),
            12 => Ok(Self::TooLarge),
            other => bail!("unknown status byte: {other}"),
        }
    }

    /// Get the byte representation. Status stays a single wire byte —
    /// only `route` grew to u32, not this.
    pub const fn as_byte(self) -> u8 {
        self as u8
    }

    /// Get the status name for logging
    pub fn name(self) -> &'static str {
        match self {
            Self::Ok => "OK",
            Self::Error => "ERROR",
            Self::NotFound => "NOT_FOUND",
            Self::Unauthorized => "UNAUTHORIZED",
            Self::Forbidden => "FORBIDDEN",
            Self::BadRequest => "BAD_REQUEST",
            Self::InternalError => "INTERNAL_ERROR",
            Self::Timeout => "TIMEOUT",
            Self::RateLimited => "RATE_LIMITED",
            Self::NotImplemented => "NOT_IMPLEMENTED",
            Self::ServiceUnavailable => "SERVICE_UNAVAILABLE",
            Self::Conflict => "CONFLICT",
            Self::TooLarge => "TOO_LARGE",
        }
    }

    /// Check if this status is a success
    pub fn is_success(self) -> bool {
        matches!(self, Self::Ok)
    }

    /// Check if this status is an error
    pub fn is_error(self) -> bool {
        !self.is_success()
    }

    /// Convert a non-Ok status into a descriptive error
    pub fn into_error(self, payload: &[u8]) -> anyhow::Error {
        let message = match self {
            Self::Ok => return anyhow::anyhow!("into_error called on Status::Ok"),
            // For payload-based errors, try to get a clean string
            Self::NotFound => {
                let s = String::from_utf8_lossy(payload);
                let trimmed = s.trim();
                if trimmed.is_empty() {
                    "not found".to_string()
                } else {
                    trimmed.to_string()
                }
            }
            Self::BadRequest | Self::Conflict | Self::Error => {
                let s = String::from_utf8_lossy(payload);
                let trimmed = s.trim();
                if trimmed.is_empty() {
                    format!("{}", self.name().to_lowercase())
                } else {
                    trimmed.to_string()
                }
            }
            // Static errors - no payload needed
            Self::Unauthorized => "not authenticated — call login() first".to_string(),
            Self::Forbidden => "access denied — insufficient permissions".to_string(),
            Self::InternalError => "internal server error".to_string(),
            Self::Timeout => "request timeout".to_string(),
            Self::RateLimited => "rate limit exceeded".to_string(),
            Self::NotImplemented => "feature not implemented".to_string(),
            Self::ServiceUnavailable => "service temporarily unavailable".to_string(),
            Self::TooLarge => "request too large".to_string(),
        };

        anyhow::anyhow!("{}: {}", self.name(), message)
    }
}

impl fmt::Display for Status {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

impl From<Status> for u8 {
    fn from(status: Status) -> Self {
        status as u8
    }
}

// ============================================================
// Frame
// ============================================================

/// A framed message: route (u32) + status byte + u32 length + payload.
///
/// Wire layout, big-endian: [route: 4 bytes][status: 1 byte][len: 4 bytes][payload: len bytes]
/// 9-byte header total.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Frame {
    pub route: u32,
    pub status: Status,
    pub payload: Vec<u8>,
}

/// Size of the fixed frame header: 4 (route) + 1 (status) + 4 (length)
pub const FRAME_HEADER_SIZE: usize = 4 + 1 + 4;

impl Frame {
    /// Create a new frame with a payload
    pub fn new(route: u32, status: Status, payload: impl Into<Vec<u8>>) -> Self {
        Self { route, status, payload: payload.into() }
    }

    /// Create a success frame with empty payload
    pub fn ok(route: u32) -> Self {
        Self { route, status: Status::Ok, payload: Vec::new() }
    }

    /// Create an error frame with a message
    pub fn error(route: u32, message: impl Into<String>) -> Self {
        Self {
            route,
            status: Status::Error,
            payload: message.into().into_bytes(),
        }
    }

    /// Get the payload as a string (if UTF-8)
    pub fn payload_str(&self) -> &str {
        std::str::from_utf8(&self.payload).unwrap_or("invalid UTF-8")
    }

    /// Check if this frame is a success
    pub fn is_ok(&self) -> bool {
        self.status.is_success()
    }

    /// Check if this frame is an error
    pub fn is_error(&self) -> bool {
        self.status.is_error()
    }

    /// Convert error frames into an anyhow::Error
    pub fn into_result(self) -> Result<Self> {
        if self.is_error() {
            Err(self.status.into_error(&self.payload))
        } else {
            Ok(self)
        }
    }

    /// Encode this frame into bytes
    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(FRAME_HEADER_SIZE + self.payload.len());
        buf.extend_from_slice(&self.route.to_be_bytes());
        buf.push(self.status.as_byte());
        buf.extend_from_slice(&(self.payload.len() as u32).to_be_bytes());
        buf.extend_from_slice(&self.payload);
        buf
    }

    /// Decode a frame from bytes
    pub fn decode(data: &[u8]) -> Result<Self> {
        if data.len() < FRAME_HEADER_SIZE {
            bail!(
                "frame too short: {} bytes (minimum {})",
                data.len(),
                FRAME_HEADER_SIZE
            );
        }

        let route = u32::from_be_bytes(data[0..4].try_into().unwrap());
        let status = Status::from_byte(data[4])?;

        let len_bytes: [u8; 4] = data[5..9]
            .try_into()
            .map_err(|_| anyhow::anyhow!("invalid length bytes"))?;
        let len = u32::from_be_bytes(len_bytes) as usize;

        if data.len() < FRAME_HEADER_SIZE + len {
            bail!(
                "frame truncated: expected {} bytes, got {}",
                FRAME_HEADER_SIZE + len,
                data.len()
            );
        }

        let payload = data[FRAME_HEADER_SIZE..FRAME_HEADER_SIZE + len].to_vec();

        Ok(Self { route, status, payload })
    }

    /// Get the approximate size of this frame (header + payload)
    pub fn frame_size(&self) -> usize {
        FRAME_HEADER_SIZE + self.payload.len()
    }
}

// ============================================================
// Frame Reader/Writer
// ============================================================

/// Maximum frame size (16 MB)
pub const MAX_FRAME_SIZE: usize = 16 * 1024 * 1024;

/// Write a framed message to a QUIC send stream
pub async fn write_frame(
    send: &mut quinn::SendStream,
    route: u32,
    status: Status,
    payload: &[u8],
) -> Result<()> {
    // Validate payload size
    if payload.len() > MAX_FRAME_SIZE {
        bail!(
            "payload too large: {} bytes (max {})",
            payload.len(),
            MAX_FRAME_SIZE
        );
    }

    let mut frame = Vec::with_capacity(FRAME_HEADER_SIZE + payload.len());
    frame.extend_from_slice(&route.to_be_bytes());
    frame.push(status.as_byte());
    frame.extend_from_slice(&(payload.len() as u32).to_be_bytes());
    frame.extend_from_slice(payload);

    send.write_all(&frame)
        .await
        .context("failed to write frame to stream")?;
    Ok(())
}

/// Read a framed message from a QUIC receive stream
pub async fn read_frame(recv: &mut quinn::RecvStream) -> Result<Option<Frame>> {
    // Read route (4 bytes, big-endian u32)
    let mut route_buf = [0u8; 4];
    match recv.read_exact(&mut route_buf).await {
        Ok(()) => {}
        Err(quinn::ReadExactError::FinishedEarly(0)) => return Ok(None),
        Err(e) => return Err(e).context("failed to read route bytes"),
    }
    let route = u32::from_be_bytes(route_buf);

    // Read status byte (still a single byte)
    let mut status_buf = [0u8; 1];
    recv.read_exact(&mut status_buf)
        .await
        .context("failed to read status byte")?;
    let status = Status::from_byte(status_buf[0])?;

    // Read length prefix (4 bytes, big-endian u32)
    let mut len_buf = [0u8; 4];
    recv.read_exact(&mut len_buf)
        .await
        .context("failed to read frame length prefix")?;
    let len = u32::from_be_bytes(len_buf) as usize;

    // Validate length
    if len > MAX_FRAME_SIZE {
        bail!("frame too large: {len} bytes (max {MAX_FRAME_SIZE})");
    }

    // Read payload
    let mut payload = vec![0u8; len];
    recv.read_exact(&mut payload)
        .await
        .context("failed to read frame payload")?;

    Ok(Some(Frame { route, status, payload }))
}

/// Convenience function to read a frame and convert errors
pub async fn read_frame_result(recv: &mut quinn::RecvStream) -> Result<Frame> {
    match read_frame(recv).await? {
        Some(frame) => frame.into_result(),
        None => bail!("connection closed without frame"),
    }
}

// ============================================================
// Tests
// ============================================================

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

    #[test]
    fn frame_encoding_round_trip() {
        let original = Frame::new(0x01, Status::Ok, b"hello world".to_vec());
        let encoded = original.encode();
        let decoded = Frame::decode(&encoded).unwrap();
        assert_eq!(original, decoded);
    }

    #[test]
    fn frame_encoding_has_correct_length_prefix() {
        let payload = b"hello world";
        let frame = Frame::new(0x01, Status::Ok, payload.to_vec());
        let encoded = frame.encode();

        assert_eq!(u32::from_be_bytes(encoded[0..4].try_into().unwrap()), 0x01);
        assert_eq!(encoded[4], Status::Ok.as_byte());
        let len_bytes: [u8; 4] = encoded[5..9].try_into().unwrap();
        let len = u32::from_be_bytes(len_bytes) as usize;
        assert_eq!(len, payload.len());
        assert_eq!(&encoded[9..], payload);
    }

    #[test]
    fn empty_payload_encodes_zero_length() {
        let frame = Frame::new(0x02, Status::NotFound, Vec::new());
        let encoded = frame.encode();
        let len_bytes: [u8; 4] = encoded[5..9].try_into().unwrap();
        assert_eq!(u32::from_be_bytes(len_bytes), 0);
        assert_eq!(encoded.len(), FRAME_HEADER_SIZE);
    }

    #[test]
    fn status_round_trips() {
        for status in [
            Status::Ok,
            Status::Error,
            Status::NotFound,
            Status::Unauthorized,
            Status::Forbidden,
            Status::BadRequest,
            Status::InternalError,
            Status::Timeout,
            Status::RateLimited,
            Status::NotImplemented,
            Status::ServiceUnavailable,
            Status::Conflict,
            Status::TooLarge,
        ] {
            let byte = status.as_byte();
            let recovered = Status::from_byte(byte).expect("valid status byte");
            assert_eq!(status, recovered);
            assert_eq!(status.name(), recovered.name());
        }
    }

    #[test]
    fn unknown_status_byte_is_rejected() {
        assert!(Status::from_byte(99).is_err());
    }

    #[test]
    fn frame_into_result_ok() {
        let frame = Frame::ok(0x01);
        let result = frame.into_result();
        assert!(result.is_ok());
    }

    #[test]
    fn frame_into_result_error() {
        let frame = Frame::error(0x01, "something went wrong");
        let result = frame.into_result();
        assert!(result.is_err());
    }

    #[test]
    fn error_status_converts_to_message() {
        let payload = b"custom error message";
        let err = Status::Error.into_error(payload);
        assert_eq!(err.to_string(), "ERROR: custom error message");
    }

    #[test]
    fn frame_size_calculation() {
        let frame = Frame::new(0x01, Status::Ok, b"hello".to_vec());
        assert_eq!(frame.frame_size(), FRAME_HEADER_SIZE + 5); // header + payload
    }

    #[test]
    fn decode_truncated_frame_fails() {
        // header claims a 5-byte payload but none is present
        let mut data = vec![0u8; FRAME_HEADER_SIZE];
        data[5..9].copy_from_slice(&5u32.to_be_bytes());
        assert!(Frame::decode(&data).is_err());
    }

    #[test]
    fn decode_too_short_fails() {
        let data = vec![0x01, 0x00]; // Too short
        assert!(Frame::decode(&data).is_err());
    }
}