fix: don't panic the LSP server on malformed client input

.github/workflows/tests.yaml

@@ -187,7 +187,9 @@ jobs:
         include:
           - target: x86_64-unknown-linux-gnu
             os: ubuntu-24.04
-            features: default,test-dbs-external
+            # `lsp` is off by default; include it here so the LSP stub's tests
+            # actually run somewhere.
+            features: default,test-dbs-external,lsp
           # Only run wasm on ubuntu, given it's the same rust target. (There is
           # a possibility of having a failure on just one platform, but it's
           # quite unlikely. If we do observe this, we can add those tests them

prqlc/prqlc/src/cli/lsp.rs

@@ -1,11 +1,9 @@
 use std::error::Error;
 
 use lsp_types::OneOf;
-use lsp_types::{
-    request::GotoDefinition, GotoDefinitionResponse, InitializeParams, ServerCapabilities,
-};
+use lsp_types::{request::GotoDefinition, GotoDefinitionResponse, ServerCapabilities};
 
-use lsp_server::{Connection, ExtractError, Message, Request, RequestId, Response};
+use lsp_server::{Connection, ErrorCode, ExtractError, Message, Request, RequestId, Response};
 
 pub fn run() -> Result<(), Box<dyn Error + Sync + Send>> {
     // Note that  we must have our logging only write out to stderr.
@@ -20,29 +18,26 @@ pub fn run() -> Result<(), Box<dyn Error + Sync + Send>> {
         definition_provider: Some(OneOf::Left(true)),
         ..Default::default()
     })
-    .unwrap();
-    let initialization_params = match connection.initialize(server_capabilities) {
-        Ok(it) => it,
-        Err(e) => {
-            if e.channel_is_disconnected() {
-                io_threads.join()?;
-            }
-            return Err(e.into());
+    .expect("`ServerCapabilities` is always serializable");
+    // The client's initialization params aren't used yet, so we deliberately
+    // don't deserialize them: a client that sends a payload our `lsp-types`
+    // version doesn't model would otherwise take the server down before the
+    // main loop even starts.
+    if let Err(e) = connection.initialize(server_capabilities) {
+        if e.channel_is_disconnected() {
+            io_threads.join()?;
         }
-    };
-    main_loop(connection, initialization_params)?;
+        return Err(e.into());
+    }
+    main_loop(connection)?;
     io_threads.join()?;
 
     // Shut down gracefully.
     eprintln!("shutting down server");
     Ok(())
 }
 
