skyzen 0.2.1

//! WebSocket types for Durable Object hibernation.

#[cfg(feature = "ws")]
use std::sync::Arc;

#[cfg(feature = "ws")]
use http_kit::error::BoxHttpError;
#[cfg(feature = "ws")]
use http_kit::http_error;
use http_kit::ws::WebSocketMessage;
use serde::{de::DeserializeOwned, Serialize};
#[cfg(feature = "ws")]
use skyzen_core::Responder;
#[cfg(all(feature = "ws", target_arch = "wasm32"))]
use wasm_bindgen::JsCast;

use super::error::DurableObjectError;

/// A hibernation WebSocket event delivered to [`DurableObject::websocket`].
///
/// [`DurableObject::websocket`]: super::DurableObject::websocket
#[derive(Debug)]
pub enum WebSocketEvent {
    /// A message was received from the client.
    Message(WebSocketMessage),
    /// The connection was closed by the client.
    Close {
        /// The close code.
        code: u16,
        /// The close reason.
        reason: String,
        /// Whether the close was clean.
        was_clean: bool,
    },
    /// An error occurred on the connection.
    Error(String),
}

/// Opaque handle to a hibernating WebSocket connection.
///
/// This wraps the platform-specific WebSocket object and provides
/// methods to send messages, manage tags, and handle attachments.
///
/// The inner representation is platform-specific and set by the runtime glue.
pub struct WebSocketConnection {
    /// Platform-specific inner handle, set by the cloudflare glue layer.
    inner: Box<dyn WebSocketConnectionInner>,
}

impl std::fmt::Debug for WebSocketConnection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WebSocketConnection")
            .finish_non_exhaustive()
    }
}

/// Platform-specific WebSocket connection operations.
///
/// Implemented by the cloudflare glue layer (or test mocks).
pub trait WebSocketConnectionInner: Send + Sync {
    /// Send a text message.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] when the runtime send fails.
    fn send_text(&self, text: &str) -> Result<(), DurableObjectError>;
    /// Send binary data.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] when the runtime send fails.
    fn send_binary(&self, data: &[u8]) -> Result<(), DurableObjectError>;
    /// Close the connection.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] when the runtime close fails.
    fn close(&self, code: u16, reason: &str) -> Result<(), DurableObjectError>;
    /// Get tags for this connection.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] when tag lookup fails.
    fn tags(&self) -> Result<Vec<String>, DurableObjectError>;
    /// Get the typed attachment.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] when attachment retrieval fails.
    fn get_attachment_raw(&self) -> Result<Option<Vec<u8>>, DurableObjectError>;
    /// Set the typed attachment.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] when attachment persistence fails.
    fn set_attachment_raw(&self, data: &[u8]) -> Result<(), DurableObjectError>;
}

impl WebSocketConnection {
    /// Create from a platform-specific inner handle.
    #[must_use]
    pub fn new(inner: Box<dyn WebSocketConnectionInner>) -> Self {
        Self { inner }
    }

    /// Send a text message to this connection.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] if the send fails.
    pub fn send_text(&self, text: &str) -> Result<(), DurableObjectError> {
        self.inner.send_text(text)
    }

    /// Send binary data to this connection.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] if the send fails.
    pub fn send_binary(&self, data: &[u8]) -> Result<(), DurableObjectError> {
        self.inner.send_binary(data)
    }

    /// Serialize a value to JSON and send it as text.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] if serialization or send fails.
    pub fn send_json<T: Serialize>(&self, value: &T) -> Result<(), DurableObjectError> {
        let json = serde_json::to_string(value)
            .map_err(|e| DurableObjectError::Serialization(e.to_string()))?;
        self.send_text(&json)
    }

    /// Close this connection with a code and reason.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] if the close fails.
    pub fn close(&self, code: u16, reason: &str) -> Result<(), DurableObjectError> {
        self.inner.close(code, reason)
    }

    /// Get the tags associated with this connection.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] if the operation fails.
    pub fn tags(&self) -> Result<Vec<String>, DurableObjectError> {
        self.inner.tags()
    }

