runa/crates/runa-client/src/lib.rs

//! Loopback WebSocket client for the versioned protocol.

use std::fmt;
use std::io;
use std::net::{Shutdown, SocketAddr, TcpStream};
use std::time::{Duration, Instant};

use runa_protocol::{
    Command as ProtocolCommand, ImageEvent, Notification, PROTOCOL_VERSION, Request, Response,
    WEBSOCKET_SUBPROTOCOL,
};
use tungstenite::client::IntoClientRequest;
use tungstenite::{Message, WebSocket, protocol::WebSocketConfig};

pub const MAX_MESSAGE_BYTES: usize = 1024 * 1024;

pub struct WebSocketClient {
    socket: WebSocket<TcpStream>,
    next_request_id: u64,
    revision: u64,
    disconnected: bool,
}

impl WebSocketClient {
    /// Connect only to loopback, without authentication or automatic retries.
    pub fn connect(address: SocketAddr) -> io::Result<Self> {
        if !address.ip().is_loopback() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "server address must be a loopback address",
            ));
        }
        let stream = TcpStream::connect_timeout(&address, Duration::from_secs(5))?;
        stream.set_read_timeout(Some(Duration::from_secs(5)))?;
        stream.set_write_timeout(Some(Duration::from_secs(5)))?;
        let mut request = format!("ws://{address}/")
            .into_client_request()
            .map_err(io::Error::other)?;
        request.headers_mut().insert(
            "sec-websocket-protocol",
            WEBSOCKET_SUBPROTOCOL.parse().map_err(io::Error::other)?,
        );
        let config = WebSocketConfig::default()
            .max_message_size(Some(MAX_MESSAGE_BYTES))
            .max_frame_size(Some(MAX_MESSAGE_BYTES))
            .max_write_buffer_size(2 * MAX_MESSAGE_BYTES);
        let (socket, response) =
            tungstenite::client::client_with_config(request, stream, Some(config))
                .map_err(io::Error::other)?;
        if response
            .headers()
            .get("sec-websocket-protocol")
            .and_then(|value| value.to_str().ok())
            != Some(WEBSOCKET_SUBPROTOCOL)
        {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "server did not negotiate runa-debug",
            ));
        }
        socket
            .get_ref()
            .set_read_timeout(Some(Duration::from_secs(30)))?;
        socket
            .get_ref()
            .set_write_timeout(Some(Duration::from_secs(30)))?;
        Ok(Self {
            socket,
            next_request_id: 0,
            revision: 0,
            disconnected: false,
        })
    }

    pub fn request(&mut self, command: ProtocolCommand) -> Result<Response, ClientError> {
        if self.disconnected {
            return Err(io::Error::new(
                io::ErrorKind::NotConnected,
                "WebSocket client is disconnected; reconnect explicitly",
            )
            .into());
        }
        let result = self.request_inner(command);
        if result.is_err() {
            self.disconnected = true;
            // A lost response may follow a committed operation. Never retry it
            // or reuse a stream whose request/response alignment is uncertain.
            let _ = self.socket.get_ref().shutdown(Shutdown::Both);
        }
        result
    }

    fn request_inner(&mut self, command: ProtocolCommand) -> Result<Response, ClientError> {
        self.next_request_id = self
            .next_request_id
            .checked_add(1)
            .ok_or(ClientError::RequestIdsExhausted)?;
        let request_id = format!("runa-client-{}", self.next_request_id);
        let encoded = serde_json::to_string(&Request {
            version: PROTOCOL_VERSION,
            request_id: request_id.clone(),
            command,
        })?;
        if encoded.len() > MAX_MESSAGE_BYTES {
            return Err(ClientError::RequestTooLarge);
        }
        self.socket
            .send(Message::Text(encoded.into()))
            .map_err(io::Error::other)?;
        let deadline = Instant::now() + Duration::from_secs(30);
        loop {
            let remaining = deadline
                .checked_duration_since(Instant::now())
                .filter(|duration| !duration.is_zero())
                .ok_or_else(|| {
                    io::Error::new(io::ErrorKind::TimedOut, "WebSocket response timed out")
                })?;
            self.socket.get_ref().set_read_timeout(Some(remaining))?;
            match self.socket.read().map_err(io::Error::other)? {
                Message::Text(text) => {
                    let value: serde_json::Value = serde_json::from_str(&text)?;
                    if value.get("event").is_some() {
                        let event: Notification = serde_json::from_value(value)?;
                        if event.version != PROTOCOL_VERSION {
                            return Err(ClientError::UnsupportedVersion(event.version));
                        }
                        let ImageEvent::ImageChanged { revision } = event.event;
                        if revision <= self.revision {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                "non-increasing image notification revision",
                            )
                            .into());
                        }
                        self.revision = revision;
                        continue;
                    }
                    let response: Response = serde_json::from_value(value)?;
                    if response.version != PROTOCOL_VERSION {
                        return Err(ClientError::UnsupportedVersion(response.version));
                    }
                    if response.request_id != request_id {
                        return Err(ClientError::CorrelationMismatch {
                            expected: request_id,
                            actual: response.request_id,
                        });
                    }
                    return Ok(response);
                }
                Message::Ping(_) => self.socket.flush().map_err(io::Error::other)?,
                Message::Pong(_) => {}
                _ => {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        "expected a WebSocket text protocol message",
                    )
                    .into());
                }
            }
        }
    }
}

