flaron-sdk 1.0.0

//! WebSocket primitives for flares that handle WebSocket upgrades.
//!
//! A WebSocket-enabled flare exports `ws_open`, `ws_message`, and `ws_close`
//! instead of (or in addition to) `handle_request`. Inside those exports the
//! event accessors return data the host wired up before the call:
//!
//! * [`event_type`] - `"open"`, `"message"`, or `"close"`.
//! * [`event_data`] - the message payload (binary or UTF-8 text bytes).
//! * [`conn_id`]    - the per-connection identifier the host issues.
//! * [`close_code`] - close code from the peer (only meaningful in
//!   `ws_close`).
//!
//! Use [`send`] to push frames to the peer and [`close`] to terminate the
//! connection from the flare side.

use crate::{ffi, mem};

/// Errors returned by [`send`] when the host refuses the frame.
#[derive(Debug, thiserror::Error)]
pub enum WsSendError {
    /// Per-invocation send-rate limit reached. Drop the frame and back off.
    #[error("ws send: per-invocation rate limit reached")]
    SendLimitReached,

    /// Frame exceeds the host's per-message size cap.
    #[error("ws send: message too large")]
    MessageTooLarge,

    /// Generic send failure (connection closed, write error, etc.).
    #[error("ws send: write failed")]
    SendError,

    /// Unknown error code returned by the host.
    #[error("ws send: unknown error code {0}")]
    Unknown(i32),
}

impl WsSendError {
    fn from_code(code: i32) -> Self {
        match code {
            1 => Self::SendLimitReached,
            2 => Self::MessageTooLarge,
            3 => Self::SendError,
            _ => Self::Unknown(code),
        }
    }
}

/// Send a frame to the connected peer.
///
/// The host treats the bytes as opaque - pass UTF-8 for a text frame or
/// arbitrary bytes for a binary frame.
pub fn send(data: &[u8]) -> Result<(), WsSendError> {
    let (data_ptr, data_len) = mem::host_arg_bytes(data);
    let code = unsafe { ffi::ws_send(data_ptr, data_len) };
    if code == 0 {
        Ok(())
    } else {
        Err(WsSendError::from_code(code))
    }
}

/// Convenience: send a text frame.
pub fn send_text(text: &str) -> Result<(), WsSendError> {
    send(text.as_bytes())
}

/// Close the connection with the given WebSocket status code.
///
/// Common codes:
/// * `1000` - Normal Closure
/// * `1001` - Going Away
/// * `1002` - Protocol Error
/// * `1008` - Policy Violation
/// * `1011` - Internal Error
pub fn close(code: u16) {
    unsafe { ffi::ws_close_conn(code as i32) }
}

/// Per-connection identifier issued by the host. Stable for the lifetime of
/// a single WebSocket connection.
pub fn conn_id() -> String {
    // SAFETY: host writes a valid UTF-8 connection ID into the bump arena.
    unsafe { mem::read_packed_string(ffi::ws_conn_id()) }.unwrap_or_default()
}

/// Type of the current WebSocket event: `"open"`, `"message"`, or `"close"`.
pub fn event_type() -> String {
    // SAFETY: host writes a valid UTF-8 event type into the bump arena.
    unsafe { mem::read_packed_string(ffi::ws_event_type()) }.unwrap_or_default()
}

/// Payload bytes for a `"message"` event. Returns an empty `Vec<u8>` for
/// `"open"` and `"close"` events.
pub fn event_data() -> Vec<u8> {
    // SAFETY: host writes the event payload bytes into the bump arena.
    unsafe { mem::read_packed_bytes(ffi::ws_event_data()) }.unwrap_or_default()
}

/// Convenience: event payload interpreted as a UTF-8 string. Invalid UTF-8
/// is replaced with the Unicode replacement character.
pub fn event_text() -> String {
    String::from_utf8_lossy(&event_data()).into_owned()
}

/// Close code provided by the remote peer (only meaningful inside a
/// `ws_close` handler - `0` otherwise).
pub fn close_code() -> u16 {
    unsafe { ffi::ws_close_code() as u16 }
}

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

    #[test]
    fn send_records_payload() {
        test_host::reset();
        send(b"hello").unwrap();
        let sends = test_host::read_mock(|m| m.ws_sends.clone());
        assert_eq!(sends, vec![b"hello".to_vec()]);
    }

    #[test]
    fn send_text_records_utf8_bytes() {
        test_host::reset();
        send_text("héllo").unwrap();
        let sends = test_host::read_mock(|m| m.ws_sends.clone());
        assert_eq!(sends, vec!["héllo".as_bytes().to_vec()]);
    }

    #[test]
    fn send_maps_error_codes() {
        for (code, expected) in [
            (1, WsSendError::SendLimitReached),
            (2, WsSendError::MessageTooLarge),
            (3, WsSendError::SendError),
        ] {
            test_host::reset();
            test_host::with_mock(|m| m.ws_send_error = code);
            let err = send(b"x").unwrap_err();
            assert!(
                std::mem::discriminant(&err) == std::mem::discriminant(&expected),
                "code {} mismatch",
                code,
            );
        }
    }

    #[test]
    fn send_unknown_error_code() {
        test_host::reset();
        test_host::with_mock(|m| m.ws_send_error = 99);
        match send(b"x").unwrap_err() {
            WsSendError::Unknown(99) => {}
            other => panic!("expected Unknown(99), got {:?}", other),
        }
    }

    #[test]
    fn close_records_code() {
        test_host::reset();
        close(1000);
        close(1011);
        assert_eq!(
            test_host::read_mock(|m| m.ws_closes.clone()),
            vec![1000, 1011]
        );
    }

    #[test]
    fn conn_id_returns_host_value() {
        test_host::reset();
        test_host::with_mock(|m| m.ws_conn_id = Some("conn-abc-123".into()));
        assert_eq!(conn_id(), "conn-abc-123");
    }

    #[test]
    fn conn_id_empty_when_unset() {
        test_host::reset();
        assert_eq!(conn_id(), "");
    }

    #[test]
    fn event_type_open_message_close() {
        for ty in ["open", "message", "close"] {
            test_host::reset();
            test_host::with_mock(|m| m.ws_event_type = Some(ty.into()));
            assert_eq!(event_type(), ty);
        }
    }

    #[test]
    fn event_data_returns_payload_bytes() {
        test_host::reset();
        test_host::with_mock(|m| m.ws_event_data = Some(vec![1, 2, 3, 4]));
        assert_eq!(event_data(), vec![1, 2, 3, 4]);
    }

    #[test]
    fn event_data_empty_for_open_close() {
        test_host::reset();
        assert!(event_data().is_empty());
    }

    #[test]
    fn event_text_decodes_utf8() {
        test_host::reset();
        test_host::with_mock(|m| m.ws_event_data = Some("héllo".as_bytes().to_vec()));
        assert_eq!(event_text(), "héllo");
    }

    #[test]
    fn close_code_returns_host_value() {
        test_host::reset();
        test_host::with_mock(|m| m.ws_close_code = 1006);
        assert_eq!(close_code(), 1006);
    }
}