lsp-max 26.7.3

//! Generic server for multiplexing bidirectional streams through a transport.

#[cfg(all(feature = "runtime-agnostic", not(feature = "runtime-tokio")))]
use async_codec_lite::{FramedRead, FramedWrite};
#[cfg(all(feature = "runtime-agnostic", not(feature = "runtime-tokio")))]
use futures::io::{AsyncRead, AsyncWrite};

#[cfg(feature = "runtime-tokio")]
use tokio::io::{AsyncRead, AsyncWrite};
#[cfg(feature = "runtime-tokio")]
use tokio_util::codec::{FramedRead, FramedWrite};

use std::task::Poll;

use futures::channel::mpsc;
use futures::{future, join, stream, FutureExt, Sink, SinkExt, Stream, StreamExt, TryFutureExt};
use tower::Service;
use tracing::error;

/// Drives `poll_ready` to completion; returns `Err` if the service has exited.
async fn await_ready<S>(service: &mut S) -> Result<(), S::Error>
where
    S: Service<Request>,
{
    future::poll_fn(|cx| service.poll_ready(cx)).await
}

/// Probes `poll_ready` once without blocking; returns `Some(err)` if the service
/// has exited, `None` if it is still pending or ready without error.
async fn probe_exited<S>(service: &mut S) -> Option<S::Error>
where
    S: Service<Request>,
{
    future::poll_fn(|cx| match service.poll_ready(cx) {
        Poll::Ready(Err(err)) => Poll::Ready(Some(err)),
        _ => Poll::Ready(None),
    })
    .await
}

use crate::codec::{LanguageServerCodec, ParseError};
use crate::jsonrpc::{Error, Id, Message, Request, Response};
use crate::service::{ClientSocket, ExitedError, RequestStream, ResponseSink};

const DEFAULT_MAX_CONCURRENCY: usize = 4;
const MESSAGE_QUEUE_SIZE: usize = 100;
/// Bounded re-polls after stdin EOF to observe the service's `Exited`
/// transition (driven by an in-flight `exit` call future) before falling
/// back to an abnormal-termination exit code.
const MAX_EXIT_OBSERVE_POLLS: usize = 16;

/// Trait implemented by client loopback sockets.
///
/// This socket handles the server-to-client half of the bidirectional communication stream.
///
/// See also: `examples/transport_utilities_explained.rs` — a run-to-exit witness that
/// demonstrates `Loopback`, `ExitedError`, and `ClientSocket` with real `assert!`s.
pub trait Loopback {
    /// Yields a stream of pending server-to-client requests.
    type RequestStream: Stream<Item = Request>;
    /// Routes client-to-server responses back to the server.
    type ResponseSink: Sink<Response> + Unpin;

    /// Splits this socket into two halves capable of operating independently.
    ///
    /// The two halves returned implement the [`Stream`] and [`Sink`] traits, respectively.
    fn split(self) -> (Self::RequestStream, Self::ResponseSink);
}

impl Loopback for ClientSocket {
    type RequestStream = RequestStream;
    type ResponseSink = ResponseSink;

    #[inline]
    fn split(self) -> (Self::RequestStream, Self::ResponseSink) {
        self.split()
    }
}

/// Server for processing requests and responses on standard I/O or TCP.
#[derive(Debug)]
pub struct Server<I, O, L = ClientSocket> {
    stdin: I,
    stdout: O,
    loopback: L,
    max_concurrency: usize,
}

