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

//! Stdio and loopback TCP clients for the versioned JSON-lines 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;

use runa_protocol::{Command as ProtocolCommand, PROTOCOL_VERSION, Request, Response};

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 TcpClient {
    client: ProtocolClient<BufReader<TcpStream>, BufWriter<TcpStream>>,
    disconnected: bool,
}

impl TcpClient {
    /// 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(30)))?;
        stream.set_write_timeout(Some(Duration::from_secs(30)))?;
        Ok(Self {
            client: ProtocolClient::new(
                BufReader::new(stream.try_clone()?),
                BufWriter::new(stream),
            ),
            disconnected: false,
        })
    }

    pub fn request(&mut self, command: ProtocolCommand) -> Result<Response, ClientError> {
        if self.disconnected {
            return Err(io::Error::new(
                io::ErrorKind::NotConnected,
                "TCP client is disconnected; reconnect explicitly",
            )
            .into());
        }
        let result = self.client.request(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.client.writer.get_ref().shutdown(Shutdown::Both);
        }
        result
    }
}

#[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;

    #[test]
    fn tcp_client_rejects_non_loopback_addresses() {
        let error = TcpClient::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 tcp_client_reports_disconnect_and_does_not_retry() {
        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 reader = BufReader::new(stream);
            let mut line = String::new();
            reader.read_line(&mut line).unwrap();
            let request: Request = serde_json::from_str(&line).unwrap();
            assert_eq!(request.command, ProtocolCommand::ListActors);
            // Close after receiving the request, without sending a response.
        });
        let mut client = TcpClient::connect(address).unwrap();
        assert!(matches!(client.request(ProtocolCommand::ListActors),
            Err(ClientError::Io(error)) if error.kind() == io::ErrorKind::UnexpectedEof));
        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));
    }
}