wsb-daemon 0.275.7622

//! The language server protocol's framing: `Content-Length` headers, a blank line, then one
//! JSON message of exactly that many bytes.

use serde_json::Value;
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt};

/// Largest message read from a server: a whole worktree's symbols fit; a runaway does not.
pub const MAX_MESSAGE: usize = 64 * 1024 * 1024;

#[derive(Debug)]
pub enum FramingError {
    Io(std::io::Error),
    NoLength,
    Header(String),
    TooLarge(usize),
    Truncated,
    Json(serde_json::Error),
}

impl std::fmt::Display for FramingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Io(e) => write!(f, "reading the server's output: {e}"),
            Self::NoLength => write!(f, "a message header ended without Content-Length"),
            Self::Header(h) => write!(f, "invalid header {h:?}"),
            Self::TooLarge(n) => write!(f, "a message of {n} bytes is over the {MAX_MESSAGE}-byte limit"),
            Self::Truncated => write!(f, "the output ended in the middle of a message"),
            Self::Json(e) => write!(f, "a message is not JSON: {e}"),
        }
    }
}

impl From<std::io::Error> for FramingError {
    fn from(e: std::io::Error) -> Self {
        Self::Io(e)
    }
}

impl From<serde_json::Error> for FramingError {
    fn from(e: serde_json::Error) -> Self {
        Self::Json(e)
    }
}

/// The next message, or `None` when the output ended between messages.
pub async fn read<R: AsyncBufRead + Unpin>(r: &mut R) -> Result<Option<Value>, FramingError> {
    let mut length: Option<usize> = None;
    let mut first = true;
    loop {
        let mut line = String::new();
        if r.read_line(&mut line).await? == 0 {
            return if first { Ok(None) } else { Err(FramingError::Truncated) };
        }
        first = false;
        let line = line.trim_end_matches(['\r', '\n']);
        if line.is_empty() {
            break;
        }
        let (name, value) = line.split_once(':').ok_or_else(|| FramingError::Header(line.to_owned()))?;
        if name.trim().eq_ignore_ascii_case("content-length") {
            let n: usize = value.trim().parse().map_err(|_| FramingError::Header(line.to_owned()))?;
            if n > MAX_MESSAGE {
                return Err(FramingError::TooLarge(n));
            }
            length = Some(n);
        }
    }
    let length = length.ok_or(FramingError::NoLength)?;
    let mut body = vec![0u8; length];
    r.read_exact(&mut body).await.map_err(|e| match e.kind() {
        std::io::ErrorKind::UnexpectedEof => FramingError::Truncated,
        _ => FramingError::Io(e),
    })?;
    Ok(Some(serde_json::from_slice(&body)?))
}

/// Write one message and flush it.
pub async fn write<W: AsyncWrite + Unpin>(w: &mut W, message: &Value) -> std::io::Result<()> {
    let body = serde_json::to_vec(message).map_err(std::io::Error::other)?;
    w.write_all(format!("Content-Length: {}\r\n\r\n", body.len()).as_bytes()).await?;
    w.write_all(&body).await?;
    w.flush().await
}

#[cfg(test)]
#[path = "framing_tests.rs"]
mod tests;