impl<I, O, L> Server<I, O, L>
where
    I: AsyncRead + Unpin,
    O: AsyncWrite,
    L: Loopback,
    <L::ResponseSink as Sink<Response>>::Error: std::error::Error,
{
    /// Creates a new `Server` with the given `stdin` and `stdout` handles.
    pub fn new(stdin: I, stdout: O, socket: L) -> Self {
        Server {
            stdin,
            stdout,
            loopback: socket,
            max_concurrency: DEFAULT_MAX_CONCURRENCY,
        }
    }

    /// Sets the server concurrency limit to `max`.
    ///
    /// This setting specifies how many incoming requests may be processed concurrently. Setting
    /// this value to `1` forces all requests to be processed sequentially, thereby implicitly
    /// disabling support for the [`$/cancelRequest`] notification.
    ///
    /// [`$/cancelRequest`]: https://microsoft.github.io/language-server-protocol/specification#cancelRequest
    ///
    /// If not explicitly specified, `max` defaults to 4.
    ///
    /// # Preference over standard `tower` middleware
    ///
    /// The [`ConcurrencyLimit`] and [`Buffer`] middlewares provided by `tower` rely on
    /// [`tokio::spawn`] in common usage, while this library aims to be executor agnostic and to
    /// support exotic targets currently incompatible with `tokio`, such as WASM. As such, `Server`
    /// includes its own concurrency facilities that don't require a global executor to be present.
    ///
    /// [`ConcurrencyLimit`]: https://docs.rs/tower/latest/tower/limit/concurrency/struct.ConcurrencyLimit.html
    /// [`Buffer`]: https://docs.rs/tower/latest/tower/buffer/index.html
    /// [`tokio::spawn`]: https://docs.rs/tokio/latest/tokio/fn.spawn.html
    pub fn concurrency_level(mut self, max: usize) -> Self {
        self.max_concurrency = max;
        self
    }

    /// Spawns the service with messages read through `stdin` and responses written to `stdout`.
    pub async fn serve<T>(self, mut service: T) -> Result<(), T::Error>
    where
        T: Service<Request, Response = Option<Response>> + Send + 'static,
        T::Error: std::error::Error + Send + Sync + 'static + From<ExitedError>,
        T::Future: Send,
    {
        let (client_requests, mut client_responses) = self.loopback.split();
        let (client_requests, client_abort) = stream::abortable(client_requests);
        let (mut responses_tx, responses_rx) = mpsc::channel(MESSAGE_QUEUE_SIZE);
        let (mut server_tasks_tx, server_tasks_rx) = mpsc::channel(MESSAGE_QUEUE_SIZE);

        let mut framed_stdin = FramedRead::new(self.stdin, LanguageServerCodec::default());
        let framed_stdout = FramedWrite::new(self.stdout, LanguageServerCodec::default());

        let process_server_tasks = server_tasks_rx
            .buffer_unordered(self.max_concurrency)
            .filter_map(future::ready)
            .map(|res| Ok(Message::Response(res)))
            .forward(responses_tx.clone().sink_map_err(|_| unreachable!()))
            .map(|_| ());

        let print_output = stream::select(responses_rx, client_requests.map(Message::Request))
            .map(Ok)
            .forward(framed_stdout.sink_map_err(|e| error!("failed to encode message: {}", e)))
            .map(|_| ());

        let read_input = async {
            while let Some(msg) = framed_stdin.next().await {
                if let Ok(ref m) = msg {
                    match m {
                        Message::Request(req) => tracing::trace!(
                            "--- Server::serve read_input got Request: method={}, id={:?}",
                            req.method(),
                            req.id()
                        ),
                        Message::Response(res) => tracing::trace!(
                            "--- Server::serve read_input got Response: id={:?}",
                            res.id()
                        ),
                    }
                }
                match msg {
                    Ok(Message::Request(req)) => {
                        if let Err(err) = await_ready(&mut service).await {
                            error!("{}", display_sources(&err));
                            return Err(err);
                        }

                        let fut = service.call(req).unwrap_or_else(|err| {
                            error!("{}", display_sources(&err));
                            None
                        });

                        if let Err(e) = server_tasks_tx.send(fut).await {
                            error!("transport send failed: {}", e);
                            return Ok(());
                        }
                    }
                    Ok(Message::Response(res)) => {
                        if let Err(err) = client_responses.send(res).await {
                            error!("{}", display_sources(&err));
                            return Ok(());
                        }
                    }
                    Err(err) => {
                        error!("failed to decode message: {}", err);
                        let res = Response::from_error(Id::Null, to_jsonrpc_error(err));
                        if let Err(e) = responses_tx.send(Message::Response(res)).await {
                            error!("transport send failed: {}", e);
                            return Ok(());
                        }
                    }
                }
            }

            server_tasks_tx.disconnect();
            responses_tx.disconnect();
            client_abort.abort();

            // After stdin EOF, surface the state machine's STORED exit code.
            //
            // An `exit` notification dispatched just before EOF transitions the
            // service to `Exited` via its `call` future, at which point
            // `poll_ready` reports `Err(ExitedError(get_exit_code()))` carrying
            // the lawful code (0 on the ShutDown -> Exited path, 1 otherwise).
            // There is a window where the `exit` call future has not yet driven
            // the transition when EOF is first observed; polling once and falling
            // back to a hardcoded `1` would discard a lawful `0`. Yield and
            // re-poll a bounded number of times so the stored code is propagated
            // rather than a fixed fallback.
            for _ in 0..MAX_EXIT_OBSERVE_POLLS {
                if let Some(err) = probe_exited(&mut service).await {
                    return Err(err);
                }

                // Not yet observed as Exited; let any in-flight `exit` call
                // future make progress, then re-poll for the stored code.
                tokio::task::yield_now().await;
            }

            // No lawful Exited transition was observed (e.g. EOF without a prior
            // `exit`/`shutdown`): per LSP, this is an abnormal termination.
            Err(T::Error::from(ExitedError(1)))
        };

        let (_, res, _) = join!(print_output, read_input, process_server_tasks);
        res
    }
}

