clusterflux-public/crates/clusterflux-control/src/lib.rs

#[cfg(not(target_arch = "wasm32"))]
use std::io::{BufRead, BufReader, Read, Write};
#[cfg(not(target_arch = "wasm32"))]
use std::net::TcpStream;
use std::net::{IpAddr, ToSocketAddrs};
use std::time::Duration;

use serde_json::Value;
use thiserror::Error;

pub const CONTROL_API_PATH: &str = "/api/v1/control";
pub const MAX_CONTROL_FRAME_BYTES: usize = 1024 * 1024;

#[derive(Debug, Error)]
pub enum ControlTransportError {
    #[error("invalid coordinator endpoint: {0}")]
    InvalidEndpoint(String),
    #[error("insecure remote coordinator endpoint is forbidden: {0}")]
    InsecureRemote(String),
    #[error("coordinator transport I/O failed: {0}")]
    Io(#[from] std::io::Error),
    #[error("coordinator transport JSON failed: {0}")]
    Json(#[from] serde_json::Error),
    #[error("coordinator HTTP request failed: {0}")]
    Http(String),
    #[error("coordinator control frame exceeds {MAX_CONTROL_FRAME_BYTES} bytes")]
    FrameTooLarge,
    #[error("coordinator closed the local control session without a response")]
    Closed,
    #[error("coordinator network transport is unavailable inside a Wasm guest")]
    UnavailableInWasm,
}

enum ControlTransport {
    #[cfg(not(target_arch = "wasm32"))]
    Https { agent: ureq::Agent, url: String },
    #[cfg(not(target_arch = "wasm32"))]
    LoopbackJsonLine {
        writer: TcpStream,
        reader: BufReader<TcpStream>,
    },
    #[cfg(target_arch = "wasm32")]
    #[allow(dead_code)]
    Unavailable,
}

pub struct ControlSession {
    transport: ControlTransport,
    requests: u64,
}

impl ControlSession {
    pub fn connect(endpoint: &str) -> Result<Self, ControlTransportError> {
        Self::connect_with_timeouts(endpoint, Duration::from_secs(10), Duration::from_secs(30))
    }

    pub fn connect_with_timeouts(
        endpoint: &str,
        connect_timeout: Duration,
        io_timeout: Duration,
    ) -> Result<Self, ControlTransportError> {
        #[cfg(target_arch = "wasm32")]
        {
            let _ = (endpoint, connect_timeout, io_timeout);
            return Err(ControlTransportError::UnavailableInWasm);
        }

        #[cfg(not(target_arch = "wasm32"))]
        {
            let endpoint = endpoint.trim();
            if endpoint.starts_with("https://") || endpoint.starts_with("http://") {
                let url = control_api_url(endpoint)?;
                if endpoint.starts_with("http://") && !endpoint_is_loopback(endpoint) {
                    return Err(ControlTransportError::InsecureRemote(endpoint.to_owned()));
                }
                let agent = ureq::AgentBuilder::new()
                    .timeout_connect(connect_timeout)
                    .timeout_read(io_timeout)
                    .timeout_write(io_timeout)
                    .build();
                return Ok(Self {
                    transport: ControlTransport::Https { agent, url },
                    requests: 0,
                });
            }

            let loopback_address = endpoint
                .strip_prefix("clusterflux+tcp://")
                .unwrap_or(endpoint);
            if !endpoint_is_loopback(loopback_address) {
                return Err(ControlTransportError::InsecureRemote(endpoint.to_owned()));
            }
            let writer = TcpStream::connect(loopback_address)?;
            writer.set_read_timeout(Some(io_timeout))?;
            writer.set_write_timeout(Some(io_timeout))?;
            let reader = BufReader::new(writer.try_clone()?);
            Ok(Self {
                transport: ControlTransport::LoopbackJsonLine { writer, reader },
                requests: 0,
            })
        }
    }

    pub fn request(&mut self, value: &Value) -> Result<Value, ControlTransportError> {
        #[cfg(target_arch = "wasm32")]
        {
            let _ = (&self.transport, self.requests, value);
            return Err(ControlTransportError::UnavailableInWasm);
        }

        #[cfg(not(target_arch = "wasm32"))]
        {
            let encoded = serde_json::to_vec(value)?;
            if encoded.len() > MAX_CONTROL_FRAME_BYTES {
                return Err(ControlTransportError::FrameTooLarge);
            }
            let response = match &mut self.transport {
                #[cfg(not(target_arch = "wasm32"))]
                ControlTransport::Https { agent, url } => {
                    let response = agent
                        .post(url)
                        .set("Content-Type", "application/json")
                        .set("Accept", "application/json")
                        .send_bytes(&encoded)
                        .map_err(|error| ControlTransportError::Http(error.to_string()))?;
                    if response.status() != 200 {
                        return Err(ControlTransportError::Http(format!(
                            "coordinator returned HTTP {} {}",
                            response.status(),
                            response.status_text()
                        )));
                    }
                    let mut bytes = Vec::new();
                    response
                        .into_reader()
                        .take((MAX_CONTROL_FRAME_BYTES + 1) as u64)
                        .read_to_end(&mut bytes)?;
                    if bytes.len() > MAX_CONTROL_FRAME_BYTES {
                        return Err(ControlTransportError::FrameTooLarge);
                    }
                    serde_json::from_slice(&bytes)?
                }
                ControlTransport::LoopbackJsonLine { writer, reader } => {
                    writer.write_all(&encoded)?;
                    writer.write_all(b"\n")?;
                    writer.flush()?;
                    let mut bytes = Vec::new();
                    reader
                        .take((MAX_CONTROL_FRAME_BYTES + 1) as u64)
                        .read_until(b'\n', &mut bytes)?;
                    if bytes.is_empty() {
                        return Err(ControlTransportError::Closed);
                    }
                    if bytes.len() > MAX_CONTROL_FRAME_BYTES {
                        return Err(ControlTransportError::FrameTooLarge);
                    }
                    serde_json::from_slice(&bytes)?
                }
            };
            self.requests += 1;
            Ok(response)
        }
    }

    pub fn requests(&self) -> u64 {
        self.requests
    }
}

pub fn control_api_url(endpoint: &str) -> Result<String, ControlTransportError> {
    let endpoint = endpoint.trim().trim_end_matches('/');
    if !(endpoint.starts_with("https://") || endpoint.starts_with("http://")) {
        return Err(ControlTransportError::InvalidEndpoint(endpoint.to_owned()));
    }
    if endpoint.ends_with(CONTROL_API_PATH) {
        Ok(endpoint.to_owned())
    } else {
        Ok(format!("{endpoint}{CONTROL_API_PATH}"))
    }
}

pub fn endpoint_identity(endpoint: &str) -> Result<String, ControlTransportError> {
    let endpoint = endpoint.trim();
    if endpoint.starts_with("https://") || endpoint.starts_with("http://") {
        if endpoint.starts_with("http://") && !endpoint_is_loopback(endpoint) {
            return Err(ControlTransportError::InsecureRemote(endpoint.to_owned()));
        }
        return control_api_url(endpoint);
    }
    let loopback_address = endpoint
        .strip_prefix("clusterflux+tcp://")
        .unwrap_or(endpoint);
    if endpoint_is_loopback(loopback_address) {
        return Ok(format!("clusterflux+tcp://{loopback_address}"));
    }
    Err(ControlTransportError::InsecureRemote(endpoint.to_owned()))
}

pub fn endpoint_is_loopback(endpoint: &str) -> bool {
    let authority = endpoint
        .trim()
        .strip_prefix("https://")
        .or_else(|| endpoint.trim().strip_prefix("http://"))
        .or_else(|| endpoint.trim().strip_prefix("clusterflux+tcp://"))
        .unwrap_or(endpoint.trim())
        .split('/')
        .next()
        .unwrap_or_default();
    let host = if authority.starts_with('[') {
        authority
            .strip_prefix('[')
            .and_then(|value| value.split_once(']'))
            .map(|(host, _)| host)
            .unwrap_or(authority)
    } else {
        authority
            .rsplit_once(':')
            .map(|(host, _)| host)
            .unwrap_or(authority)
    };
    if host.eq_ignore_ascii_case("localhost") {
        return true;
    }
    if host
        .parse::<IpAddr>()
        .is_ok_and(|address| address.is_loopback())
    {
        return true;
    }
    authority
        .to_socket_addrs()
        .ok()
        .is_some_and(|mut addresses| addresses.all(|address| address.ip().is_loopback()))
}

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

    #[test]
    fn hosted_endpoints_are_real_https_api_urls() {
        assert_eq!(
            control_api_url("https://clusterflux.example").unwrap(),
            "https://clusterflux.example/api/v1/control"
        );
        assert_eq!(
            endpoint_identity("https://clusterflux.example/api/v1/control").unwrap(),
            "https://clusterflux.example/api/v1/control"
        );
    }

    #[test]
    fn plaintext_transport_is_restricted_to_loopback() {
        assert!(endpoint_is_loopback("127.0.0.1:7999"));
        assert!(endpoint_is_loopback("clusterflux+tcp://127.0.0.1:7999"));
        assert!(endpoint_is_loopback("http://[::1]:7999"));
        assert!(matches!(
            ControlSession::connect("http://example.com:7999"),
            Err(ControlTransportError::InsecureRemote(_))
        ));
        assert!(matches!(
            ControlSession::connect("example.com:7999"),
            Err(ControlTransportError::InsecureRemote(_))
        ));
    }
}