    /// Get the per-connection typed attachment (survives hibernation).
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] if deserialization fails.
    pub fn attachment<T: DeserializeOwned>(&self) -> Result<Option<T>, DurableObjectError> {
        self.inner
            .get_attachment_raw()?
            .map(|bytes| {
                serde_json::from_slice(&bytes)
                    .map_err(|e| DurableObjectError::Serialization(e.to_string()))
            })
            .transpose()
    }

    /// Set the per-connection typed attachment (survives hibernation).
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] if serialization fails.
    pub fn set_attachment<T: Serialize>(&self, value: &T) -> Result<(), DurableObjectError> {
        let bytes = serde_json::to_vec(value)
            .map_err(|e| DurableObjectError::Serialization(e.to_string()))?;
        self.inner.set_attachment_raw(&bytes)
    }
}

/// Responder that accepts a WebSocket for Hibernation.
///
/// Construct this in a handler, add tags, and return it.
/// During `respond_to()`, the runtime reads `DurableObjectState` from
/// request extensions to call `acceptWebSocket(ws, tags)`.
#[derive(Debug, Default)]
pub struct HibernationWebSocketUpgrade {
    tags: Vec<String>,
    #[cfg(feature = "ws")]
    subprotocol: Subprotocol,
}

/// What a hibernation handshake answers the client's subprotocol offer with.
#[cfg(feature = "ws")]
#[derive(Debug, Default)]
enum Subprotocol {
    /// The client's offer goes unanswered — correct only when it made none.
    #[default]
    Unanswered,
    /// Answer with the first of these the client offered.
    Negotiated(Vec<String>),
    /// Answer with exactly this value.
    Exact(http_kit::header::HeaderValue),
}

#[cfg(feature = "ws")]
impl Subprotocol {
    /// The value to echo on the `101`, given what the client offered.
    fn answer(&self, request: &crate::Request) -> Option<http_kit::header::HeaderValue> {
        match self {
            Self::Unanswered => None,
            Self::Negotiated(supported) => {
                crate::websocket::select_offered_protocol(request.headers(), supported)
            }
            Self::Exact(protocol) => Some(protocol.clone()),
        }
    }
}

#[cfg(all(feature = "ws", not(target_arch = "wasm32")))]
type NativeDurableObjectWebSocketAcceptFn = dyn Fn(
        crate::websocket::WebSocket,
        Vec<String>,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>
    + Send
    + Sync;

/// Native simulator adapter that hands an upgraded connection to its Durable Object instance.
#[cfg(all(feature = "ws", not(target_arch = "wasm32")))]
#[derive(Clone)]
pub struct NativeDurableObjectState {
    accept_websocket: Arc<NativeDurableObjectWebSocketAcceptFn>,
}

#[cfg(all(feature = "ws", not(target_arch = "wasm32")))]
impl std::fmt::Debug for NativeDurableObjectState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NativeDurableObjectState")
            .finish_non_exhaustive()
    }
}

#[cfg(all(feature = "ws", not(target_arch = "wasm32")))]
impl NativeDurableObjectState {
    pub(crate) fn new<F, Fut>(accept_websocket: F) -> Self
    where
        F: Fn(crate::websocket::WebSocket, Vec<String>) -> Fut + Send + Sync + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        Self {
            accept_websocket: Arc::new(move |websocket, tags| {
                Box::pin(accept_websocket(websocket, tags))
            }),
        }
    }

    fn accept_websocket(
        &self,
        websocket: crate::websocket::WebSocket,
        tags: Vec<String>,
    ) -> impl std::future::Future<Output = ()> + Send {
        (self.accept_websocket)(websocket, tags)
    }
}

impl HibernationWebSocketUpgrade {
    /// Create a new hibernation WebSocket upgrade.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a tag to categorize this connection.
    #[must_use]
    pub fn tag(mut self, tag: impl Into<String>) -> Self {
        self.tags.push(tag.into());
        self
    }

    /// Get the tags that will be applied to this connection.
    #[must_use]
    pub fn tags(&self) -> &[String] {
        &self.tags
    }

    /// Consume and return the tags.
    #[must_use]
    pub fn into_tags(self) -> Vec<String> {
        self.tags
    }

    /// Answer the handshake with the first of `protocols` the client also offered.
    ///
    /// RFC 6455 §4.1 requires the echo whenever the client offered a subprotocol: without it the
    /// client fails the connection, so a hibernating socket that ignores the offer never opens.
    #[cfg(feature = "ws")]
    #[must_use]
    pub fn protocols<I, S>(mut self, protocols: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        self.subprotocol = Subprotocol::Negotiated(
            protocols
                .into_iter()
                .map(|protocol| protocol.as_ref().to_owned())
                .collect(),
        );
        self
    }

