zesdex/crates/zesdex-backend/src/app/lsp/client.rs

use std::io::{BufRead, BufReader, Read, Write};
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

use serde_json::{json, Value};

const LSP_INIT_TIMEOUT_MS: u64 = 60_000;
const LSP_CALL_TIMEOUT_MS: u64 = 30_000;
const LSP_DIAGNOSTICS_TIMEOUT_MS: u64 = 10_000;

pub struct LspClient {
    stdin: std::process::ChildStdin,
    stdout: BufReader<std::process::ChildStdout>,
    next_id: u64,
    server_capabilities: Value,
}

fn file_path_to_uri(path: &str) -> String {
    let abs_path = std::path::Path::new(path);
    let abs_path = if abs_path.is_relative() {
        match std::env::current_dir() {
            Ok(cwd) => cwd.join(path),
            Err(_) => abs_path.to_path_buf(),
        }
    } else {
        abs_path.to_path_buf()
    };
    let canonical = abs_path.canonicalize().unwrap_or(abs_path);
    let path_str = canonical.to_string_lossy();
    if cfg!(windows) {
        let path_str = path_str.replace('\\', "/");
        if path_str.starts_with('/') {
            format!("file://{path_str}")
        } else {
            format!("file:///{path_str}")
        }
    } else {
        format!("file://{path_str}")
    }
}

impl LspClient {
    pub fn spawn(command: &str, args: &[String]) -> anyhow::Result<Self> {
        let mut cmd = Command::new(command);
        cmd.args(args);
        cmd.stdin(Stdio::piped());
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());