#[derive(Debug)]
pub enum ClientError {
    Io(io::Error),
    Serialization(serde_json::Error),
    RequestIdsExhausted,
    RequestTooLarge,
    UnsupportedVersion(u16),
    CorrelationMismatch { expected: String, actual: String },
}

impl fmt::Display for ClientError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Io(error) => write!(formatter, "protocol transport failed: {error}"),
            Self::Serialization(error) => write!(formatter, "protocol JSON failed: {error}"),
            Self::RequestIdsExhausted => write!(formatter, "protocol request IDs exhausted"),
            Self::RequestTooLarge => write!(
                formatter,
                "protocol request exceeds {MAX_MESSAGE_BYTES} bytes"
            ),
            Self::UnsupportedVersion(version) => {
                write!(
                    formatter,
                    "server returned unsupported protocol version {version}"
                )
            }
            Self::CorrelationMismatch { expected, actual } => write!(
                formatter,
                "response request ID mismatch: expected `{expected}`, got `{actual}`"
            ),
        }
    }
}

impl std::error::Error for ClientError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Io(error) => Some(error),
            Self::Serialization(error) => Some(error),
            _ => None,
        }
    }
}

impl From<io::Error> for ClientError {
    fn from(error: io::Error) -> Self {
        Self::Io(error)
    }
}

impl From<serde_json::Error> for ClientError {
    fn from(error: serde_json::Error) -> Self {
        Self::Serialization(error)
    }
}

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

    struct DebugProtocol;
    impl tungstenite::handshake::server::Callback for DebugProtocol {
        fn on_request(
            self,
            _: &tungstenite::handshake::server::Request,
            mut response: tungstenite::handshake::server::Response,
        ) -> Result<
            tungstenite::handshake::server::Response,
            tungstenite::handshake::server::ErrorResponse,
        > {
            response.headers_mut().insert(
                "sec-websocket-protocol",
                WEBSOCKET_SUBPROTOCOL.parse().unwrap(),
            );
            Ok(response)
        }
    }

    fn websocket_peer(
        handler: impl FnOnce(WebSocket<TcpStream>, Request) + Send + 'static,
    ) -> (SocketAddr, std::thread::JoinHandle<()>) {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let address = listener.local_addr().unwrap();
        let peer = std::thread::spawn(move || {
            let (stream, _) = listener.accept().unwrap();
            let mut socket = tungstenite::accept_hdr(stream, DebugProtocol).unwrap();
            let message = socket.read().unwrap();
            let request = serde_json::from_str(message.to_text().unwrap()).unwrap();
            handler(socket, request);
        });
        (address, peer)
    }

    #[test]
    fn websocket_client_validates_notifications_without_losing_response_correlation() {
        let (address, peer) = websocket_peer(|mut socket, request| {
            socket
                .send(Message::Text(
                    serde_json::to_string(&Notification {
                        version: PROTOCOL_VERSION,
                        event: ImageEvent::ImageChanged { revision: 42 },
                    })
                    .unwrap()
                    .into(),
                ))
                .unwrap();
            socket
                .send(Message::Text(
                    serde_json::to_string(&Response {
                        version: PROTOCOL_VERSION,
                        request_id: request.request_id,
                        result: Ok(Reply::Actors(Vec::new())),
                    })
                    .unwrap()
                    .into(),
                ))
                .unwrap();
        });
        let mut client = WebSocketClient::connect(address).unwrap();
        assert_eq!(
            client.request(ProtocolCommand::ListActors).unwrap().result,
            Ok(Reply::Actors(Vec::new()))
        );
        assert_eq!(client.revision, 42);
        peer.join().unwrap();
    }

    #[test]
    fn websocket_client_rejects_bad_versions_correlations_and_notifications() {
        for message in [
            serde_json::json!({"version": 2, "request_id": "runa-client-1",
                "result": {"Ok": {"kind": "actors", "data": []}}}),
            serde_json::json!({"version": 3, "request_id": "wrong",
                "result": {"Ok": {"kind": "actors", "data": []}}}),
            serde_json::json!({"version": 3, "event": {"event": "image_changed", "data": {"revision": 0}}}),
            serde_json::json!({"version": 3, "event": {"event": "unknown"}}),
        ] {
            let (address, peer) = websocket_peer(move |mut socket, _| {
                socket
                    .send(Message::Text(message.to_string().into()))
                    .unwrap();
            });
            let mut client = WebSocketClient::connect(address).unwrap();
            assert!(client.request(ProtocolCommand::ListActors).is_err());
            assert!(matches!(client.request(ProtocolCommand::ListActors),
                Err(ClientError::Io(error)) if error.kind() == io::ErrorKind::NotConnected));
            peer.join().unwrap();
        }
    }

    #[test]
    fn websocket_client_rejects_non_loopback_addresses() {
        let error = WebSocketClient::connect("192.0.2.1:3737".parse().unwrap())
            .err()
            .expect("non-loopback connection must be rejected");
        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn websocket_client_reports_disconnect_and_does_not_retry() {
        let (address, peer) = websocket_peer(|_socket, request| {
            assert_eq!(request.command, ProtocolCommand::ListActors);
            // Close after receiving the request, without sending a response.
        });
        let mut client = WebSocketClient::connect(address).unwrap();
        assert!(matches!(
            client.request(ProtocolCommand::ListActors),
            Err(ClientError::Io(_))
        ));
        assert!(matches!(client.request(ProtocolCommand::ListActors),
            Err(ClientError::Io(error)) if error.kind() == io::ErrorKind::NotConnected));
        peer.join().unwrap();
    }
}