-fn main_loop(
-    connection: Connection,
-    params: serde_json::Value,
-) -> Result<(), Box<dyn Error + Sync + Send>> {
-    let _params: InitializeParams = serde_json::from_value(params).unwrap();
+fn main_loop(connection: Connection) -> Result<(), Box<dyn Error + Sync + Send>> {
     eprintln!("starting main loop");
     for msg in &connection.receiver {
         eprintln!("got msg: {msg:?}");
@@ -52,6 +47,7 @@ fn main_loop(
                     return Ok(());
                 }
                 eprintln!("got request: {req:?}");
+                let id = req.id.clone();
                 match cast::<GotoDefinition>(req) {
                     Ok((id, params)) => {
                         eprintln!("got gotoDefinition request #{id}: {params:?}");
@@ -60,7 +56,19 @@ fn main_loop(
                         connection.sender.send(Message::Response(resp))?;
                         continue;
                     }
-                    Err(err @ ExtractError::JsonError { .. }) => panic!("{err:?}"),
+                    // Params that don't deserialize are a client-side problem —
+                    // often just a version skew against our `lsp-types`. Reply
+                    // with an error for that one request rather than taking the
+                    // server down.
+                    Err(ExtractError::JsonError { method, error }) => {
+                        let resp = Response::new_err(
+                            id,
+                            ErrorCode::InvalidParams as i32,
+                            format!("invalid params for `{method}`: {error}"),
+                        );
+                        connection.sender.send(Message::Response(resp))?;
+                        continue;
+                    }
                     Err(ExtractError::MethodMismatch(req)) => req,
                 };
                 // ...
@@ -90,3 +98,62 @@ where
 {
     req.extract(R::METHOD)
 }
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use super::main_loop;
+    use lsp_server::{Connection, ErrorCode, Message, Notification, Request, RequestId};
+
+    /// Long enough that a loaded CI runner won't trip it, short enough that a
+    /// regression which hangs the loop fails the test rather than running out
+    /// the job's clock — nextest has no `terminate-after` configured, so an
+    /// unbounded wait here is never killed.
+    const REPLY_TIMEOUT: Duration = Duration::from_secs(10);
+
+    /// A `textDocument/definition` request whose params fail to deserialize
+    /// used to `panic!`, taking the whole server down. It should get an error
+    /// response and the loop should keep running.
+    #[test]
+    fn invalid_params_get_an_error_response() {
+        let (server, client) = Connection::memory();
+        // The result comes back over a channel rather than a `JoinHandle` so
+        // the wait for the loop to finish can be bounded too.
+        let (done_tx, done_rx) = std::sync::mpsc::channel();
+        std::thread::spawn(move || done_tx.send(main_loop(server)));
+
+        client
+            .sender
+            .send(Message::Request(Request {
+                id: RequestId::from(1),
+                method: "textDocument/definition".to_string(),
+                params: serde_json::json!({ "not": "a position" }),
+            }))
+            .unwrap();
+
+        let reply = client
+            .receiver
+            .recv_timeout(REPLY_TIMEOUT)
+            .expect("server should reply to a bad request rather than hang");
+        let Message::Response(resp) = reply else {
+            panic!("expected a response");
+        };
+        assert_eq!(resp.id, RequestId::from(1));
+        let err = resp.response_result.unwrap_err();
+        assert_eq!(err.code, ErrorCode::InvalidParams as i32);
+
+        // The loop survived the bad request and still handles later messages.
+        client
+            .sender
+            .send(Message::Notification(Notification {
+                method: "exit".to_string(),
+                params: serde_json::Value::Null,
+            }))
+            .unwrap();
+        done_rx
+            .recv_timeout(REPLY_TIMEOUT)
+            .expect("server should exit on `exit` rather than hang")
+            .unwrap();
+    }
+}

prqlc/prqlc/src/cli/test.rs

@@ -773,3 +773,52 @@ fn lsp_ignores_non_exit_notification() {
     shutting down server
     "###);
 }
+
+/// `initialize` params that our pinned `lsp-types` can't model — here, ones
+/// missing the required `capabilities` field — must not take the server down.
+/// Before the fix, they were deserialized into `InitializeParams` and
+/// `unwrap`ed, so the server panicked before the main loop even started.
+#[cfg(feature = "lsp")]
+#[test]
+fn lsp_unmodelled_initialize_params() {
+    let init = serde_json::to_string(&lsp_server::Message::Request(lsp_server::Request {
+        method: "initialize".into(),
+        id: lsp_server::RequestId::from(1),
+        params: serde_json::json!({}),
+    }))
+    .unwrap();
+    let initialized = serde_json::to_string(&lsp_server::Message::Notification(
+        lsp_server::Notification {
+            method: "initialized".into(),
+            params: serde_json::json!({}),
+        },
+    ))
+    .unwrap();
+    let exit = serde_json::to_string(&lsp_server::Message::Notification(
+        lsp_server::Notification {
+            method: "exit".into(),
+            params: serde_json::Value::Null,
+        },
+    ))
+    .unwrap();
+
+    assert_cmd_snapshot!(prqlc_command().args(["lsp"])
+        .pass_stdin(format!("Content-Length: {}\r\n\r\n{}Content-Length: {}\r\n\r\n{}Content-Length: {}\r\n\r\n{}",
+            init.len(), init,
+            initialized.len(), initialized,
+            exit.len(), exit))
+        , @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    Content-Length: 78
+
+    {"jsonrpc":"2.0","id":1,"result":{"capabilities":{"definitionProvider":true}}}
+    ----- stderr -----
+    starting PRQL LSP server
+    starting main loop
+    got msg: Notification(Notification { method: "exit", params: Null })
+    got notification: Notification { method: "exit", params: Null }
+    shutting down server
+    "###);
+}