        let mut child = cmd
            .spawn()
            .map_err(|e| anyhow::anyhow!("failed to spawn LSP server '{command}': {e}"))?;

        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| anyhow::anyhow!("failed to capture stdin for LSP server"))?;
        let stdout = BufReader::new(
            child
                .stdout
                .take()
                .ok_or_else(|| anyhow::anyhow!("failed to capture stdout for LSP server"))?,
        );

        let mut client = LspClient {
            stdin,
            stdout,
            next_id: 0,
            server_capabilities: Value::Null,
        };

        let init_params = json!({
            "processId": std::process::id(),
            "clientInfo": {
                "name": "zesdex",
                "version": "0.1.0"
            },
            "capabilities": {
                "textDocument": {
                    "synchronization": {
                        "dynamicRegistration": true,
                        "willSave": false,
                        "willSaveWaitUntil": false,
                        "didSave": false
                    },
                    "hover": {
                        "dynamicRegistration": true,
                        "contentFormat": ["plaintext", "markdown"]
                    },
                    "completion": {
                        "dynamicRegistration": true,
                        "completionItem": {
                            "snippetSupport": false
                        }
                    },
                    "definition": {
                        "dynamicRegistration": true
                    },
                    "references": {
                        "dynamicRegistration": true
                    },
                    "documentSymbol": {
                        "dynamicRegistration": true,
                        "hierarchicalDocumentSymbolSupport": true
                    }
                },
                "workspace": {
                    "workspaceFolders": true
                },
                "general": {
                    "positionEncodings": ["utf-16"]
                }
            }
        });

        let result = client.call_with_timeout(
            "initialize",
            &init_params,
            Duration::from_millis(LSP_INIT_TIMEOUT_MS),
        )?;
        client.server_capabilities = result.get("capabilities").cloned().unwrap_or_default();

        client.notify("initialized", &json!({}))?;

        Ok(client)
    }

    pub fn server_capabilities(&self) -> &Value {
        &self.server_capabilities
    }

    pub fn call(&mut self, method: &str, params: &Value) -> anyhow::Result<Value> {
        self.call_with_timeout(method, params, Duration::from_millis(LSP_CALL_TIMEOUT_MS))
    }

    fn call_with_timeout(
        &mut self,
        method: &str,
        params: &Value,
        timeout: Duration,
    ) -> anyhow::Result<Value> {
        self.next_id += 1;
        let id = self.next_id;
        let req = json!({
            "jsonrpc": "2.0",
            "id": id,
            "method": method,
            "params": params
        });
        self.send_frame(&req)?;
        self.read_response(id, timeout)
    }

    pub fn notify(&mut self, method: &str, params: &Value) -> anyhow::Result<()> {
        let req = json!({
            "jsonrpc": "2.0",
            "method": method,
            "params": params
        });
        self.send_frame(&req)
    }

    fn send_frame(&mut self, msg: &Value) -> anyhow::Result<()> {
        let body = serde_json::to_string(msg)
            .map_err(|e| anyhow::anyhow!("failed to serialize LSP message: {e}"))?;
        let header = format!("Content-Length: {}\r\n\r\n", body.len());
        self.stdin
            .write_all(header.as_bytes())
            .map_err(|e| anyhow::anyhow!("failed to write LSP frame header: {e}"))?;
        self.stdin
            .write_all(body.as_bytes())
            .map_err(|e| anyhow::anyhow!("failed to write LSP frame body: {e}"))?;
        self.stdin
            .flush()
            .map_err(|e| anyhow::anyhow!("failed to flush LSP stdin: {e}"))?;
        Ok(())
    }

    fn read_response(&mut self, expected_id: u64, timeout: Duration) -> anyhow::Result<Value> {
        let deadline = Instant::now() + timeout;
        loop {
            if Instant::now() > deadline {
                anyhow::bail!("LSP call timed out after {}ms", timeout.as_millis());
            }
            let frame = self.read_frame()?;
            if frame.get("id") == Some(&json!(expected_id)) {
                if let Some(err) = frame.get("error") {
                    let code = err
                        .get("code")
                        .and_then(serde_json::Value::as_i64)
                        .unwrap_or(0);
                    let msg = err
                        .get("message")
                        .and_then(|m| m.as_str())
                        .unwrap_or("unknown error");
                    anyhow::bail!("LSP error {code}: {msg}");
                }
                return Ok(frame.get("result").cloned().unwrap_or(Value::Null));
            }
        }
    }

    pub fn read_notification(&mut self, method: &str, timeout: Duration) -> anyhow::Result<Value> {
        let deadline = Instant::now() + timeout;
        loop {
            if Instant::now() > deadline {
                anyhow::bail!("timed out waiting for LSP notification '{method}'");
            }
            let frame = self.read_frame()?;
            if frame.get("method") == Some(&json!(method)) {
                return Ok(frame.get("params").cloned().unwrap_or(Value::Null));
            }
        }
    }

    fn read_frame(&mut self) -> anyhow::Result<Value> {
        let mut content_length: Option<usize> = None;
        loop {
            let mut line = String::new();
            match self.stdout.read_line(&mut line) {
                Ok(0) => anyhow::bail!("LSP server closed the connection"),
                Ok(_) => {}
                Err(e) => anyhow::bail!("LSP read error: {e}"),
            }
            let trimmed = line.trim();
            if trimmed.is_empty() {
                break;
            }
            if let Some(len_str) = trimmed.strip_prefix("Content-Length: ") {
                // Cap Content-Length at 64 MiB to prevent OOM from a
                // malicious or misconfigured LSP server (CWE-400).
                const MAX_CONTENT_LENGTH: usize = 64 * 1024 * 1024;
                let length: usize = len_str.trim().parse::<usize>().map_err(|e| {
                    anyhow::anyhow!("invalid Content-Length '{}': {}", len_str.trim(), e)
                })?;
                if length > MAX_CONTENT_LENGTH {
                    anyhow::bail!(
                        "Content-Length {length} exceeds maximum allowed size of {MAX_CONTENT_LENGTH} bytes",
                    );
                }
                content_length = Some(length);
            }
        }

        let length = content_length
            .ok_or_else(|| anyhow::anyhow!("missing Content-Length header in LSP response"))?;

        let mut body = vec![0u8; length];
        self.stdout
            .read_exact(&mut body)
            .map_err(|e| anyhow::anyhow!("failed to read LSP body ({length} bytes): {e}"))?;

        let json_str = String::from_utf8(body)
            .map_err(|e| anyhow::anyhow!("invalid UTF-8 in LSP response: {e}"))?;

        serde_json::from_str(&json_str)
            .map_err(|e| anyhow::anyhow!("invalid JSON in LSP response: {e}"))
    }

    pub fn did_open(
        &mut self,
        uri: &str,
        language_id: &str,
        version: i32,
        text: &str,
    ) -> anyhow::Result<()> {
        self.notify(
            "textDocument/didOpen",
            &json!({
                "textDocument": {
                    "uri": uri,
                    "languageId": language_id,
                    "version": version,
                    "text": text
                }
            }),
        )
    }

    pub fn did_change(&mut self, uri: &str, version: i32, text: &str) -> anyhow::Result<()> {
        self.notify(
            "textDocument/didChange",
            &json!({
                "textDocument": {
                    "uri": uri,
                    "version": version
                },
                "contentChanges": [{
                    "text": text
                }]
            }),
        )
    }

    pub fn did_close(&mut self, uri: &str) -> anyhow::Result<()> {
        self.notify(
            "textDocument/didClose",
            &json!({
                "textDocument": {
                    "uri": uri
                }
            }),
        )
    }

    pub fn hover(&mut self, uri: &str, line: u32, character: u32) -> anyhow::Result<Value> {
        self.call(
            "textDocument/hover",
            &json!({
                "textDocument": { "uri": uri },
                "position": { "line": line, "character": character }
            }),
        )
    }

    pub fn completion(&mut self, uri: &str, line: u32, character: u32) -> anyhow::Result<Value> {
        self.call(
            "textDocument/completion",
            &json!({
                "textDocument": { "uri": uri },
                "position": { "line": line, "character": character }
            }),
        )
    }

    pub fn goto_definition(
        &mut self,
        uri: &str,
        line: u32,
        character: u32,
    ) -> anyhow::Result<Value> {
        self.call(
            "textDocument/definition",
            &json!({
                "textDocument": { "uri": uri },
                "position": { "line": line, "character": character }
            }),
        )
    }

    pub fn references(&mut self, uri: &str, line: u32, character: u32) -> anyhow::Result<Value> {
        self.call(
            "textDocument/references",
            &json!({
                "textDocument": { "uri": uri },
                "position": { "line": line, "character": character },
                "context": {
                    "includeDeclaration": true
                }
            }),
        )
    }

    pub fn collect_diagnostics(
        &mut self,
        uri: &str,
        language_id: &str,
        text: &str,
    ) -> anyhow::Result<Value> {
        self.did_open(uri, language_id, 1, text)?;
        let result = self.read_notification(
            "textDocument/publishDiagnostics",
            Duration::from_millis(LSP_DIAGNOSTICS_TIMEOUT_MS),
        );
        self.did_close(uri)?;
        match result {
            Ok(params) => Ok(params
                .get("diagnostics")
                .cloned()
                .unwrap_or_else(|| json!([]))),
            Err(e) => Err(e),
        }
    }

    pub fn shutdown(&mut self) {
        let _ = self.call_with_timeout("shutdown", &json!({}), Duration::from_secs(5));
        let _ = self.notify("exit", &json!({}));
    }
}

impl Drop for LspClient {
    fn drop(&mut self) {
        let _ = self.notify("exit", &json!({}));
    }
}

pub fn path_to_lsp_uri(path: &str) -> String {
    file_path_to_uri(path)
}