    /// Answer the handshake with exactly this subprotocol, whatever the client offered.
    ///
    /// This is how a browser authenticates a Durable Object socket. It cannot attach an
    /// `Authorization` header to a `WebSocket`, so the credential travels in the subprotocol list
    /// (`new WebSocket(url, ["app.bearer." + token])`) and the server echoes the accepted token
    /// back — which the client requires before it will open the socket.
    ///
    /// A [`HeaderValue`](http_kit::header::HeaderValue) rather than a string, because a
    /// subprotocol that cannot be sent as a header value is not an answer, and the handler is
    /// where that is known.
    #[cfg(feature = "ws")]
    #[must_use]
    pub fn protocol(mut self, protocol: http_kit::header::HeaderValue) -> Self {
        self.subprotocol = Subprotocol::Exact(protocol);
        self
    }
}

#[cfg(all(feature = "ws", target_arch = "wasm32"))]
type DurableObjectWebSocketAcceptFn =
    dyn Fn(&web_sys::WebSocket, &[String]) -> Result<(), DurableObjectError> + Send + Sync;

/// Wrapper for client websocket stored in response extensions.
#[cfg(all(feature = "ws", target_arch = "wasm32"))]
#[derive(Clone)]
pub struct DurableClientWebSocket(pub web_sys::WebSocket);

#[cfg(all(feature = "ws", target_arch = "wasm32"))]
impl std::fmt::Debug for DurableClientWebSocket {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DurableClientWebSocket")
            .finish_non_exhaustive()
    }
}

// SAFETY: WebAssembly workers are single-threaded; JS handles are safe to mark Send/Sync.
#[cfg(all(feature = "ws", target_arch = "wasm32"))]
unsafe impl Send for DurableClientWebSocket {}
#[cfg(all(feature = "ws", target_arch = "wasm32"))]
unsafe impl Sync for DurableClientWebSocket {}

/// Internal extension value injected by Durable Object runtime glue.
///
/// `HibernationWebSocketUpgrade` reads this from request extensions and
/// calls `acceptWebSocket(server, tags)` through it.
#[cfg(all(feature = "ws", target_arch = "wasm32"))]
#[derive(Clone)]
pub struct DurableObjectState {
    accept_websocket: Arc<DurableObjectWebSocketAcceptFn>,
}

#[cfg(all(feature = "ws", target_arch = "wasm32"))]
impl std::fmt::Debug for DurableObjectState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DurableObjectState").finish_non_exhaustive()
    }
}

#[cfg(all(feature = "ws", target_arch = "wasm32"))]
impl DurableObjectState {
    /// Create a state adapter used by `HibernationWebSocketUpgrade`.
    #[must_use]
    pub fn new(
        accept_websocket: impl Fn(&web_sys::WebSocket, &[String]) -> Result<(), DurableObjectError>
            + Send
            + Sync
            + 'static,
    ) -> Self {
        Self {
            accept_websocket: Arc::new(accept_websocket),
        }
    }

    /// Accept a server websocket with optional tags.
    ///
    /// # Errors
    ///
    /// Returns [`DurableObjectError`] when runtime websocket acceptance fails.
    pub fn accept_websocket(
        &self,
        websocket: &web_sys::WebSocket,
        tags: &[String],
    ) -> Result<(), DurableObjectError> {
        (self.accept_websocket)(websocket, tags)
    }
}

#[cfg(all(feature = "ws", target_arch = "wasm32"))]
http_error!(
    HibernationDurableStateMissing,
    http_kit::StatusCode::INTERNAL_SERVER_ERROR,
    "DurableObjectState missing in request extensions for hibernation websocket upgrade."
);

#[cfg(all(feature = "ws", not(target_arch = "wasm32")))]
http_error!(
    NativeHibernationDurableStateMissing,
    http_kit::StatusCode::INTERNAL_SERVER_ERROR,
    "NativeDurableObjectState missing in request extensions for hibernation websocket upgrade."
);

#[cfg(all(feature = "ws", not(target_arch = "wasm32")))]
impl Responder for HibernationWebSocketUpgrade {
    type Error = BoxHttpError;