fn display_sources(error: &dyn std::error::Error) -> String {
    if let Some(source) = error.source() {
        format!("{}: {}", error, display_sources(source))
    } else {
        error.to_string()
    }
}

#[cfg(feature = "runtime-tokio")]
fn to_jsonrpc_error(err: ParseError) -> Error {
    match err {
        ParseError::Body(err) if err.is_data() => Error::invalid_request(),
        _ => Error::parse_error(),
    }
}

#[cfg(all(feature = "runtime-agnostic", not(feature = "runtime-tokio")))]
fn to_jsonrpc_error(err: impl std::error::Error) -> Error {
    match err.source().and_then(|e| e.downcast_ref()) {
        Some(ParseError::Body(err)) if err.is_data() => Error::invalid_request(),
        _ => Error::parse_error(),
    }
}

#[cfg(test)]
mod tests {
    use std::task::{Context, Poll};

    #[cfg(all(feature = "runtime-agnostic", not(feature = "runtime-tokio")))]
    use futures::io::Cursor;
    #[cfg(feature = "runtime-tokio")]
    use std::io::Cursor;

    use futures::future::Ready;
    use futures::{future, sink, stream};

    use super::*;

    const REQUEST: &str = r#"{"jsonrpc":"2.0","method":"initialize","params":{},"id":1}"#;
    const RESPONSE: &str = r#"{"jsonrpc":"2.0","result":{"capabilities":{}},"id":1}"#;

    #[derive(Debug)]
    struct MockService;

    impl Service<Request> for MockService {
        type Response = Option<Response>;
        type Error = ExitedError;
        type Future = Ready<Result<Self::Response, Self::Error>>;

        fn poll_ready(&mut self, _: &mut Context) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn call(&mut self, _: Request) -> Self::Future {
            let response = serde_json::from_str(RESPONSE).unwrap();
            future::ok(Some(response))
        }
    }

    struct MockLoopback(Vec<Request>);

    impl Loopback for MockLoopback {
        type RequestStream = stream::Iter<std::vec::IntoIter<Request>>;
        type ResponseSink = sink::Drain<Response>;

        fn split(self) -> (Self::RequestStream, Self::ResponseSink) {
            (stream::iter(self.0), sink::drain())
        }
    }

    fn mock_request() -> Vec<u8> {
        format!("Content-Length: {}\r\n\r\n{}", REQUEST.len(), REQUEST).into_bytes()
    }

    fn mock_response() -> Vec<u8> {
        format!("Content-Length: {}\r\n\r\n{}", RESPONSE.len(), RESPONSE).into_bytes()
    }

    fn mock_stdio() -> (Cursor<Vec<u8>>, Vec<u8>) {
        (Cursor::new(mock_request()), Vec::new())
    }

