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

//! Stdio and loopback WebSocket clients for the versioned protocol.

use std::ffi::OsStr;
use std::fmt;
use std::io::{self, BufRead, BufReader, BufWriter, Write};
use std::net::{Shutdown, SocketAddr, TcpStream};
use std::process::{Child, ChildStdin, ChildStdout, Command, ExitStatus, Stdio};
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_LINE_BYTES: usize = 1024 * 1024;

pub struct LocalClient {
    child: Child,
    client: Option<ProtocolClient<BufReader<ChildStdout>, BufWriter<ChildStdin>>>,
}

pub struct ProtocolClient<R, W> {
    reader: R,
    writer: W,
    next_request_id: u64,
}

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_LINE_BYTES))
            .max_frame_size(Some(MAX_MESSAGE_LINE_BYTES))
            .max_write_buffer_size(2 * MAX_MESSAGE_LINE_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_LINE_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,
    ResponseTooLarge,
    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_LINE_BYTES} bytes"
            ),
            Self::ResponseTooLarge => write!(
                formatter,
                "protocol response exceeds {MAX_MESSAGE_LINE_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)
    }
}

impl LocalClient {
    pub fn spawn(program: impl AsRef<OsStr>, args: &[impl AsRef<OsStr>]) -> io::Result<Self> {
        let mut child = Command::new(program)
            .args(args)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .spawn()?;
        let stdin = child.stdin.take().ok_or_else(|| {
            io::Error::other("child stdin was not available after enabling piped stdin")
        })?;
        let stdout = child.stdout.take().ok_or_else(|| {
            io::Error::other("child stdout was not available after enabling piped stdout")
        })?;

        Ok(Self {
            child,
            client: Some(ProtocolClient::new(
                BufReader::new(stdout),
                BufWriter::new(stdin),
            )),
        })
    }

    pub fn request(&mut self, command: ProtocolCommand) -> Result<Response, ClientError> {
        self.client
            .as_mut()
            .ok_or_else(|| {
                ClientError::Io(io::Error::new(
                    io::ErrorKind::BrokenPipe,
                    "client is shut down",
                ))
            })?
            .request(command)
    }

    pub fn shutdown(mut self) -> io::Result<ExitStatus> {
        self.client.take();
        self.child.wait()
    }
}

impl<R: BufRead, W: Write> ProtocolClient<R, W> {
    pub fn new(reader: R, writer: W) -> Self {
        Self {
            reader,
            writer,
            next_request_id: 0,
        }
    }

    pub fn request(&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 request = Request {
            version: PROTOCOL_VERSION,
            request_id: request_id.clone(),
            command,
        };
        let encoded = serde_json::to_vec(&request)?;
        if encoded.len() > MAX_MESSAGE_LINE_BYTES {
            return Err(ClientError::RequestTooLarge);
        }

        self.writer.write_all(&encoded)?;
        self.writer.write_all(b"\n")?;
        self.writer.flush()?;

        let mut line = Vec::new();
        let has_line = read_bounded_line(&mut self.reader, &mut line).map_err(|error| {
            if error.kind() == io::ErrorKind::InvalidData {
                ClientError::ResponseTooLarge
            } else {
                ClientError::Io(error)
            }
        })?;
        if !has_line {
            return Err(ClientError::Io(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "server closed the connection before responding",
            )));
        }
        let response: Response = serde_json::from_slice(&line)?;
        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,
            });
        }
        Ok(response)
    }
}

impl Drop for LocalClient {
    fn drop(&mut self) {
        if self.child.try_wait().ok().flatten().is_none() {
            let _ = self.child.kill();
            let _ = self.child.wait();
        }
    }
}

fn read_bounded_line<R: BufRead>(reader: &mut R, line: &mut Vec<u8>) -> io::Result<bool> {
    line.clear();
    loop {
        let available = reader.fill_buf()?;
        if available.is_empty() {
            return Ok(!line.is_empty());
        }

        let newline = available.iter().position(|byte| *byte == b'\n');
        let payload_len = newline.unwrap_or(available.len());
        if line.len().saturating_add(payload_len) > MAX_MESSAGE_LINE_BYTES {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("response line exceeds {MAX_MESSAGE_LINE_BYTES} bytes"),
            ));
        }

        line.extend_from_slice(&available[..payload_len]);
        let consumed = payload_len + usize::from(newline.is_some());
        reader.consume(consumed);
        if newline.is_some() {
            return Ok(true);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use runa_protocol::{ProtocolError, Reply};
    use std::io::Cursor;

    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();
    }

    fn response(request_id: &str) -> Vec<u8> {
        let response = Response {
            version: PROTOCOL_VERSION,
            request_id: request_id.to_owned(),
            result: Ok(Reply::Acknowledged),
        };
        let mut bytes = serde_json::to_vec(&response).unwrap();
        bytes.push(b'\n');
        bytes
    }

    #[test]
    fn client_writes_correlated_requests_and_validates_responses() {
        let mut written = Vec::new();
        let mut client = ProtocolClient::new(Cursor::new(response("runa-client-1")), &mut written);
        let result = client.request(ProtocolCommand::GetBehaviors).unwrap();
        assert_eq!(result.request_id, "runa-client-1");
        assert_eq!(result.result, Ok(Reply::Acknowledged));

        let request: Request = serde_json::from_slice(
            written
                .strip_suffix(b"\n")
                .expect("client writes a newline"),
        )
        .unwrap();
        assert_eq!(request.request_id, "runa-client-1");
        assert_eq!(request.command, ProtocolCommand::GetBehaviors);
    }

    #[test]
    fn client_rejects_mismatched_response_ids_and_versions() {
        let mut client = ProtocolClient::new(Cursor::new(response("wrong-id")), Vec::new());
        assert!(matches!(
            client.request(ProtocolCommand::GetBehaviors),
            Err(ClientError::CorrelationMismatch { .. })
        ));

        let mut unsupported = serde_json::to_vec(&Response {
            version: PROTOCOL_VERSION + 1,
            request_id: "runa-client-1".to_owned(),
            result: Err(ProtocolError::new(
                runa_protocol::ErrorCode::UnsupportedVersion,
                "unsupported",
            )),
        })
        .unwrap();
        unsupported.push(b'\n');
        let mut client = ProtocolClient::new(Cursor::new(unsupported), Vec::new());
        assert!(matches!(
            client.request(ProtocolCommand::GetBehaviors),
            Err(ClientError::UnsupportedVersion(version)) if version == PROTOCOL_VERSION + 1
        ));
    }

    #[test]
    fn client_enforces_request_and_response_size_limits() {
        let mut client = ProtocolClient::new(Cursor::new(response("runa-client-1")), Vec::new());
        let error = client
            .request(ProtocolCommand::PublishBehavior {
                name: "large".to_owned(),
                expected_revision: None,
                source: "x".repeat(MAX_MESSAGE_LINE_BYTES),
            })
            .unwrap_err();
        assert!(matches!(error, ClientError::RequestTooLarge));

        let oversized_response = vec![b'x'; MAX_MESSAGE_LINE_BYTES + 1];
        let mut client = ProtocolClient::new(Cursor::new(oversized_response), Vec::new());
        let error = client.request(ProtocolCommand::GetBehaviors).unwrap_err();
        assert!(matches!(error, ClientError::ResponseTooLarge));
    }
}