    fn respond_to(
        self,
        request: &crate::Request,
        response: &mut crate::Response,
    ) -> Result<(), Self::Error> {
        let state = request
            .extensions()
            .get::<NativeDurableObjectState>()
            .cloned()
            .ok_or_else(|| Box::new(NativeHibernationDurableStateMissing::new()) as BoxHttpError)?;
        let mut upgrade = crate::websocket::upgrade_from_request(request)
            .map_err(|error| Box::new(error) as BoxHttpError)?;
        if let Some(protocol) = self.subprotocol.answer(request) {
            upgrade = upgrade.protocol(protocol);
        }
        let tags = self.into_tags();
        let responder =
            upgrade.on_upgrade(move |websocket| state.accept_websocket(websocket, tags));
        responder
            .respond_to(request, response)
            .map_err(|error| Box::new(error) as BoxHttpError)
    }
}

#[cfg(all(feature = "ws", target_arch = "wasm32"))]
http_error!(
    HibernationWebSocketAcceptFailed,
    http_kit::StatusCode::INTERNAL_SERVER_ERROR,
    "Failed to accept hibernation websocket connection."
);

#[cfg(all(feature = "ws", target_arch = "wasm32"))]
impl Responder for HibernationWebSocketUpgrade {
    type Error = BoxHttpError;