    #[tokio::test(flavor = "current_thread")]
    async fn serves_on_stdio() {
        let (mut stdin, mut stdout) = mock_stdio();
        let res = Server::new(&mut stdin, &mut stdout, MockLoopback(vec![]))
            .serve(MockService)
            .await;

        assert_eq!(res, Err(ExitedError(1)));
        assert_eq!(stdin.position(), 80);
        assert_eq!(stdout, mock_response());
    }

    #[tokio::test(flavor = "current_thread")]
    async fn interleaves_messages() {
        let socket = MockLoopback(vec![serde_json::from_str(REQUEST).unwrap()]);

        let (mut stdin, mut stdout) = mock_stdio();
        let res = Server::new(&mut stdin, &mut stdout, socket)
            .serve(MockService)
            .await;

        assert_eq!(res, Err(ExitedError(1)));
        assert_eq!(stdin.position(), 80);
        let output: Vec<_> = mock_request().into_iter().chain(mock_response()).collect();
        assert_eq!(stdout, output);
    }

    #[tokio::test(flavor = "current_thread")]
    async fn handles_invalid_json() {
        let invalid = r#"{"jsonrpc":"2.0","method":"#;
        let message = format!("Content-Length: {}\r\n\r\n{}", invalid.len(), invalid).into_bytes();
        let (mut stdin, mut stdout) = (Cursor::new(message), Vec::new());

        let res = Server::new(&mut stdin, &mut stdout, MockLoopback(vec![]))
            .serve(MockService)
            .await;

        assert_eq!(res, Err(ExitedError(1)));
        assert_eq!(stdin.position(), 48);
        let err = r#"{"jsonrpc":"2.0","error":{"code":-32700,"message":"Parse error"},"id":null}"#;
        let output = format!("Content-Length: {}\r\n\r\n{}", err.len(), err).into_bytes();
        assert_eq!(stdout, output);
    }

    /// Service that always returns None (no response), used to exercise the
    /// server_tasks channel send path without blocking on a response.
    #[derive(Debug)]
    struct NoneService;

    impl Service<Request> for NoneService {
        type Response = Option<Response>;
        type Error = ExitedError;
        type Future = Ready<Result<Self::Response, Self::Error>>;

        fn poll_ready(&mut self, _: &mut Context) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn call(&mut self, _: Request) -> Self::Future {
            future::ok(None)
        }
    }

    /// Proves that multiple valid requests sent through the transport do NOT panic when
    /// the server_tasks channel send is exercised — all complete without unwrap panics.
    /// This guards against the defect where server_tasks_tx.send().await.unwrap() would
    /// panic if the receiving half was dropped during a shutdown race.
    #[tokio::test(flavor = "current_thread")]
    async fn no_panic_on_server_tasks_channel_send() {
        let mut input = Vec::new();
        for _ in 0..5 {
            input.extend_from_slice(&mock_request());
        }
        let (mut stdin, mut stdout) = (Cursor::new(input), Vec::new());

        // NoneService returns None for every request so futures resolve immediately.
        let res = Server::new(&mut stdin, &mut stdout, MockLoopback(vec![]))
            .serve(NoneService)
            .await;

        // Server exits cleanly after EOF — no panic occurred.
        assert_eq!(res, Err(ExitedError(1)));
    }

    /// Proves that two consecutive invalid JSON messages both go through the responses_tx
    /// send path without panicking. This guards against the defect where
    /// responses_tx.send().await.unwrap() would panic if the receiver was dropped.
    #[tokio::test(flavor = "current_thread")]
    async fn no_panic_on_responses_tx_send_two_invalid_messages() {
        let invalid = r#"{"jsonrpc":"2.0","method":"#;
        let framed = format!("Content-Length: {}\r\n\r\n{}", invalid.len(), invalid).into_bytes();
        let mut input = framed.clone();
        input.extend_from_slice(&framed);

        let (mut stdin, mut stdout) = (Cursor::new(input), Vec::new());

        let res = Server::new(&mut stdin, &mut stdout, MockLoopback(vec![]))
            .serve(MockService)
            .await;

        // Both sends must complete without panic; result is Ok(()) or ExitedError.
        assert!(res == Ok(()) || res == Err(ExitedError(1)));
    }
}