    fn respond_to(
        self,
        request: &crate::Request,
        response: &mut crate::Response,
    ) -> Result<(), Self::Error> {
        let durable_state = request
            .extensions()
            .get::<DurableObjectState>()
            .cloned()
            .ok_or_else(|| Box::new(HibernationDurableStateMissing::new()) as BoxHttpError)?;

        let pair = crate::websocket::ffi::WebSocketPair::new();
        let client = pair.client();
        let server = pair.server();
        let server_socket: web_sys::WebSocket = server.unchecked_into();

        durable_state
            .accept_websocket(&server_socket, self.tags())
            .map_err(|error| {
                tracing::error!(%error, "failed to accept hibernation websocket");
                Box::new(HibernationWebSocketAcceptFailed::new()) as BoxHttpError
            })?;

        *response.status_mut() = http_kit::StatusCode::SWITCHING_PROTOCOLS;

        // The answer to the client's subprotocol offer is an ordinary response header; the Worker
        // runtime carries the whole header map onto the `101` it builds for the host.
        if let Some(protocol) = self.subprotocol.answer(request) {
            response
                .headers_mut()
                .insert(http_kit::header::SEC_WEBSOCKET_PROTOCOL, protocol);
        }

        let client_socket: web_sys::WebSocket = client.unchecked_into();
        response
            .extensions_mut()
            .insert(DurableClientWebSocket(client_socket));

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use serde::{Deserialize, Serialize};

    use super::{
        DurableObjectError, HibernationWebSocketUpgrade, WebSocketConnection,
        WebSocketConnectionInner,
    };

    #[derive(Debug, Default)]
    struct MockSocketState {
        text_messages: Vec<String>,
        binary_messages: Vec<Vec<u8>>,
        close_calls: Vec<(u16, String)>,
        tags: Vec<String>,
        attachment: Option<Vec<u8>>,
        fail_text: Option<DurableObjectError>,
        fail_binary: Option<DurableObjectError>,
        fail_close: Option<DurableObjectError>,
        fail_tags: Option<DurableObjectError>,
        fail_get_attachment: Option<DurableObjectError>,
        fail_set_attachment: Option<DurableObjectError>,
    }

    #[derive(Debug, Clone)]
    struct MockSocketInner {
        state: Arc<Mutex<MockSocketState>>,
    }

    impl MockSocketInner {
        fn error_clone(error: &DurableObjectError) -> DurableObjectError {
            match error {
                DurableObjectError::Runtime(message) => {
                    DurableObjectError::Runtime(message.clone())
                }
                DurableObjectError::Serialization(message) => {
                    DurableObjectError::Serialization(message.clone())
                }
                DurableObjectError::WebSocket(message) => {
                    DurableObjectError::WebSocket(message.clone())
                }
            }
        }
    }

    impl WebSocketConnectionInner for MockSocketInner {
        fn send_text(&self, text: &str) -> Result<(), DurableObjectError> {
            {
                let mut state = self.state.lock().unwrap();
                if let Some(error) = &state.fail_text {
                    return Err(Self::error_clone(error));
                }
                state.text_messages.push(text.to_owned());
            }
            Ok(())
        }

        fn send_binary(&self, data: &[u8]) -> Result<(), DurableObjectError> {
            {
                let mut state = self.state.lock().unwrap();
                if let Some(error) = &state.fail_binary {
                    return Err(Self::error_clone(error));
                }
                state.binary_messages.push(data.to_vec());
            }
            Ok(())
        }

        fn close(&self, code: u16, reason: &str) -> Result<(), DurableObjectError> {
            {
                let mut state = self.state.lock().unwrap();
                if let Some(error) = &state.fail_close {
                    return Err(Self::error_clone(error));
                }
                state.close_calls.push((code, reason.to_owned()));
            }
            Ok(())
        }

        fn tags(&self) -> Result<Vec<String>, DurableObjectError> {
            let state = self.state.lock().unwrap();
            if let Some(error) = &state.fail_tags {
                return Err(Self::error_clone(error));
            }
            Ok(state.tags.clone())
        }

        fn get_attachment_raw(&self) -> Result<Option<Vec<u8>>, DurableObjectError> {
            let state = self.state.lock().unwrap();
            if let Some(error) = &state.fail_get_attachment {
                return Err(Self::error_clone(error));
            }
            Ok(state.attachment.clone())
        }

        fn set_attachment_raw(&self, data: &[u8]) -> Result<(), DurableObjectError> {
            {
                let mut state = self.state.lock().unwrap();
                if let Some(error) = &state.fail_set_attachment {
                    return Err(Self::error_clone(error));
                }
                state.attachment = Some(data.to_vec());
            }
            Ok(())
        }
    }

    #[derive(Debug, Serialize, Deserialize, PartialEq)]
    struct Attachment {
        room: String,
        revision: u32,
    }

    #[derive(Debug)]
    struct AlwaysFails;

    impl Serialize for AlwaysFails {
        fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
        where
            S: serde::Serializer,
        {
            Err(serde::ser::Error::custom("expected serialization failure"))
        }
    }

    fn connection_with_state(state: Arc<Mutex<MockSocketState>>) -> WebSocketConnection {
        WebSocketConnection::new(Box::new(MockSocketInner { state }))
    }

    #[test]
    fn send_json_serializes_payload_as_text_message() {
        let state = Arc::new(Mutex::new(MockSocketState::default()));
        let connection = connection_with_state(Arc::clone(&state));

        connection
            .send_json(&Attachment {
                room: "room-1".to_owned(),
                revision: 7,
            })
            .unwrap();

        assert_eq!(
            state.lock().unwrap().text_messages,
            vec![r#"{"room":"room-1","revision":7}"#.to_owned()]
        );
    }

    #[test]
    fn send_json_rejects_non_finite_numbers() {
        let connection = connection_with_state(Arc::new(Mutex::new(MockSocketState::default())));

        let error = connection.send_json(&AlwaysFails).unwrap_err();

        assert!(matches!(error, DurableObjectError::Serialization(_)));
    }

    #[test]
    fn attachment_round_trips_through_raw_storage() {
        let state = Arc::new(Mutex::new(MockSocketState::default()));
        let connection = connection_with_state(Arc::clone(&state));
        let attachment = Attachment {
            room: "general".to_owned(),
            revision: 3,
        };

        connection.set_attachment(&attachment).unwrap();
        let restored = connection.attachment::<Attachment>().unwrap().unwrap();

        assert_eq!(restored, attachment);
        assert!(state.lock().unwrap().attachment.is_some());
    }

    #[test]
    fn attachment_reports_deserialization_errors_for_invalid_bytes() {
        let state = Arc::new(Mutex::new(MockSocketState {
            attachment: Some(b"{".to_vec()),
            ..MockSocketState::default()
        }));
        let connection = connection_with_state(state);

        let error = connection.attachment::<Attachment>().unwrap_err();

        assert!(matches!(error, DurableObjectError::Serialization(_)));
    }

    #[test]
    fn hibernation_upgrade_preserves_tag_order() {
        let tags = HibernationWebSocketUpgrade::new()
            .tag("room-1")
            .tag("presence")
            .into_tags();

        assert_eq!(tags, vec!["room-1".to_owned(), "presence".to_owned()]);
    }
}